import math
import warnings
import re
from typing import Tuple, Dict, List, Any, Union, Optional
import numpy as np
import pandas as pd
import numexpr as ne
from scipy.optimize import minimize, basinhopping
from . import loss as loss_module

ne.set_num_threads(1)

def _prepare_data_for_ne(X: pd.DataFrame, variables: List[str]) -> Dict[str, np.ndarray]:
    """Creates a dictionary of variables for numexpr, ensuring standardized names (x1, x2, ...)."""
    # `variables` is the list of expected column names in X, e.g., ['log(A)', 'B']
    # The skeleton is already in terms of x1, x2...
    
    # Create a mapping from the processed column names to x1, x2...
    rename_map = dict(zip(variables, [f"x{i+1}" for i in range(len(variables))]))
    
    # Extract the required columns and rename them for numexpr
    renamed_X_data = {}
    for original_name, standardized_name in rename_map.items():
        if original_name in X.columns:
            renamed_X_data[standardized_name] = X[original_name].values
        else:
            raise KeyError(f"Variable '{original_name}' not found in the provided DataFrame columns.")
            
    return renamed_X_data

def _replace_params_in_skeleton(skeleton: str, params_values: Dict[str, float]) -> str:
    """Replaces parameter placeholders (c1, p1) with their numeric values."""
    for name, value in params_values.items():
        skeleton = re.sub(r'\b' + name + r'\b', str(value), skeleton)
    return skeleton

def make_predictions(skeleton: str, X: pd.DataFrame, groups: np.ndarray, fitted_params: Dict, variables: List[str]) -> np.ndarray:
    """Makes predictions for given data using a fitted equation skeleton and parameters."""
    all_predictions = np.zeros(len(X))
    global_params = {k: v for k, v in fitted_params.items() if not isinstance(v, dict)}

    for group in np.unique(groups):
        mask = (groups == group)
        if not np.any(mask): continue
        
        X_group = X.loc[mask]
        group_params = {k: v.get(group, 1.0) for k, v in fitted_params.items() if isinstance(v, dict)}
        current_params = {**global_params, **group_params}
        equation_to_eval = _replace_params_in_skeleton(skeleton, current_params)
        
        local_dict = _prepare_data_for_ne(X_group, variables)
        try:
            with warnings.catch_warnings():
                warnings.simplefilter('ignore')
                all_predictions[mask] = ne.evaluate(equation_to_eval, local_dict=local_dict)
        except Exception as e:
            print(f"Error evaluating for group {group}: {equation_to_eval}\nError: {e}")
            all_predictions[mask] = np.nan
            
    return all_predictions

def fit_equation_parameters(skeleton: str, X: pd.DataFrame, y: np.ndarray, groups: np.ndarray,
                            c_params: List[str], u_params: List[str], loss_name: str, 
                            variables: List[str], initial_params: Dict = None, force_consistent_sign = False,
                            basinhop=False, minimizer_kwargs: Optional[Dict] = None) -> Dict:
    """Fits global (c) and group-level (u) parameters using a two-step optimization."""
    minimizer_kwargs = minimizer_kwargs or {}
    loss_fn = loss_module.get_loss(loss_name)
    unique_groups = np.unique(groups)
    initial_params = initial_params or {}
    
    # Step 1: Fit P parameters per group
    initial_cs = {c: initial_params.get(c, 1.0) for c in c_params}
    fitted_ps = {u: {} for u in u_params}
    if u_params:
        for group in unique_groups:
            mask = (groups == group)
            X_group, y_group = (X.loc[mask], y[mask])
            local_dict = _prepare_data_for_ne(X_group, variables)
            
            def objective_p(p_vals):
                p_dict = dict(zip(u_params, p_vals))
                eq_eval = _replace_params_in_skeleton(skeleton, {**initial_cs, **p_dict})
                try:
                    with warnings.catch_warnings():
                        warnings.simplefilter('ignore')
                        preds = ne.evaluate(eq_eval, local_dict=local_dict)
                        if preds.shape == ():
                            preds = np.repeat(preds.item(), len(y_group))
                        y_arr = np.asarray(y_group)
                        return np.inf if not np.all(np.isfinite(preds)) else loss_fn(y_arr, preds)
                except Exception: return np.inf

            
            u_start = [initial_params.get(u, {}).get(group, 1.0) for u in u_params]
            kwargs = {"method": "Powell"}
            if force_consistent_sign:
                kwargs.update({"bounds": [(1e-10, None)] * len(u_start)})
            kwargs.update(minimizer_kwargs)
            if basinhop:
                def take_step(x):
                    random_array = np.random.randn(*x.shape)
                    return random_array
                kwargs["options"] = {"maxiter": 10, "ftol": 1e-3}
                res_p = basinhopping(objective_p, u_start, niter=4, T=0, take_step=take_step, niter_success=2, minimizer_kwargs=kwargs)
                u_start = res_p.x
            kwargs["options"] = {"maxiter": 10*len(u_start), "ftol": 1e-6}
            res_p = minimize(objective_p, u_start, **kwargs)
            for i, u in enumerate(u_params): fitted_ps[u][group] = res_p.x[i]
    
    if not c_params:
        return fitted_ps

    # Step 2: Fit all parameters together
    # Step 2: Jointly fit all parameters together.
    # The warm start (x0) is constructed from `initial_params` for C's and
    # the results of Step 1 for P's.
    x0 = [initial_params.get(c, 1.0) for c in c_params]
    bounds = [(None, None)] * len(x0)
    group_map = []
    for u in u_params:
        for group in unique_groups:
            x0.append(fitted_ps.get(u, {}).get(group, 1.0))
            group_map.append({'u': u, 'g': group})
            bounds.append((1e-10, None))
        
    kwargs = {
        "method": "Powell"
    }
    if force_consistent_sign:
        kwargs.update({"bounds": bounds})
    kwargs.update(minimizer_kwargs)

    def objective_all(all_vals):
        total_loss, n_c = 0.0, len(c_params)
        c_dict = dict(zip(c_params, all_vals[:n_c]))
        u_struct = {u: {} for u in u_params}
        for i, m in enumerate(group_map): u_struct[m['u']][m['g']] = all_vals[n_c + i]
            
        for group in unique_groups:
            mask = (groups == group)
            X_group, y_group = (X.loc[mask], y[mask])
            group_p_vals = {u: u_struct.get(u, {}).get(group, 1.0) for u in u_params}
            eq_eval = _replace_params_in_skeleton(skeleton, {**c_dict, **group_p_vals})
            local_dict = _prepare_data_for_ne(X_group, variables)
            try:
                with warnings.catch_warnings():
                    warnings.simplefilter('ignore')
                    preds = ne.evaluate(eq_eval, local_dict=local_dict)
                    if preds.shape == ():
                        preds = np.repeat(preds.item(), len(y_group))
                    y_arr = np.asarray(y_group)
                    total_loss += (np.inf if not np.all(np.isfinite(preds)) else loss_fn(y_arr, preds)) * len(y_group)
            except Exception: total_loss += np.inf * len(y_group)
        return total_loss / len(y) if len(y) > 0 else 0.0

    kwargs["options"] = {"maxiter": 10*len(c_params), "ftol": 1e-6}
    res_all = minimize(objective_all, np.array(x0), **kwargs)
    
    # Unpack final results
    final_params, n_c = {}, len(c_params)
    for i, c in enumerate(c_params): final_params[c] = res_all.x[i]
    for u in u_params: final_params[u] = {}
    for i, m in enumerate(group_map): final_params[m['u']][m['g']] = res_all.x[n_c + i]
    return final_params