from __future__ import annotations
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable
import json
import platform
import sys
import time
import numpy as np
import pandas as pd
import sklearn

from .config import load_config, resolve_registered_path
from .data import PreparedData, load_and_validate, sha256_file
from .metrics import calibration, metric_bundle, paired_galaxy_bootstrap, per_galaxy_metrics
from .models import FeatureSpec, build_pipeline, feature_registry
from .splits import build_group_folds, equal_group_weights
from .exports import export_run

Progress = Callable[[float, str], None]


@dataclass
class ModelRun:
    name: str
    predictions: np.ndarray
    metrics: dict[str, float]
    fold_metrics: pd.DataFrame
    galaxy_metrics: pd.DataFrame
    train_metrics: dict[str, float]


def _notify(callback: Progress | None, fraction: float, message: str):
    if callback:
        callback(max(0.0, min(1.0, fraction)), message)


def cross_validate_model(
    data: pd.DataFrame,
    y: np.ndarray,
    groups: np.ndarray,
    folds,
    spec: FeatureSpec,
    config: dict[str, Any],
) -> ModelRun:
    predictions = np.full(len(data), np.nan, dtype=float)
    fold_rows = []
    for fold_id, (train_idx, test_idx) in enumerate(folds, start=1):
        pipe = build_pipeline(spec, config)
        fit_kwargs = {}
        if config["cross_validation"].get("equal_galaxy_fit_weights", True):
            fit_kwargs["model__sample_weight"] = equal_group_weights(groups[train_idx])
        pipe.fit(data.iloc[train_idx], y[train_idx], **fit_kwargs)
        pred = pipe.predict(data.iloc[test_idx])
        predictions[test_idx] = pred
        fold_metric = metric_bundle(y[test_idx], pred, groups[test_idx])
        fold_rows.append({"model": spec.name, "fold": fold_id, **fold_metric,
                          "test_rows": int(len(test_idx)),
                          "test_galaxies": int(len(np.unique(groups[test_idx])))})
    if not np.isfinite(predictions).all():
        raise RuntimeError(f"Model {spec.name} did not produce complete OOF predictions.")
    overall = metric_bundle(y, predictions, groups)
    overall.update(calibration(y, predictions))
    galaxy = per_galaxy_metrics(y, predictions, groups, spec.name)

    full_pipe = build_pipeline(spec, config)
    fit_kwargs = {}
    if config["cross_validation"].get("equal_galaxy_fit_weights", True):
        fit_kwargs["model__sample_weight"] = equal_group_weights(groups)
    full_pipe.fit(data, y, **fit_kwargs)
    train_pred = full_pipe.predict(data)
    train_metrics = metric_bundle(y, train_pred, groups)
    return ModelRun(spec.name, predictions, overall, pd.DataFrame(fold_rows), galaxy, train_metrics)


def cross_validate_score_only(
    data: pd.DataFrame,
    y: np.ndarray,
    groups: np.ndarray,
    folds,
    spec: FeatureSpec,
    config: dict[str, Any],
) -> float:
    predictions = np.full(len(data), np.nan, dtype=float)
    for train_idx, test_idx in folds:
        pipe = build_pipeline(spec, config)
        fit_kwargs = {}
        if config["cross_validation"].get("equal_galaxy_fit_weights", True):
            fit_kwargs["model__sample_weight"] = equal_group_weights(groups[train_idx])
        pipe.fit(data.iloc[train_idx], y[train_idx], **fit_kwargs)
        predictions[test_idx] = pipe.predict(data.iloc[test_idx])
    if not np.isfinite(predictions).all():
        raise RuntimeError(f"Model {spec.name} did not produce complete OOF predictions.")
    return metric_bundle(y, predictions, groups)["galaxy_balanced_rmse"]


def _signature_permuted_data(data: pd.DataFrame, structural_columns: list[str], group_col: str, rng) -> pd.DataFrame:
    result = data.copy()
    galaxy_rows = result.groupby(group_col, sort=True).first()
    galaxy_ids = galaxy_rows.index.to_numpy()
    sources = rng.permutation(galaxy_ids)
    mapping = dict(zip(galaxy_ids, sources))
    for col in structural_columns:
        source_values = galaxy_rows[col]
        result[col] = result[group_col].map(lambda g: source_values.loc[mapping[g]])
    return result


def _target_curve_permutation(data: pd.DataFrame, y: np.ndarray, group_col: str, radius_col: str, rng) -> np.ndarray:
    frame = pd.DataFrame({group_col: data[group_col].astype(str), radius_col: data[radius_col].astype(float), "y": y})
    curves = {}
    for galaxy, group in frame.groupby(group_col):
        ordered = group.sort_values(radius_col)
        curves[galaxy] = (ordered[radius_col].to_numpy(float), ordered["y"].to_numpy(float))
    galaxies = np.array(sorted(curves))
    sources = rng.permutation(galaxies)
    mapping = dict(zip(galaxies, sources))
    y_perm = np.empty(len(frame), dtype=float)
    for galaxy, idx in frame.groupby(group_col).groups.items():
        source_r, source_y = curves[mapping[galaxy]]
        receiver_r = frame.loc[idx, radius_col].to_numpy(float)
        y_perm[np.asarray(idx, dtype=int)] = np.interp(receiver_r, source_r, source_y)
    return y_perm


def _one_sided_p(null_values: np.ndarray, observed: float) -> float:
    return float((1 + np.sum(null_values >= observed)) / (len(null_values) + 1))


def determine_verdict(comparison: dict[str, Any], config: dict[str, Any]) -> tuple[str, list[str]]:
    delta = comparison["delta_galaxy_balanced_rmse"]
    rel = comparison["relative_rmse_improvement"]
    train_delta = comparison["train_delta_galaxy_balanced_rmse"]
    lower = comparison["bootstrap_ci_lower"]
    p_sig = comparison["signature_null_p_value"]
    p_target = comparison["target_null_p_value"]
    rules = config["verdict"]
    reasons = []

    if delta <= 0:
        if train_delta > 0:
            return "NON_GENERALIZING_SIGNAL", [
                "The structural model improved full-data fit but not galaxy-held-out performance."
            ]
        return "NO_INCREMENTAL_SIGNAL", [
            "M_IP did not improve galaxy-balanced held-out RMSE over M2."
        ]

    if p_sig > rules["weak_p_value"] and p_target > rules["weak_p_value"]:
        return "NULL_CONSISTENT", [
            "The observed incremental gain is not distinguished from either registered null."
        ]

    robust = (
        rel >= rules["minimum_relative_rmse_improvement"]
        and p_sig <= rules["robust_p_value"]
        and p_target <= rules["robust_p_value"]
        and (lower > 0 if rules.get("require_positive_bootstrap_lower_bound", True) else True)
    )
    if robust:
        reasons.extend([
            "M_IP improved galaxy-balanced held-out RMSE over M2.",
            "The paired galaxy bootstrap lower bound is positive.",
            "Both registered null p-values meet the robust threshold.",
        ])
        return "ROBUST_INCREMENTAL_SIGNAL", reasons

    reasons.append("A positive held-out increment exists but one or more robustness gates were not met.")
    return "WEAK_INCREMENTAL_SIGNAL", reasons


def run_analysis(
    radial_path: str | Path,
    signature_path: str | Path,
    config_path: str | Path | None,
    output_dir: str | Path,
    overrides: dict[str, Any] | None = None,
    progress: Progress | None = None,
) -> dict[str, Any]:
    started = time.time()
    registered_config = load_config()
    config = load_config(config_path, overrides)
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    _notify(progress, 0.02, "Validating registered inputs")
    prepared = load_and_validate(radial_path, signature_path, config)
    registered_radial = resolve_registered_path(registered_config, "radial")
    registered_signatures = resolve_registered_path(registered_config, "signatures")
    registered_inputs_match = (
        sha256_file(radial_path) == sha256_file(registered_radial)
        and sha256_file(signature_path) == sha256_file(registered_signatures)
    )
    protocol_status = (
        "REGISTERED_EXECUTION"
        if config == registered_config and registered_inputs_match
        else "ENGINEERING_OVERRIDE"
    )
    data = prepared.joined.reset_index(drop=True)
    group_col = config["columns"]["group"]
    target_col = config["columns"]["target"]
    radius_col = config["columns"]["radius"]
    y = data[target_col].to_numpy(float)
    groups = data[group_col].astype(str).to_numpy()
    folds, fold_registry = build_group_folds(data, config)
    registry = feature_registry(config)

    primary_models = ["M0", "M1", "M2", "M_I", "M_P", "M_IP"]
    model_runs: dict[str, ModelRun] = {}
    for i, name in enumerate(primary_models):
        _notify(progress, 0.06 + 0.22 * i / len(primary_models), f"Running grouped validation: {name}")
        model_runs[name] = cross_validate_model(data, y, groups, folds, registry[name], config)

    model_comparison_rows = []
    fold_metrics = []
    galaxy_metrics = []
    for name in primary_models:
        run = model_runs[name]
        row = {"model": name, **run.metrics}
        row.update({f"train_{k}": v for k, v in run.train_metrics.items()})
        model_comparison_rows.append(row)
        fold_metrics.append(run.fold_metrics)
        galaxy_metrics.append(run.galaxy_metrics)
    model_comparison = pd.DataFrame(model_comparison_rows)
    fold_metrics_df = pd.concat(fold_metrics, ignore_index=True)
    galaxy_metrics_df = pd.concat(galaxy_metrics, ignore_index=True)

    _notify(progress, 0.30, "Running feature ablations")
    ablation_names = ["M_I", "M_P", "M_IP", "M_IP_MINUS_AK", "M_IP_MINUS_RHO", "M_IP_MINUS_P_NUMERIC", "M_IP_MINUS_P_CATEGORICAL"]
    ablation_rows = []
    for name in ablation_names:
        run = model_runs.get(name) or cross_validate_model(data, y, groups, folds, registry[name], config)
        ablation_rows.append({"model": name, **run.metrics})
    feature_ablation = pd.DataFrame(ablation_rows)

    baseline = model_runs["M2"]
    combined = model_runs["M_IP"]
    bootstrap = paired_galaxy_bootstrap(
        baseline.galaxy_metrics,
        combined.galaxy_metrics,
        int(config["uncertainty"]["galaxy_bootstrap_repetitions"]),
        int(config["random_seed"]) + 701,
        float(config["uncertainty"]["confidence_level"]),
    )
    observed_delta = baseline.metrics["galaxy_balanced_rmse"] - combined.metrics["galaxy_balanced_rmse"]

    fg = config["feature_groups"]
    structural_columns = (
        fg["struc_i_Ak_numeric"] + fg["struc_i_rho_numeric"] + fg["struc_i_categorical"]
        + fg["struc_perc_numeric"] + fg["struc_perc_categorical"]
    )

    sig_n = int(config["nulls"]["signature_permutations"])
    target_n = int(config["nulls"]["target_curve_permutations"])
    rng_sig = np.random.default_rng(int(config["random_seed"]) + 1101)
    sig_deltas = np.empty(sig_n, dtype=float)
    for i in range(sig_n):
        permuted = _signature_permuted_data(data, structural_columns, group_col, rng_sig)
        perm_rmse = cross_validate_score_only(permuted, y, groups, folds, registry["M_IP"], config)
        sig_deltas[i] = baseline.metrics["galaxy_balanced_rmse"] - perm_rmse
        if i % max(1, sig_n // 20) == 0:
            _notify(progress, 0.38 + 0.20 * (i + 1) / max(1, sig_n), f"Signature null {i+1}/{sig_n}")

    rng_target = np.random.default_rng(int(config["random_seed"]) + 2201)
    target_deltas = np.empty(target_n, dtype=float)
    for i in range(target_n):
        y_perm = _target_curve_permutation(data, y, group_col, radius_col, rng_target)
        base_rmse = cross_validate_score_only(data, y_perm, groups, folds, registry["M2"], config)
        comb_rmse = cross_validate_score_only(data, y_perm, groups, folds, registry["M_IP"], config)
        target_deltas[i] = base_rmse - comb_rmse
        if i % max(1, target_n // 20) == 0:
            _notify(progress, 0.59 + 0.25 * (i + 1) / max(1, target_n), f"Target-curve null {i+1}/{target_n}")

    comparison = {
        "baseline_model": "M2",
        "structural_model": "M_IP",
        "baseline_galaxy_balanced_rmse": baseline.metrics["galaxy_balanced_rmse"],
        "structural_galaxy_balanced_rmse": combined.metrics["galaxy_balanced_rmse"],
        "delta_galaxy_balanced_rmse": observed_delta,
        "relative_rmse_improvement": observed_delta / baseline.metrics["galaxy_balanced_rmse"],
        "train_delta_galaxy_balanced_rmse": baseline.train_metrics["galaxy_balanced_rmse"] - combined.train_metrics["galaxy_balanced_rmse"],
        **bootstrap,
        "signature_null_p_value": _one_sided_p(sig_deltas, observed_delta),
        "target_null_p_value": _one_sided_p(target_deltas, observed_delta),
        "signature_null_mean_delta": float(np.mean(sig_deltas)) if sig_n else float("nan"),
        "target_null_mean_delta": float(np.mean(target_deltas)) if target_n else float("nan"),
    }
    verdict, verdict_reasons = determine_verdict(comparison, config)

    heldout = data[["row_id", group_col, "radius_kpc", radius_col, config["columns"]["local_gbar"], target_col]].copy()
    for name in primary_models:
        heldout[f"pred_{name}"] = model_runs[name].predictions
        heldout[f"residual_{name}"] = y - model_runs[name].predictions

    null_distribution = pd.DataFrame({
        "permutation": np.arange(1, max(sig_n, target_n) + 1),
        "signature_delta": np.pad(sig_deltas, (0, max(sig_n, target_n)-sig_n), constant_values=np.nan),
        "target_delta": np.pad(target_deltas, (0, max(sig_n, target_n)-target_n), constant_values=np.nan),
    })
    null_summary = pd.DataFrame([
        {"null": "cross_galaxy_signature", "repetitions": sig_n, "mean_delta": float(np.mean(sig_deltas)), "sd_delta": float(np.std(sig_deltas, ddof=1)), "p_value": comparison["signature_null_p_value"]},
        {"null": "galaxy_target_curve", "repetitions": target_n, "mean_delta": float(np.mean(target_deltas)), "sd_delta": float(np.std(target_deltas, ddof=1)), "p_value": comparison["target_null_p_value"]},
    ])

    result = {
        "chamber": config["chamber"],
        "version": config["version"],
        "timestamp_utc": datetime.now(timezone.utc).isoformat(),
        "protocol_status": protocol_status,
        "verdict": verdict,
        "verdict_reasons": verdict_reasons,
        "input_validation": prepared.report,
        "comparison": comparison,
        "model_metrics": {name: model_runs[name].metrics for name in primary_models},
        "run_seconds": float(time.time() - started),
        "environment": {
            "python": sys.version.split()[0],
            "platform": platform.platform(),
            "numpy": np.__version__,
            "pandas": pd.__version__,
            "scikit_learn": sklearn.__version__,
        },
        "interpretive_boundary": config["interpretive_boundary"],
    }

    _notify(progress, 0.88, "Writing result records and figures")
    export_run(
        output_dir=output_dir,
        config=config,
        result=result,
        fold_registry=fold_registry,
        model_comparison=model_comparison,
        fold_metrics=fold_metrics_df,
        galaxy_metrics=galaxy_metrics_df,
        heldout=heldout,
        feature_ablation=feature_ablation,
        null_summary=null_summary,
        null_distribution=null_distribution,
        comparison=comparison,
        radial_path=Path(radial_path),
        signature_path=Path(signature_path),
    )
    _notify(progress, 1.0, f"Complete: {verdict}")
    return result
