from __future__ import annotations
from typing import Any
import numpy as np
import pandas as pd
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score


def metric_bundle(y_true, y_pred, groups) -> dict[str, float]:
    y_true = np.asarray(y_true, dtype=float)
    y_pred = np.asarray(y_pred, dtype=float)
    groups = np.asarray(groups).astype(str)
    residual = y_true - y_pred
    per_group = []
    for group in np.unique(groups):
        mask = groups == group
        mse = float(np.mean(residual[mask] ** 2))
        mae = float(np.mean(np.abs(residual[mask])))
        per_group.append((group, mse, mae, int(mask.sum())))
    gb_rmse = float(np.sqrt(np.mean([row[1] for row in per_group])))
    gb_mae = float(np.mean([row[2] for row in per_group]))
    median_group_rmse = float(np.median([np.sqrt(row[1]) for row in per_group]))
    return {
        "row_rmse": float(np.sqrt(mean_squared_error(y_true, y_pred))),
        "row_mae": float(mean_absolute_error(y_true, y_pred)),
        "row_r2": float(r2_score(y_true, y_pred)),
        "galaxy_balanced_rmse": gb_rmse,
        "galaxy_balanced_mae": gb_mae,
        "median_galaxy_rmse": median_group_rmse,
    }


def per_galaxy_metrics(y_true, y_pred, groups, model_name: str) -> pd.DataFrame:
    y_true = np.asarray(y_true, dtype=float)
    y_pred = np.asarray(y_pred, dtype=float)
    groups = np.asarray(groups).astype(str)
    rows = []
    for group in sorted(np.unique(groups)):
        mask = groups == group
        err = y_true[mask] - y_pred[mask]
        rows.append({
            "galaxy_id": group,
            "model": model_name,
            "n_rows": int(mask.sum()),
            "rmse": float(np.sqrt(np.mean(err ** 2))),
            "mse": float(np.mean(err ** 2)),
            "mae": float(np.mean(np.abs(err))),
            "bias": float(np.mean(y_pred[mask] - y_true[mask])),
        })
    return pd.DataFrame(rows)


def calibration(y_true, y_pred) -> dict[str, float]:
    y_true = np.asarray(y_true, dtype=float)
    y_pred = np.asarray(y_pred, dtype=float)
    if np.std(y_pred) == 0:
        return {"calibration_intercept": float("nan"), "calibration_slope": float("nan")}
    slope, intercept = np.polyfit(y_pred, y_true, 1)
    return {"calibration_intercept": float(intercept), "calibration_slope": float(slope)}


def paired_galaxy_bootstrap(
    galaxy_metrics_base: pd.DataFrame,
    galaxy_metrics_model: pd.DataFrame,
    repetitions: int,
    seed: int,
    confidence: float,
) -> dict[str, float]:
    base = galaxy_metrics_base.set_index("galaxy_id")["mse"]
    model = galaxy_metrics_model.set_index("galaxy_id")["mse"]
    common = base.index.intersection(model.index)
    b = base.loc[common].to_numpy(float)
    m = model.loc[common].to_numpy(float)
    rng = np.random.default_rng(seed)
    deltas = np.empty(repetitions, dtype=float)
    n = len(common)
    for i in range(repetitions):
        idx = rng.integers(0, n, size=n)
        deltas[i] = np.sqrt(np.mean(b[idx])) - np.sqrt(np.mean(m[idx]))
    alpha = 1.0 - confidence
    return {
        "bootstrap_delta_mean": float(np.mean(deltas)),
        "bootstrap_ci_lower": float(np.quantile(deltas, alpha / 2)),
        "bootstrap_ci_upper": float(np.quantile(deltas, 1 - alpha / 2)),
        "bootstrap_positive_fraction": float(np.mean(deltas > 0)),
    }
