"""Analyze a trained HiCon-GNN checkpoint."""

from __future__ import annotations

import argparse
import hashlib
import os
from pathlib import Path
from typing import List, Sequence, Tuple

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
from torch.utils.data import DataLoader

from concrete_gnn import (
    DEFAULT_SCHEMA,
    TABULAR_DIM,
    IntegratedMultiscaleModel,
    SchemaSpec,
    Standardizer,
    StrengthHead,
    TabularANN,
    ConcreteMixDataset,
    collate_multiscale,
    collect_strength_predictions,
    load_split_indices,
    random_split_indices,
    read_table,
    split_indices_sha256,
    strength_column,
)


MIX_FEATURE_COLS: List[str] = [
    "cement_kg_m3",
    "slag_kg_m3",
    "fly_ash_kg_m3",
    "silica_fume_kg_m3",
    "metakaolin_kg_m3",
    "limestone_powder_kg_m3",
    "other_scm_kg_m3",
    "water_kg_m3",
    "superplasticizer_kg_m3",
    "coarse_aggregate_kg_m3",
    "fine_aggregate_kg_m3",
    "fibre_content_kg_m3",
    "fibre_length_mm",
    "fibre_diameter_mm",
    "fibre_tensile_strength_mpa",
    "fibre_modulus_gpa",
    "age_days",
]

OUTPUT_ROOT = Path("outputs")


def _make_base_model(name: str) -> torch.nn.Module:
    if name == "hierarchical":
        return IntegratedMultiscaleModel(schema=DEFAULT_SCHEMA)
    if name == "tabular_ann":
        return TabularANN(in_dim=TABULAR_DIM, schema=DEFAULT_SCHEMA)
    raise ValueError(f"Unknown model name {name!r}")


def _resolve_device(force_cpu: bool) -> torch.device:
    if force_cpu:
        return torch.device("cpu")
    if torch.cuda.is_available():
        return torch.device("cuda")
    mps = getattr(torch.backends, "mps", None)
    if mps is not None and mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def _load_checkpoint(
    path: Path, device: torch.device
) -> Tuple[StrengthHead, Standardizer, str, dict]:
    checkpoint = torch.load(path, map_location=device, weights_only=True)
    name = checkpoint.get("model_name") or path.stem
    standardizer = Standardizer(
        mean=torch.tensor(float(checkpoint["target_mean"])),
        std=torch.tensor(float(checkpoint["target_std"])),
    )
    base = _make_base_model(name)
    base.load_state_dict(checkpoint["base_model_state_dict"])
    model = StrengthHead(base, standardizer).to(device)
    model.eval()
    return model, standardizer, name, checkpoint.get("run_metadata", {})


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


def _build_loader(
    df: pd.DataFrame,
    indices: Sequence[int],
    standardizer: Standardizer,
    seed: int,
    batch_size: int,
    schema: SchemaSpec,
) -> DataLoader:
    ds = ConcreteMixDataset(
        df, indices, seed=seed, target_standardizer=standardizer, schema=schema
    )
    return DataLoader(ds, batch_size=batch_size, shuffle=False, collate_fn=collate_multiscale)


def _predict(
    model: StrengthHead, loader: DataLoader, device: torch.device
) -> Tuple[np.ndarray, np.ndarray]:
    pred, true = collect_strength_predictions(model, loader, device)
    return pred.numpy(), true.numpy()


@torch.no_grad()
def _predict_all(
    model: StrengthHead, loader: DataLoader, device: torch.device
) -> Tuple[np.ndarray, np.ndarray]:
    """Return all three predictions and their available measured targets."""

    model.eval()
    predictions, targets = [], []
    for batch in loader:
        batch = batch.to(device)
        model(batch)
        predictions.append(model.concrete_pred.detach().cpu())
        compressive = model.standardizer.decode(batch.concrete_target[:, 0]).cpu()
        targets.append(
            torch.cat([compressive[:, None], batch.concrete_aux_target.cpu()], dim=1)
        )
    return torch.cat(predictions).numpy(), torch.cat(targets).numpy()


# Permutation importance


def run_permutation_importance(
    model: StrengthHead,
    standardizer: Standardizer,
    df: pd.DataFrame,
    test_idx: Sequence[int],
    args,
    device: torch.device,
    out_dir: Path,
) -> None:
    print("\n[perm] Permutation feature importance")
    test_df = df.iloc[list(test_idx)].reset_index(drop=True)
    base_loader = _build_loader(test_df, range(len(test_df)), standardizer,
                                args.seed + 20, args.batch_size,
                                schema=model.base.schema)
    base_pred, true = _predict(model, base_loader, device)
    base_rmse = float(np.sqrt(((base_pred - true) ** 2).mean()))
    print(f"  baseline test RMSE = {base_rmse:.3f} MPa")

    rng = np.random.default_rng(args.seed)
    rows = []
    for col in MIX_FEATURE_COLS:
        permuted = test_df.copy()
        permuted[col] = rng.permutation(permuted[col].to_numpy())
        loader = _build_loader(permuted, range(len(permuted)), standardizer,
                               args.seed + 20, args.batch_size,
                               schema=model.base.schema)
        pred, _ = _predict(model, loader, device)
        rmse = float(np.sqrt(((pred - true) ** 2).mean()))
        delta = rmse - base_rmse
        rows.append({"feature": col, "rmse": rmse, "rmse_increase": delta})
        print(f"  {col:<32} rmse={rmse:7.3f}  delta={delta:+.3f}")
    table = pd.DataFrame(rows).sort_values("rmse_increase", ascending=False)
    csv_path = out_dir / "permutation_importance.csv"
    table.to_csv(csv_path, index=False)
    print(f"  wrote {csv_path}")

    fig, ax = plt.subplots(figsize=(7, 6))
    ax.barh(table["feature"], table["rmse_increase"], color="steelblue")
    ax.invert_yaxis()
    ax.set_xlabel("RMSE increase when feature shuffled (MPa)")
    ax.set_title("Permutation feature importance")
    ax.axvline(0, color="black", linewidth=0.5)
    fig.tight_layout()
    fig_path = out_dir / "permutation_importance.png"
    fig.savefig(fig_path, dpi=150)
    plt.close(fig)
    print(f"  wrote {fig_path}")


# Residual diagnostics


def run_residual_diagnostics(
    model: StrengthHead,
    standardizer: Standardizer,
    df: pd.DataFrame,
    test_idx: Sequence[int],
    args,
    device: torch.device,
    out_dir: Path,
) -> None:
    print("\n[residuals] Parity + residual diagnostics")
    full_test_df = df.iloc[list(test_idx)].reset_index(drop=True)
    loader = _build_loader(full_test_df, range(len(full_test_df)), standardizer,
                           args.seed + 20, args.batch_size,
                           schema=model.base.schema)
    all_pred, all_true = _predict_all(model, loader, device)
    has_target = np.isfinite(all_true[:, 0])
    pred = all_pred[has_target, 0]
    true = all_true[has_target, 0]
    resid = pred - true
    test_df = full_test_df[has_target].reset_index(drop=True)

    fig, ax = plt.subplots(figsize=(6, 6))
    ax.scatter(true, pred, s=12, alpha=0.5, color="steelblue")
    lo = float(min(true.min(), pred.min())) - 5
    hi = float(max(true.max(), pred.max())) + 5
    ax.plot([lo, hi], [lo, hi], "k--", linewidth=1)
    ax.set_xlim(lo, hi); ax.set_ylim(lo, hi)
    ax.set_xlabel("true compressive strength (MPa)")
    ax.set_ylabel("predicted compressive strength (MPa)")
    ax.set_title("Compressive-strength parity")
    fig.tight_layout()
    fig.savefig(out_dir / "residuals_parity.png", dpi=150)
    plt.close(fig)

    fig, ax = plt.subplots(figsize=(7, 4))
    ax.scatter(pred, resid, s=12, alpha=0.5, color="steelblue")
    ax.axhline(0, color="black", linewidth=0.5)
    ax.set_xlabel("predicted (MPa)")
    ax.set_ylabel("residual = pred - true (MPa)")
    ax.set_title("Residuals vs predicted")
    fig.tight_layout()
    fig.savefig(out_dir / "residuals_vs_pred.png", dpi=150)
    plt.close(fig)

    diag_cols = ["cement_kg_m3", "water_kg_m3", "fibre_content_kg_m3", "age_days"]
    fig, axes = plt.subplots(2, 2, figsize=(10, 7), sharey=True)
    for ax, col in zip(axes.flat, diag_cols):
        ax.scatter(test_df[col], resid, s=10, alpha=0.5, color="steelblue")
        ax.axhline(0, color="black", linewidth=0.5)
        ax.set_xlabel(col); ax.set_ylabel("residual (MPa)")
    fig.suptitle("Residuals vs mix-design features")
    fig.tight_layout()
    fig.savefig(out_dir / "residuals_vs_features.png", dpi=150)
    plt.close(fig)

    predictions = pd.DataFrame({"row_index": list(test_idx)})
    target_names = ("compressive", "tensile", "flexural")
    summary_rows = []
    for col, target_name in enumerate(target_names):
        measured = all_true[:, col]
        predicted = all_pred[:, col]
        mask = np.isfinite(measured)
        predictions[f"measured_{target_name}_mpa"] = measured
        predictions[f"predicted_{target_name}_mpa"] = predicted
        predictions[f"residual_{target_name}_mpa"] = np.where(
            mask, predicted - measured, np.nan
        )
        if mask.any():
            residual = predicted[mask] - measured[mask]
            ss_res = float(np.sum(residual ** 2))
            ss_tot = float(np.sum((measured[mask] - measured[mask].mean()) ** 2))
            summary_rows.append(
                {
                    "target": target_name,
                    "n": int(mask.sum()),
                    "rmse": float(np.sqrt(np.mean(residual ** 2))),
                    "mae": float(np.mean(np.abs(residual))),
                    "bias": float(np.mean(residual)),
                    "r2": 1.0 - ss_res / max(ss_tot, 1.0e-12),
                }
            )
    predictions.to_csv(out_dir / "reference_predictions.csv", index=False)
    summary = pd.DataFrame(summary_rows)
    summary.to_csv(out_dir / "residuals_summary.csv", index=False)
    print(summary.to_string(index=False))
    print("  wrote residual plots, reference_predictions.csv and residuals_summary.csv")


# KernelSHAP


def run_shap(
    model: StrengthHead,
    standardizer: Standardizer,
    df: pd.DataFrame,
    train_idx: Sequence[int],
    test_idx: Sequence[int],
    args,
    device: torch.device,
    out_dir: Path,
) -> None:
    try:
        import shap
    except ModuleNotFoundError as exc:
        raise SystemExit("Install requirements.txt to run KernelSHAP") from exc

    print("\n[shap] KernelSHAP on mix-design columns")
    print(
        f"  background={args.shap_background}  samples={args.shap_samples}  "
        f"nsamples={args.shap_nsamples}"
    )
    rng = np.random.default_rng(args.seed)

    train_rows = df.iloc[list(train_idx)].reset_index(drop=True)
    test_rows = df.iloc[list(test_idx)].reset_index(drop=True)

    train_features = train_rows[MIX_FEATURE_COLS].apply(pd.to_numeric, errors="coerce")
    medians = train_features.median().fillna(0.0)
    train_X = train_features.fillna(medians).to_numpy(dtype=np.float32)
    bg_n = min(args.shap_background, len(train_X))
    background = shap.kmeans(train_X, bg_n)

    n_explain = min(args.shap_samples, len(test_rows))
    sample_idx = rng.choice(len(test_rows), size=n_explain, replace=False)
    sample_df = test_rows.iloc[sample_idx].reset_index(drop=True)
    sample_X = (
        sample_df[MIX_FEATURE_COLS]
        .apply(pd.to_numeric, errors="coerce")
        .fillna(medians)
        .to_numpy(dtype=np.float32)
    )

    template = test_rows.iloc[0].copy()

    def model_fn(rows: np.ndarray) -> np.ndarray:
        rows = np.asarray(rows, dtype=np.float32)
        coalition_df = pd.DataFrame(np.tile(template.to_numpy(), (len(rows), 1)),
                                    columns=df.columns)
        for i, col in enumerate(MIX_FEATURE_COLS):
            coalition_df[col] = rows[:, i]
        coalition_df[strength_column(coalition_df)] = 0.0
        loader = _build_loader(coalition_df, range(len(coalition_df)),
                               standardizer, args.seed + 999, args.batch_size,
                               schema=model.base.schema)
        pred, _ = _predict(model, loader, device)
        return pred

    explainer = shap.KernelExplainer(model_fn, background)
    shap_values = explainer.shap_values(sample_X, nsamples=args.shap_nsamples,
                                        silent=False)
    shap_values = np.asarray(shap_values)
    print(f"  computed SHAP values, shape={shap_values.shape}")

    out_csv = out_dir / "shap_values.csv"
    pd.DataFrame(shap_values, columns=MIX_FEATURE_COLS).assign(
        true_mpa=sample_df["compressive_strength_mpa"].to_numpy(),
    ).to_csv(out_csv, index=False)

    # summary plots
    shap.summary_plot(shap_values, sample_X, feature_names=MIX_FEATURE_COLS,
                      show=False, plot_type="dot")
    plt.tight_layout()
    plt.savefig(out_dir / "shap_summary_beeswarm.png", dpi=150)
    plt.close()

    shap.summary_plot(shap_values, sample_X, feature_names=MIX_FEATURE_COLS,
                      show=False, plot_type="bar")
    plt.tight_layout()
    plt.savefig(out_dir / "shap_summary_bar.png", dpi=150)
    plt.close()
    print(f"  wrote {out_csv}, shap_summary_beeswarm.png, shap_summary_bar.png")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--checkpoint",
                        default=str(OUTPUT_ROOT / "checkpoints_reference" / "hierarchical.pt"),
                        help="Path to a *.pt checkpoint saved by train_real_data.py")
    parser.add_argument("--data", default=os.path.join("data", "HiConGNN_processed_database.csv"))
    parser.add_argument(
        "--split", default=None, help="Manifest path required by manifest checkpoints"
    )
    parser.add_argument("--out", default=str(OUTPUT_ROOT / "analysis"))
    parser.add_argument("--seed", type=int, default=17,
                        help="Graph-generation and analysis seed used during training")
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument("--analyses", default="residuals",
                        help="Comma-separated subset of: residuals, perm, shap")
    parser.add_argument("--shap-samples", type=int, default=30,
                        help="Number of test rows to explain with KernelSHAP")
    parser.add_argument("--shap-background", type=int, default=20,
                        help="Background coalition size (k-means summary of train set)")
    parser.add_argument("--shap-nsamples", type=int, default=200,
                        help="KernelSHAP coalitions per sample")
    parser.add_argument("--cpu", action="store_true")
    args = parser.parse_args()

    device = _resolve_device(args.cpu)
    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    print(f"device={device}  checkpoint={args.checkpoint}  out={out_dir}")
    model, standardizer, name, metadata = _load_checkpoint(Path(args.checkpoint), device)
    print(f"  loaded model_name={name}  target_mean={float(standardizer.mean):.3f}  "
          f"target_std={float(standardizer.std):.3f}")

    expected_data = metadata.get("data_sha256")
    if expected_data and _sha256(args.data) != expected_data:
        raise ValueError("data file does not match the checkpoint metadata")

    df = read_table(args.data)
    split_strategy = metadata.get("split_strategy")
    if split_strategy is None and metadata.get("split_sha256"):
        split_strategy = "manifest"

    if split_strategy == "row_random":
        if args.split:
            raise ValueError("--split is only valid for manifest checkpoints")
        split_seed = metadata.get("split_seed")
        if split_seed is None:
            raise ValueError("row-random checkpoint is missing split_seed")
        train_idx, val_idx, test_idx = random_split_indices(len(df), int(split_seed))
    elif split_strategy == "manifest":
        if not args.split:
            raise ValueError("manifest checkpoint requires --split")
        expected_split = metadata.get("split_sha256")
        if not expected_split:
            raise ValueError("manifest checkpoint is missing split_sha256")
        if _sha256(args.split) != expected_split:
            raise ValueError("split file does not match the checkpoint metadata")
        train_idx, val_idx, test_idx = load_split_indices(args.split, len(df))
    else:
        raise ValueError(f"unknown or missing split_strategy: {split_strategy!r}")

    expected_indices = metadata.get("split_indices_sha256")
    actual_indices = split_indices_sha256(train_idx, val_idx, test_idx)
    if expected_indices and actual_indices != expected_indices:
        raise ValueError("split indices do not match the checkpoint metadata")
    if split_strategy == "row_random" and not expected_indices:
        raise ValueError("row-random checkpoint is missing split_indices_sha256")

    print(
        f"  rows={len(df)}  split={split_strategy}  train={len(train_idx)}  "
        f"val={len(val_idx)}  test={len(test_idx)}"
    )

    todo = [a.strip() for a in args.analyses.split(",") if a.strip()]
    unknown = sorted(set(todo) - {"perm", "residuals", "shap"})
    if unknown:
        parser.error(f"unknown analyses: {', '.join(unknown)}")
    if "perm" in todo:
        run_permutation_importance(model, standardizer, df, test_idx, args, device, out_dir)
    if "residuals" in todo:
        run_residual_diagnostics(model, standardizer, df, test_idx, args, device, out_dir)
    if "shap" in todo:
        run_shap(model, standardizer, df, train_idx, test_idx, args, device, out_dir)


if __name__ == "__main__":
    main()
