from __future__ import annotations
from pathlib import Path
from typing import Any
import hashlib
import json
import shutil
import zipfile
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt


def _write_json(path: Path, value: Any):
    path.write_text(json.dumps(value, indent=2, ensure_ascii=False, allow_nan=False), encoding="utf-8")


def _safe_json(value):
    if isinstance(value, dict):
        return {k: _safe_json(v) for k, v in value.items()}
    if isinstance(value, list):
        return [_safe_json(v) for v in value]
    if isinstance(value, float) and not np.isfinite(value):
        return None
    return value


def _hash(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def export_run(
    output_dir: Path,
    config: dict[str, Any],
    result: dict[str, Any],
    fold_registry: pd.DataFrame,
    model_comparison: pd.DataFrame,
    fold_metrics: pd.DataFrame,
    galaxy_metrics: pd.DataFrame,
    heldout: pd.DataFrame,
    feature_ablation: pd.DataFrame,
    null_summary: pd.DataFrame,
    null_distribution: pd.DataFrame,
    comparison: dict[str, Any],
    radial_path: Path,
    signature_path: Path,
):
    output_dir.mkdir(parents=True, exist_ok=True)
    figures = output_dir / "figures"
    figures.mkdir(exist_ok=True)

    fold_registry.to_csv(output_dir / "grav_bridge_i_v0_1_0_fold_registry.csv", index=False)
    model_comparison.to_csv(output_dir / "grav_bridge_i_v0_1_0_model_comparison.csv", index=False)
    fold_metrics.to_csv(output_dir / "grav_bridge_i_v0_1_0_fold_metrics.csv", index=False)
    galaxy_metrics.to_csv(output_dir / "grav_bridge_i_v0_1_0_galaxy_metrics.csv", index=False)
    heldout.to_csv(output_dir / "grav_bridge_i_v0_1_0_heldout_predictions.csv", index=False)
    feature_ablation.to_csv(output_dir / "grav_bridge_i_v0_1_0_feature_ablation.csv", index=False)
    null_summary.to_csv(output_dir / "grav_bridge_i_v0_1_0_null_summary.csv", index=False)
    null_distribution.to_csv(output_dir / "grav_bridge_i_v0_1_0_null_distributions.csv", index=False)
    _write_json(output_dir / "grav_bridge_i_v0_1_0_results.json", _safe_json(result))
    _write_json(output_dir / "grav_bridge_i_v0_1_0_config_used.json", config)

    # Separate figures: no subplots.
    order = ["M0", "M1", "M2", "M_I", "M_P", "M_IP"]
    plot_df = model_comparison.set_index("model").loc[order]
    plt.figure(figsize=(9, 5))
    plt.bar(plot_df.index, plot_df["galaxy_balanced_rmse"])
    plt.ylabel("Galaxy-balanced RMSE")
    plt.title("Held-out Model Comparison")
    plt.tight_layout()
    plt.savefig(figures / "model_comparison_rmse.png", dpi=170)
    plt.close()

    plt.figure(figsize=(8, 5))
    plt.hist(null_distribution["signature_delta"].dropna(), bins=30)
    plt.axvline(comparison["delta_galaxy_balanced_rmse"], linestyle="--")
    plt.xlabel("M2 RMSE − permuted-signature M_IP RMSE")
    plt.ylabel("Count")
    plt.title("Cross-galaxy Signature Null")
    plt.tight_layout()
    plt.savefig(figures / "signature_null_distribution.png", dpi=170)
    plt.close()

    plt.figure(figsize=(8, 5))
    plt.hist(null_distribution["target_delta"].dropna(), bins=30)
    plt.axvline(comparison["delta_galaxy_balanced_rmse"], linestyle="--")
    plt.xlabel("M2 RMSE − M_IP RMSE under target-curve permutation")
    plt.ylabel("Count")
    plt.title("Galaxy Target-Curve Null")
    plt.tight_layout()
    plt.savefig(figures / "target_null_distribution.png", dpi=170)
    plt.close()

    plt.figure(figsize=(6, 6))
    plt.scatter(heldout[config["columns"]["target"]], heldout["pred_M_IP"], s=10, alpha=0.5)
    low = min(heldout[config["columns"]["target"]].min(), heldout["pred_M_IP"].min())
    high = max(heldout[config["columns"]["target"]].max(), heldout["pred_M_IP"].max())
    plt.plot([low, high], [low, high], linestyle="--")
    plt.xlabel("Observed log10(g_obs/g_bar)")
    plt.ylabel("Held-out M_IP prediction")
    plt.title("Held-out Structural Model Predictions")
    plt.tight_layout()
    plt.savefig(figures / "heldout_observed_vs_predicted_MIP.png", dpi=170)
    plt.close()

    result_record = f"""# GRAV-BRIDGE-I v0.1.0 Result Record

## Protocol status

**{result["protocol_status"]}**

## Verdict

**{result["verdict"]}**

{chr(10).join(f"- {reason}" for reason in result["verdict_reasons"])}

## Registered question

Do chamber-defined structural signatures of a galaxy's baryonic acceleration profile improve
prediction of its observed gravitational discrepancy beyond local baryonic acceleration,
normalized radius, and ordinary ladder statistics?

## Data boundary

- Galaxies: {result["input_validation"]["galaxies"]}
- Radial rows: {result["input_validation"]["radial_rows"]}
- Grouped folds: {config["cross_validation"]["n_splits"]}
- Cross-galaxy signature permutations: {config["nulls"]["signature_permutations"]}
- Galaxy target-curve permutations: {config["nulls"]["target_curve_permutations"]}

## Decisive comparison

- M2 galaxy-balanced RMSE: {comparison["baseline_galaxy_balanced_rmse"]:.8f}
- M_IP galaxy-balanced RMSE: {comparison["structural_galaxy_balanced_rmse"]:.8f}
- RMSE delta (M2 − M_IP): {comparison["delta_galaxy_balanced_rmse"]:.8f}
- Relative improvement: {100 * comparison["relative_rmse_improvement"]:.4f}%
- Bootstrap 95% CI: [{comparison["bootstrap_ci_lower"]:.8f}, {comparison["bootstrap_ci_upper"]:.8f}]
- Signature-null p-value: {comparison["signature_null_p_value"]:.6f}
- Target-null p-value: {comparison["target_null_p_value"]:.6f}

## Interpretive boundary

This result measures incremental out-of-sample prediction only. It is not a derivation of a
field equation, evidence of causation, a lensing result, or a demonstration that dark matter
does not exist.
"""
    (output_dir / "GRAV_BRIDGE_I_RESULT_RECORD.md").write_text(result_record, encoding="utf-8")

    manifest = {
        "radial_input": {"path": str(radial_path), "sha256": _hash(radial_path)},
        "signature_input": {"path": str(signature_path), "sha256": _hash(signature_path)},
        "outputs": {},
    }
    for path in sorted(output_dir.rglob("*")):
        if path.is_file() and path.name != "GRAV_BRIDGE_I_RUN_PACKAGE.zip":
            manifest["outputs"][str(path.relative_to(output_dir))] = _hash(path)
    _write_json(output_dir / "RUN_MANIFEST_SHA256.json", manifest)

    zip_path = output_dir / "GRAV_BRIDGE_I_RUN_PACKAGE.zip"
    with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as z:
        for path in sorted(output_dir.rglob("*")):
            if path.is_file() and path != zip_path:
                z.write(path, arcname=path.relative_to(output_dir))
