from __future__ import annotations
import re 
from typing import Tuple, List, Union, Dict
import numpy as np
import pandas as pd
import sympy as sp
from . import calculator
import warnings
from typing import TYPE_CHECKING
from .utils import build_cv_folds
if TYPE_CHECKING:
    from .data_processing import DataTransformer

def skeleton_to_canonical(expr: str):
    """Recursively convert a sympy expression to a canonical tuple."""
    expr = re.sub(r'\bc\d+\b', 'CCC', expr)
    expr = re.sub(r'\bu\d+\b', 'UUU', expr)
    expr = sp.sympify(expr)
    def expr_to_canonical(expr: sp.Expr):
        if expr.is_Atom:
            return (str(expr),)
        elif isinstance(expr, sp.Add):
            args = frozenset(expr_to_canonical(arg) for arg in expr.args)
            return ('add', args)
        elif isinstance(expr, sp.Mul):
            args = sorted(expr_to_canonical(arg) for arg in expr.args)
            return ('mul', tuple(args))
        elif isinstance(expr, sp.Pow):
            base_can = expr_to_canonical(expr.base)
            exp_can = expr_to_canonical(expr.exp)
            return ('Pow', (base_can, exp_can))
        else:
            # Other functions like exp(...)
            func_name = expr.func.__name__
            args = tuple(expr_to_canonical(arg) for arg in expr.args)
            return (func_name, args)
    expr = expr_to_canonical(expr)
    return expr

class Equation:
    def __init__(self, eq_str, transformer: DataTransformer, loss = 'MSE', force_consistent_sign = False):
        self.id = None
        self.original_str = eq_str
        self.loss = loss
        self.force_consistent_sign = force_consistent_sign
        self.transformer = transformer
        # Variables are the column names of the processed data (e.g., ['tmp_1', 'tmp_2'])
        self.variables = transformer.placeholder_names

        # Standardize variable names to x1, x2, ...
        replace_map = dict(zip(self.variables, [f"x{i+1}" for i in range(len(self.variables))]))
        eq_str_std_vars = replace_feature_names(eq_str, replace_map)

        # Extract numeric constants as 'c' parameters and create the equation skeleton
        self.skeleton, self.original_cs_values = replace_and_extract_numeric_parameters(eq_str_std_vars)
        self.canonical_expression = skeleton_to_canonical(self.skeleton)

        # Identify c and u parameters from the skeleton using regex
        self.c_params = sorted(list(set(re.findall(r'\bc\d+\b', self.skeleton))))
        self.u_params = sorted(list(set(re.findall(r'\bu\d+\b', self.skeleton))))

        user_skeleton = transformer.to_user_skeleton(self)
        self.complexity = complexity_calculation(user_skeleton)

        readable_skeleton = round_constants_in_equation(move_params_to_front_str(user_skeleton, ['c', 'u']))
        readable_skeleton = readable_skeleton.replace('**', "^").replace('*', "⋅")
        # Replace ci's with βi's
        readable_skeleton = re.sub(r'\bc(\d+)\b', r'β\1', readable_skeleton)
        self.readable_skeleton = readable_skeleton

        self.fitted_params = None
        self.cv_params_per_fold = None
        self.loss_value = None

    def __str__(self):
        return self.skeleton
    
    def __eq__(self, other):
        if not isinstance(other, Equation):
            return NotImplemented
        return self.canonical_expression == other.canonical_expression
    
    def __hash__(self):
        return hash(self.canonical_expression)

    def fit(self, X: pd.DataFrame, y: pd.Series, groups: pd.Series=None, initial_params = None, basinhopping=False, **minimizer_kwargs):
        """
        Fits the global (c) and group-level (u) parameters of the equation, optionally using a warm start.

        This method uses a two-step optimization process:
        1. It first fits the group-level parameters (u) for each group independently, 
           while keeping the global parameters (c) fixed at their initial values found by the SR engine.
        2. It then uses the results as a warm start to perform a final optimization of all 
           parameters (c and u) together.

        Parameters
        ----------
        X : pd.DataFrame
            Feature data.
        y : pd.Series
            Target data.
        groups : pd.Series, optional
            Group labels for each data point. If None, performs a standard regression
            without group-level parameters, treating all data as a single group.
        **minimizer_kwargs : dict, optional
            Additional keyword arguments to pass to `scipy.optimize.minimize`.
        """
        if groups is None: groups = pd.Series('default_group', index=y.index)

        # Warning checks
        unique_groups = np.unique(groups)
        if len(unique_groups) > 1 and not self.u_params:
            warnings.warn("Multi-group data provided, but the equation has no group-level parameters (u_i). Treating as a single group.")
            groups = np.array(['default_group'] * len(y))
        elif len(unique_groups) == 1 and self.u_params:
            warnings.warn("Equation has group-level parameters (u_i), but only one group was provided. global parameter (c_i) will be fit for this single group.") 
        
        params_for_cal = initial_params
        if params_for_cal is None:
            params_for_cal = dict(zip(self.c_params, self.original_cs_values))

        self.fitted_params = calculator.fit_equation_parameters(
            skeleton=self.skeleton,
            X=X,
            y=y.values,
            groups=groups.values,
            c_params=self.c_params,
            u_params=self.u_params,
            loss_name=self.loss,
            variables=self.variables,
            initial_params=params_for_cal,
            force_consistent_sign = self.force_consistent_sign,
            basinhop=basinhopping,
            minimizer_kwargs=minimizer_kwargs
        )
        return self.fitted_params

    def predict(self, X: pd.DataFrame, groups: pd.Series = None, params_override=None, threshold=None) -> np.ndarray:
        """
        Makes predictions using the fitted equation, optionally with a specific parameter set.

        Parameters
        ----------
        X : pd.DataFrame
            Feature data for which to make predictions.
        groups : pd.Series or np.ndarray, optional
            Group labels for each data point. This is required if the model was 
            fitted with multiple groups.
        """
        params_to_use = params_override or self.fitted_params
        if params_to_use is None:
            raise ValueError("Equation not fitted. Call fit() before predict().")

        # Check if group information is needed for prediction
        has_group_params = any(isinstance(val, dict) for val in self.fitted_params.values())
        if has_group_params and groups is None:
            raise ValueError("This model was fitted with group-level parameters. You must provide a 'groups' array for prediction.")
        
        if groups is None:
            groups = pd.Series('default_group', index=X.index)
        
        predictions = calculator.make_predictions(
            skeleton=self.skeleton,
            X=X,
            groups=groups.values,
            fitted_params=params_to_use,
            variables=self.variables
        )
        return predictions
    
    def store_cv_fold_params(self, params: dict):
        """Stores the fitted parameters from a single CV fold."""
        if self.cv_params_per_fold is None:
            self.cv_params_per_fold = []
        self.cv_params_per_fold.append(params.copy())
    
    def shallow_copy(self):
        new_eq = Equation(
            eq_str=self.original_str,
            transformer=self.transformer, 
            loss=self.loss,
            force_consistent_sign=self.force_consistent_sign
        )
        new_eq.fitted_params = self.fitted_params
        return new_eq

class ClassificationEquation(Equation):
    def predict(self, X: pd.DataFrame, groups: pd.Series = None, params_override=None, threshold=None) -> np.ndarray:
        raw = super().predict(X)
        probas = sigmoid(raw)

        if threshold is not None:
            if threshold == False:
                return raw
            if not isinstance(threshold, (int, float)) or not (0 < threshold < 1):
                raise ValueError("threshold must be a numeric value strictly between 0 and 1")
            return (probas >= threshold).astype(int)

        return probas
    
    def fit_threshold(self, X, y, groups=None, k_folds=5):
        """
        Perform k-fold cross-validation to determine the optimal decision threshold 
        for binary predictions using the True Skill Statistic (TSS).

        The method splits the data into `k` folds, fits a copy of the current model 
        on the training portion of each fold, and computes the threshold that 
        maximizes the TSS on the validation portion. The final threshold is 
        obtained by averaging the best thresholds across all folds.

        Parameters
        ----------
        X : pandas.DataFrame or numpy.ndarray
            Feature matrix for training and cross-validation.
        y : pandas.Series or numpy.ndarray
            Binary target labels.
        groups : pandas.Series or numpy.ndarray, optional
            Group labels for multi-view or grouped fitting (default is None).
        k : int, optional
            Number of folds for cross-validation (default is 5).

        Returns
        -------
        float
            Best threshold value that maximizes the TSS metric across folds.

        Notes
        -----
        - TSS (True Skill Statistic) is computed as `sensitivity + specificity - 1`.
        - Within each fold, the threshold closest to 0.5 is preferred if multiple thresholds 
        achieve the same TSS.
        - The method uses `self.shallow_copy()` to fit a temporary copy of the model on each fold 
        to avoid modifying the original fitted parameters.
        """
        def get_threshold(eq: Equation, X, y, groups=None):
            """Find the threshold that maximizes the TSS metric"""
            y_probas = np.asarray(eq.predict(X, groups=groups)).ravel()
            y = np.asarray(y).ravel()
            P = np.sum(y)
            N = len(y) - P

            # Sort probas and y in descending order
            sort_idx = np.argsort(y_probas)[::-1]
            sorted_probas = y_probas[sort_idx]
            sorted_y = y[sort_idx]
            # Number of positive predictions at each threshold
            cum_count_positive_predictions = np.arange(1, len(sorted_probas) + 1)
            # Compute number of true positives and false positives at each threshold
            cum_tp = np.cumsum(sorted_y)
            cum_fp = cum_count_positive_predictions - cum_tp
            # Compute sensitivities and specificities at each threshold
            all_sensitivities = cum_tp / P
            all_specificities = 1 - cum_fp / N
            # Compute tsss at each threshold
            all_tss = all_sensitivities + all_specificities - 1 
            # Find the threshold that maximizes TSS and is closest to 0.5
            best_threshold_idx= np.argmax(all_tss - 1e-2 * np.abs(sorted_probas - 0.5))
            best_threshold = float(sorted_probas[best_threshold_idx])

            return best_threshold
        folds = build_cv_folds(X, y, groups, k_folds=k_folds, shuffle=True, random_state=0)
        thresholds = []
        eq = self.shallow_copy()

        for train_idx, val_idx in folds:
            X_train, X_val = X.iloc[train_idx], X.iloc[val_idx]
            y_train, y_val = y.iloc[train_idx], y.iloc[val_idx]
            groups_train, groups_val = groups.iloc[train_idx], groups.iloc[val_idx]

            eq.fit(X_train, y_train, groups_train, initial_params=self.fitted_params)

            # get_threshold function finds the best threshold for this fold
            threshold = get_threshold(eq, X_val, y_val, groups_val)
            thresholds.append(threshold)

        # Average thresholds across all folds
        best_threshold = np.mean(thresholds)
        return best_threshold
    
    def shallow_copy(self):
        new_eq = ClassificationEquation(
            eq_str=self.original_str,
            transformer=self.transformer, 
            loss=self.loss,
            force_consistent_sign=self.force_consistent_sign
        )
        new_eq.fitted_params = self.fitted_params
        return new_eq

def replace_and_extract_numeric_parameters(expr: str) -> Tuple[str, List[Union[int, float]]]:
    """
    Replaces numeric constants in an equation string with placeholders c1, c2, ...
    preserving their sign in the formula, and returns the new 'skeleton' equation and the list of extracted constants.
    
    Returns:
        (new_expr, constants_list)
    """
    # Regex to match signed numbers (int, float, sci notation) not in variable names
    num_pattern = re.compile(
        r"""
        (?<![A-Za-z0-9_])   # not part of a variable name before
        ([+-]?)             # optional sign (group 1)
        (                   # numeric literal (group 2)
            (?:\d+\.\d*|\.\d+|\d+)   # int or decimal
            (?:[eE][+-]?\d+)?        # optional scientific notation
        )
        (?![A-Za-z0-9_])    # not part of a variable name after
        """,
        re.VERBOSE
    )

    constants: List[Union[int, float]] = []
    counter = 0

    def repl(match):
        nonlocal counter
        start = match.start()
        # Skip if immediately preceded by '**' (exponent)
        if start >= 2 and expr[start-2:start] == '**':
            return match.group(0)  # leave unchanged
        sign = match.group(1) or ''
        num_str = match.group(2)
        # Convert to int if possible, else float
        try:
            val = int(num_str)
        except ValueError:
            val = float(num_str)
        # Skip if it's exactly 1 (e.g. the 1 in 1/X or in c1/(1 + c2*X))
        if abs(val) == 1:
            return match.group(0)  # leave unchanged
        constants.append(abs(val))
        counter += 1
        return f"{sign}c{counter}"

    new_expr = num_pattern.sub(repl, expr)
    return new_expr, constants

def custom_round(x):
    if x == 0:
        return "0"
    elif x < 0.1:
        return "{:.2e}".format(x)
    else:
        return str(round(x, 2))


def round_constants_in_equation(equation):
    """Rounds all floating point numbers in an equation string to a fixed precision."""
    # Ensure the equation is treated as a single string
    if isinstance(equation, list):
        equation = ''.join(map(str, equation))
    
    # Use regex to find all numbers in the equation and round them
    def round_match(match):
        value = float(match.group())
        return f"{custom_round(value)}"
    
    # Match numbers including those in scientific notation
    rounded_equation = re.sub(r'-?\d+\.\d+(e[+-]?\d+)?', round_match, equation)
    return rounded_equation

_ops_weights = {
    "exp": 1, "log": 1, "sin": 1, "cos": 1, 'c': 1, 'u': 1, '+': 0,
    '-': 0, '*': 0, '/': 1, '**': 1
}
_variables_weight = 1
_variables_pattern = r"\b(?!\d+\b)\w+\b"

sorted_keys = sorted(_ops_weights.keys(), key=lambda k: -len(k))    # Sort keys by length (longer patterns like '**' before '*')
pattern_str = "(" + "|".join(re.escape(k) for k in sorted_keys) + '|' + _variables_pattern + ")"
_pattern = re.compile(pattern_str)

def complexity_calculation(s):
    complexity = 0
    matches = _pattern.findall(s)
    for m in matches:
        complexity += _ops_weights.get(m, _variables_weight)
    return complexity

def replace_feature_names(expression: str, var_map) -> str:
    """
    Perform substitutions in an expression string

    Parameters
    ----------
    expr_str : str
        The expression to transform, e.g. "x1 * x2 + log(x3) + x4".
    replace_map : Dict[str, str]
        Mapping from old symbol names to new symbol names, e.g. {'x1*x2': 'A', 'log(x3)': 'B', 'x4': C}.

    Returns
    -------
    str
        The resulting expression as a string, with all replacements applied, e.g. "A + B + C"
    """
    # First replace with placeholders and only in the end with the actual replacements to circumvent circular replacement problems (e.g. {'x1':'x2', 'x2':'x1'})
    import uuid
    placeholder_map = {
        k: f"tmp_{uuid.uuid4().hex}"
        for k in var_map
    }

    # First apply regex to deal with multiplication in the key space and start with longer substitutions first (e.g. A*B)
    keys = sorted(var_map.keys(), key=len, reverse=True)
    pattern = re.compile(r'\b(?:' + '|'.join(map(re.escape, keys)) + r')\b')
    def repl(m: re.Match) -> str:
        return placeholder_map[m.group(0)]
    expression = pattern.sub(repl, expression)

    # Apply xreplace to deal with most substitutions and with functions in the key space which regex doesn't address (e.g. log(A))
    expression = sp.sympify(expression)
    replacements = {sp.sympify(k): sp.Symbol(v) for k, v in placeholder_map.items()}
    expression = str(expression.xreplace(replacements))

    # Replace the placeholders with the target replacements
    final_map = {
        placeholder_map[k]: var_map[k]
        for k in var_map
    }
    expression = sp.sympify(expression)
    replacements = {sp.sympify(k): sp.Symbol(v) for k, v in final_map.items()}
    expression = str(expression.xreplace(replacements))

    return str(sp.sympify(expression))


def move_params_to_front_str(expr_str, coef_prefixes):
    """
    Parse expr_str with sympy and return a string where, for each multiplicative
    term, any symbols whose name starts with any of coef_prefixes (and numeric
    constants) are moved to the front.

    - expr_str: e.g. "X1*X2*c1 + X2*p1*log(X1)"
    - coef_prefixes: list of prefix strings, e.g. ['c', 'u']

    Returns: rewritten expression as a string (keeps other factor ordering).
    """
    def is_coef_symbol(f, coef_prefixes):
        # consider numeric constants as coefficient-like too
        if getattr(f, 'is_Number', False):
            return True
        if isinstance(f, sp.Symbol):
            for p in coef_prefixes:
                if f.name.startswith(p):
                    return True
        return False

    def factor_to_string(f, coef_prefixes):
        # recursively build string, handling nested Mul/Add/Pow/Function
        if isinstance(f, sp.Mul):
            return term_to_string_with_coeffs_front(f, coef_prefixes)
        if isinstance(f, sp.Add):
            return '(' + expr_to_string_with_coeffs_front(f, coef_prefixes) + ')'
        if isinstance(f, sp.Pow):
            base, exp = f.args
            return factor_to_string(base, coef_prefixes) + '**' + factor_to_string(exp, coef_prefixes)
        if getattr(f, 'is_Function', False):
            fname = f.func.__name__
            args = ','.join(expr_to_string_with_coeffs_front(a, coef_prefixes) for a in f.args)
            return f"{fname}({args})"
        return str(f)

    def term_to_string_with_coeffs_front(term, coef_prefixes):
        # term can be Mul or single atom; get factors in original order (term.args)
        if isinstance(term, sp.Mul):
            args = list(term.args)
        else:
            args = [term]
        coeffs = []
        others = []
        for a in args:
            if is_coef_symbol(a, coef_prefixes):
                coeffs.append(a)
            else:
                others.append(a)
        ordered = coeffs + others
        return '*'.join(factor_to_string(a, coef_prefixes) for a in ordered)

    def expr_to_string_with_coeffs_front(expr, coef_prefixes):
        if isinstance(expr, sp.Add):
            return ' + '.join(term_to_string_with_coeffs_front(t, coef_prefixes) for t in expr.args)
        else:
            return term_to_string_with_coeffs_front(expr, coef_prefixes)

    # parse without implicit multiplication (so X1 stays 'X1')
    expr = sp.parse_expr(expr_str, evaluate=False)
    return expr_to_string_with_coeffs_front(expr, coef_prefixes)

def safe_exp(x):
    return np.exp(np.clip(x, -30, 30))

def sigmoid(x):
    return 1.0 / (1.0 + safe_exp(-x))