import pandas as pd
import numpy as np
from scipy.special import gammaln
import warnings
from .equation_processing import Equation
from .loss import get_loss
from typing import List
import math

# --- Likelihood Functions ---

def gaussian_log_likelihood(y_true, y_pred, **kwargs):
    """Computes the Gaussian log-likelihood."""
    y_true = np.asarray(y_true)
    y_pred = np.asarray(y_pred)
    n = len(y_true)
    residuals = y_true - y_pred
    sigma_sq = np.clip(np.mean(residuals**2), 1e-4 * np.var(y_true), 10 * np.var(y_true))
    return -0.5 * n * np.log(2 * np.pi * sigma_sq) - 0.5 * np.sum(residuals**2) / sigma_sq

def bernoulli_log_likelihood(y_true, y_pred, **kwargs):
    """Computes the Bernoulli log-likelihood for binary classification."""
    eps = 1e-8
    y_true = np.asarray(y_true)
    y_pred = np.asarray(y_pred)
    y_pred = np.clip(y_pred, eps, 1 - eps) # Sigmoid is assumed to be applied
    return np.sum(y_true * np.log(y_pred) + (1 - y_true) * np.log(1 - y_pred))

def poisson_log_likelihood(y_true, y_pred, **kwargs):
    """Computes the Poisson log-likelihood."""
    eps = 1e-8
    y_true = np.asarray(y_true)
    y_pred = np.asarray(y_pred)
    if np.any(y_pred <= 0):
        warnings.warn(
        f"Some predicted values are non-positive; "
        f"clipping to {eps:.0e} to avoid invalid log-likelihood computation."
    )
        y_pred = np.maximum(y_pred, eps)
    # Poisson log-likelihood formula: y*log(mu) - mu - log(y!)
    return np.sum(y_true * np.log(y_pred) - y_pred - gammaln(y_true + 1))

# metric
def r2_score(y_true, y_pred, groups=None, **kwargs):
    """Computes the R-squared score (with optional group-wise baseline)."""
    if groups is None:
        ss_res = np.sum((y_true - y_pred)**2)
        ss_tot = np.sum((y_true - np.mean(y_true))**2)
    else:
        ss_res, ss_tot = 0.0, 0.0
        for g in np.unique(groups):
            mask = groups == g
            y_g, y_pred_g = y_true[mask], y_pred[mask]
            ss_res += np.sum((y_g - y_pred_g)**2)
            ss_tot += np.sum((y_g - np.mean(y_g))**2)
    return 1 - (ss_res / (ss_tot + 1e-15))

def rmse_score(y_true, y_pred, **kwargs):
    """Computes the Root Mean Squared Error."""
    return np.sqrt(np.mean((y_true - y_pred)**2))

def mae(y_true, y_pred, **kwargs):
    """Computes the Mean Absolute Error"""
    return np.mean(np.abs(y_true-y_pred))

def f1_score(y_true, y_pred, *, threshold=0.5, **kwargs):
    """Computes the F1 score for binary classification."""
    y_true = np.asarray(y_true)
    y_pred_class = apply_threshold(y_pred, threshold)
    tp = np.sum((y_true == 1) & (y_pred_class == 1))
    fp = np.sum((y_true == 0) & (y_pred_class == 1))
    fn = np.sum((y_true == 1) & (y_pred_class == 0))
    precision = tp / (tp + fp + 1e-8)
    recall = tp / (tp + fn + 1e-8)
    return 2 * (precision * recall) / (precision + recall + 1e-8)

def poisson_deviance(y_true, y_hat, eps=1e-8):
    y_true = np.asarray(y_true)
    y_hat = np.asarray(y_hat)
    return np.sum(y_true * np.log((y_true + eps) / y_hat) - (y_true - y_hat))

def pseudo_R2_dev(y, y_pred, groups=None, eps=1e-8, **kwargs):
    """Compute Pseudo R² deviance for some model and task."""
    
    # Log likelihood of the model with only the intercept
    if groups is None:
        intercept_pred = np.full_like(y, np.mean(y), dtype=float)
    else:
        groups = np.asarray(groups)
        if groups.shape != y.shape:
            raise ValueError("groups must have the same shape as y.")
        # Map labels to contiguous ints
        unique_labels, inv = np.unique(groups, return_inverse=True)
        # compute sums and counts per label, then means
        sums = np.bincount(inv, weights=y)
        counts = np.bincount(inv)
        # safety: counts must be >0 for labels that exist; if any zero (shouldn't happen) set mean to global mean
        means = np.where(counts > 0, sums / counts, np.mean(y))
        intercept_pred = means[inv]
            
    intercept_deviance = poisson_deviance(y, intercept_pred)
    fitted_deviance = poisson_deviance(y, y_pred)

    if intercept_deviance == 0:
        return np.nan

    return 1 - (fitted_deviance / intercept_deviance)

def true_skill_statistic(y_true, y_pred, *, threshold: float = 0.5, all=False, **kwargs):
    y_true = np.asarray(y_true)
    y_pred = apply_threshold(y_pred, threshold)
    tp = np.sum((y_true == 1) & (y_pred == 1))
    tn = np.sum((y_true == 0) & (y_pred == 0))
    fp = np.sum((y_true == 0) & (y_pred == 1))
    fn = np.sum((y_true == 1) & (y_pred == 0))
    denom_pos = tp + fn
    denom_neg = tn + fp

    if denom_pos == 0 or denom_neg == 0:
        return np.nan
    sensitivity = tp / denom_pos
    specificity = tn / denom_neg
    tss = sensitivity + specificity - 1
    if all:
        return tss, sensitivity, specificity
    return tss

# --- Configuration Mappings ---

LOSS_TO_METRIC = {
    'MSE': 'R2', 'RMSE': 'R2', 'MAE': 'R2',
    'BinaryCrossEntropy': 'TSS', 'ExpLoss': 'TSS',
    'PoissonDeviance': 'PseudoR2' 
}

LOSS_TO_LIKELIHOOD = {
    'BinaryCrossEntropy': bernoulli_log_likelihood,
    'ExpLoss': bernoulli_log_likelihood,
    'PoissonDeviance': poisson_log_likelihood
}
DEFAULT_LIKELIHOOD = gaussian_log_likelihood

METRIC_FUNCTIONS = {
    'R2': r2_score, 'RMSE': rmse_score, 'MAE': mae, 
    'F1': f1_score, 'TSS': true_skill_statistic,
    'PseudoR2': pseudo_R2_dev
}
METRIC_PROPERTIES = {
    'Loss': 'minimize',
    'R2': 'maximize', 'RMSE': 'minimize', 'MAE': 'minimize', 
    'F1': 'maximize', 'TSS': 'maximize',
    'PseudoR2': 'maximize',
    'AIC': 'minimize', 'AICc': 'minimize', 'AIC_w': 'maximize', 'AICc_w': 'minimize'
}

# --- Core Evaluation Logic ---

def compute_pareto_front(df, x_col, y_col, x_higher_is_better, y_higher_is_better):
    """Identifies the Pareto optimal points from a DataFrame preserving the original index."""
    sorted_df = df.sort_values(
        by=[x_col, y_col],
        ascending=[not x_higher_is_better, not y_higher_is_better]
    )
    pareto_indices = []
    best_y = -np.inf if y_higher_is_better else np.inf

    for idx, row in sorted_df.iterrows():
        y = row[y_col]
        if (y_higher_is_better and y > best_y) or (not y_higher_is_better and y < best_y):
            pareto_indices.append(idx)
            best_y = y

    return df.loc[pareto_indices].sort_values(by=x_col)

def compute_elbow_metric(
    pareto_df: pd.DataFrame,
    metric_col: str,
    complexity_col: str = "Complexity",
    metric_higher_is_better: bool = True,
) -> pd.DataFrame:
    """
    Compute an angle-based elbow (elbow) metric for 2D Pareto front solutions.
    Based on the paper Branke et al. (2004).

    This function identifies "elbows" in the Pareto front by measuring the
    oriented angle at each solution with respect to its neighbors. A larger angle
    indicates a stronger trade-off between objectives, making the point more
    clearly identifiable as an elbow.

    **Key ideas:**
      - The angle is computed between the vectors (xi-1 → xi) and (xi → xi+1).
      - If fewer neighbors exist, artificial vertical (left) or horizontal (right)
        vectors are used as substitutes.
      - To intensify the elbow detection, two neighbors are considered on each side
        (up to four angle candidates per point), and the largest angle is assigned.
      - Duplicate points (same complexity and metric) are treated as one.
      - Angles in the Pareto front context always lie between 90° and 270°:
        - 180° → the point offers the same trade-off as its neighbors.
        - 270° → ideal "perfect" elbow (theoretical best).
        - 90° → worst possible point in terms of trade-off.

    **Normalization:**
      - The angle is linearly mapped to a range of ]-1, 1[:
            elbow_metric = (angle - 180) / (270 - 180)
        so that:
            -1 → (theoretical) worst case (90°)
             0 → flat trade-off (180°)
             1 → (theoretical) perfect elbow (270°)

    Parameters
    ----------
    pareto_df : pd.DataFrame
        DataFrame containing Pareto front points.
    metric_col : str
        Column name for the performance metric.
    complexity_col : str, default="Complexity"
        Column name for the complexity dimension.
    metric_higher_is_better : bool, default=True
        Whether higher values of the metric are better (otherwise the metric
        is negated before analysis).

    Returns
    -------
    pd.DataFrame
        A copy of the input DataFrame with an added column:
          - 'elbow_metric' : float in [-1, 1], the normalized elbow measure.
    """
    df = pareto_df.copy()
    if complexity_col not in df.columns:
        raise KeyError(f"complexity_col '{complexity_col}' not in DataFrame")
    if metric_col not in df.columns:
        raise KeyError(f"metric_col '{metric_col}' not in DataFrame")
    x = df[complexity_col].to_numpy(dtype=float)
    y_raw = df[metric_col].to_numpy(dtype=float)
    y = y_raw if metric_higher_is_better else (-y_raw)  # metric axis is flipped internally (multiplied by -1) before any computation so that the “better” direction is consistently upward.
    n = len(df)
    if n == 0:
        df["elbow_metric"] = np.nan
        return df
    pairs = list(zip(x.tolist(), y.tolist()))
    ser_pairs = pd.Series(pairs)
    codes, uniques = pd.factorize(ser_pairs, sort=False)  # codes: length n; uniques: unique tuples in order of occurrence
    uniq_list = list(uniques)
    U = len(uniq_list)
    # arrays of unique points
    ux = np.array([p[0] for p in uniq_list], dtype=float)
    uy = np.array([p[1] for p in uniq_list], dtype=float)
    # sort unique points by complexity (ux)
    sort_idx = np.argsort(ux, kind="mergesort")  # indices into uniq_list
    ux_sorted = ux[sort_idx]
    uy_sorted = uy[sort_idx]
    # mapping from unique-index to position in sorted order
    pos_for_unique_index = np.empty(U, dtype=int)
    for sorted_pos, unique_idx in enumerate(sort_idx):
        pos_for_unique_index[unique_idx] = sorted_pos
    # map each row to sorted position
    row_to_pos = pos_for_unique_index[codes]
    def vec_from_to(a_x, a_y, b_x, b_y):
        return np.array([b_x - a_x, b_y - a_y], dtype=float)
    v_vertical = np.array([0.0, 1.0])
    v_horizontal = np.array([1.0, 0.0])
    def oriented_angle(u, v):
        """
        Return the oriented angle (in degrees) between 2D vectors u and v.
        If cross < 0 (clockwise), return the direct angle (0-180).
        Otherwise, return 360 - angle (reflex angle).
        """
        u = np.array(u, dtype=float)
        v = np.array(v, dtype=float)

        # Normalize
        nu = u / np.linalg.norm(u)
        nv = v / np.linalg.norm(v)

        # Base angle magnitude (0–180)
        cos_theta = np.clip(np.dot(nu, nv), -1.0, 1.0)
        theta = np.degrees(np.arccos(cos_theta))

        # Orientation via 2D cross product
        cross = nu[0]*nv[1] - nu[1]*nv[0]

        if cross < 0:
            return theta
        else:
            return 360.0 - theta
    angles_for_unique_sorted = np.full(U, np.nan, dtype=float)
    for i in range(U):
        # use ux_sorted and uy_sorted for positions
        left_positions = []
        if i - 1 >= 0:
            left_positions.append(i - 1)
        if i - 2 >= 0:
            left_positions.append(i - 2)
        right_positions = []
        if i + 1 < U:
            right_positions.append(i + 1)
        if i + 2 < U:
            right_positions.append(i + 2)
        if len(left_positions) == 0 and len(right_positions) == 0:
            angles_for_unique_sorted[i] = float("nan")
            continue
        left_vectors = []
        if len(left_positions) == 0:
            left_vectors.append(v_vertical)
        else:
            for p in left_positions:
                v = vec_from_to(ux_sorted[p], uy_sorted[p], ux_sorted[i], uy_sorted[i])
                left_vectors.append(v)
        right_vectors = []
        if len(right_positions) == 0:
            right_vectors.append(v_horizontal)
        else:
            for q in right_positions:
                v = vec_from_to(ux_sorted[i], uy_sorted[i], ux_sorted[q], uy_sorted[q])
                right_vectors.append(v)
        angle_candidates = []
        for vL in left_vectors:
            for vR in right_vectors:
                a = oriented_angle(-vL, vR)
                if not math.isnan(a):
                    angle_candidates.append(a)
        if len(angle_candidates) == 0:
            angles_for_unique_sorted[i] = float("nan")
        else:
            angles_for_unique_sorted[i] = float(max(angle_candidates))
    elbow_angles = np.array([angles_for_unique_sorted[pos] for pos in row_to_pos], dtype=float)
    df["elbow_metric"] = (elbow_angles - 180)/(270-180)
    return df

def compute_aic(log_likelihood, k, n=None, correction=False):
    """
    Compute AIC or AICc (corrected AIC) for a fitted model
    """
    # Basic AIC
    aic_value = 2 * k - 2 * log_likelihood

    # Apply small-sample correction if requested
    if correction:
        if n is None:
            raise ValueError("n (sample size) must be provided for AICc.")
        if n - k - 1 <= 0:
            raise ValueError("n must exceed k + 1 for the correction term.")
        aic_value += (2 * k * (k + 1)) / (n - k - 1)

    return aic_value

def compute_aic_weights(aic_values):
    """
    Convert an array of AIC values into Akaike weights.
    """
    aic = np.asarray(aic_values)
    min_aic = np.nanmin(aic)
    delta = aic - min_aic
    rel_lik = np.exp(-0.5 * delta)
    weights = rel_lik / np.nansum(rel_lik)
    return weights

def evaluate_equations(predictions_list: list, num_params_list: list, y_true: np.ndarray, groups: np.ndarray,
                    partial_results_df: pd.DataFrame, loss_name: str, 
                    metrics_to_calc: List[str], likelihood_override) -> pd.DataFrame:
    """
    Computes a requested set of evaluation metrics for a list of equations.
    """
    # Generate predictions and get parameter counts from the equation objects
    output_df = pd.DataFrame(index=partial_results_df.index)
    output_df['Complexity'] = partial_results_df['Complexity']
    n = len(y_true)

    # --- AIC Calculation ---
    aic_names = [name for name in ('AIC', 'AICc') if name in metrics_to_calc]
    for aic_name in aic_names:
        metric_func = likelihood_override or LOSS_TO_LIKELIHOOD.get(loss_name, DEFAULT_LIKELIHOOD)
        log_likelihoods = [metric_func(y_true, p) for p in predictions_list]
        
        aics = []
        for k, logL in zip(num_params_list, log_likelihoods):
            aics.append(compute_aic(logL, k, n, correction=(aic_name == 'AICc')))
        
        output_df[aic_name] = aics
        output_df[f'{aic_name}_w'] = compute_aic_weights(aics)

    # --- Standard Metric Calculation ---
    for metric in metrics_to_calc:
        if metric in METRIC_FUNCTIONS:
            output_df[metric] = [METRIC_FUNCTIONS[metric](y_true, p, groups=groups) for p in predictions_list]

    # --- Elbow Calculation ---
    if 'Elbow' in metrics_to_calc:
        metric_func = get_loss(loss_name)
        new_df = output_df.copy()
        new_df['Score'] = [metric_func(y_true, p) for p in predictions_list]

        new_df = compute_pareto_front(new_df, 'Complexity', 'Score',
                                     x_higher_is_better=False, 
                                     y_higher_is_better=False)
        elbow_values = compute_elbow_metric(new_df, 'Score', 'Complexity', metric_higher_is_better=False)
        output_df['Elbow'] = elbow_values['elbow_metric']
    output_df = output_df.drop(columns='Complexity')
    return output_df

def plot_pareto_front(dfs, x_axis: str, y_axis: str, x_higher_is_better=False, y_higher_is_better=True, scatter=False):
    import matplotlib.pyplot as plt
    fig, ax = plt.subplots(figsize=(10, 6))
    if not isinstance(dfs, list):
        dfs = [('SRRegressor', dfs)]

    colors = plt.cm.tab10.colors
    markers = ['o', '+', 'x', '^', 's', 'D', 'p', 'h', '<', '>']
    last_x = max([df[x_axis].max() for _, df in dfs]) if not x_higher_is_better else min([df[x_axis].min() for _, df in dfs])

    for i, (name, df) in enumerate(dfs):
        color = colors[i % len(colors)]
        marker = markers[i % len(markers)]
        
        if scatter:
            # Plot all points (even if not in the pareto front)
            ax.scatter(df[x_axis], df[y_axis], marker=marker, color=color, label=name)
            
        # Calculate and plot Pareto front
        pareto_df = compute_pareto_front(df, x_axis, y_axis, x_higher_is_better, y_higher_is_better)[[x_axis, y_axis]]
        last_y = pareto_df[y_axis].min() if not y_higher_is_better else pareto_df[y_axis].max()
        ax.scatter(pareto_df[x_axis], pareto_df[y_axis], marker=marker, color=color, label=name)
        if (pareto_df[x_axis].iloc[-1] != last_x) or (pareto_df[y_axis].iloc[-1] != last_y):
            pareto_df = pd.concat([pareto_df, pd.DataFrame({x_axis: [last_x], y_axis: [last_y]})], ignore_index=True)
        ax.plot(pareto_df[x_axis], pareto_df[y_axis], color=color, linestyle='--', linewidth=1.5)
        temp = y_axis.split(' ')
        if len(temp) == 1:
            elbow_name = 'Elbow'
        else:
            elbow_name = f'Elbow {temp[1]}'
        # Highlight the elbow point if available
        if elbow_name in df.columns:
            idx = df[elbow_name].idxmax()
            elbow_point = df.loc[idx]
            ax.scatter(elbow_point[x_axis], elbow_point[y_axis], s=150, c='red', marker='*', label='Best Elbow', zorder=5)

    ax.set_xlabel(x_axis)
    ax.set_ylabel(y_axis)
    ax.set_title(f'Pareto Front: {y_axis} vs. {x_axis}')
    ax.legend()
    plt.show()

def apply_threshold(y_pred, threshold):
    y_pred = np.asarray(y_pred)
    # Only threshold if predictions are not already binary
    if not np.all((y_pred == 0) | (y_pred == 1)):
        y_pred = (y_pred >= threshold).astype(int)
    return y_pred