"""
CausalGrok — M1: Layer-wise Linear Probing
Nilesh

The mechanistic claim:
    Before grokking: hospital probe HIGH (model uses stain shortcut),
                     tumor probe LOW in early layers
    At transition:   hospital probe DROPS, tumor probe RISES
    After grokking:  inverted — tumor high, hospital low

    If OOD acc jump + hospital probe drop + tumor probe rise
    all happen at the same epoch → mechanistic claim proven.
    That's Figure 2 of the paper.

Usage:
    # Run on all saved checkpoints from a run
    python -m experiments.mechinterp_m1 \
        --run_dir experiments/runs/<run_id> \
        --data_root data/wilds

    # Run on latest checkpoint only (quick check while training)
    python -m experiments.mechinterp_m1 \
        --run_dir experiments/runs/<run_id> \
        --data_root data/wilds \
        --latest_only

    # Run on ALL camelyon_v2 grokking runs
    python -m experiments.mechinterp_m1 \
        --all_runs \
        --data_root data/wilds

Output per run:
    experiments/runs/<run_id>/mechinterp/
        m1_probe_heatmap.png       ← epoch × layer, hospital probe acc
        m1_tumor_heatmap.png       ← epoch × layer, tumor probe acc
        m1_probe_curves.png        ← hospital vs tumor probe over epochs (layer 4)
        m1_probe_data.json         ← raw numbers for paper tables
"""

from __future__ import annotations

import argparse
import glob
import json
import os
from typing import Dict, List, Optional

import numpy as np
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Subset
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
import matplotlib
import matplotlib.pyplot as plt
import timm
import warnings
warnings.filterwarnings("ignore")

matplotlib.rcParams.update({"font.size": 11, "figure.dpi": 150})


# ──────────────────────────────────────────────
# RESNET-18 LAYER HOOKS
# Extract features after each of the 6 measurable stages:
#   stem → layer1 → layer2 → layer3 → layer4 → avgpool
# ──────────────────────────────────────────────

LAYER_NAMES = [
    "stem",     # After initial conv + bn + relu + maxpool
    "layer1",   # ResNet block 1 (64 channels)
    "layer2",   # ResNet block 2 (128 channels)
    "layer3",   # ResNet block 3 (256 channels)
    "layer4",   # ResNet block 4 (512 channels)
    "avgpool",  # Global average pool — penultimate representation
]


def register_hooks(model):
    """
    Register forward hooks on all 6 extraction points.
    Returns (hooks, features_dict).
    """
    features = {name: [] for name in LAYER_NAMES}
    hooks    = []

    def make_hook(name):
        def hook_fn(module, input, output):
            if output.dim() == 4:
                feat = output.mean(dim=[2, 3])
            else:
                feat = output.view(output.size(0), -1)
            features[name].append(feat.detach().cpu())
        return hook_fn

    hooks.append(model.maxpool.register_forward_hook(make_hook("stem")))
    hooks.append(model.layer1.register_forward_hook(make_hook("layer1")))
    hooks.append(model.layer2.register_forward_hook(make_hook("layer2")))
    hooks.append(model.layer3.register_forward_hook(make_hook("layer3")))
    hooks.append(model.layer4.register_forward_hook(make_hook("layer4")))
    hooks.append(model.global_pool.register_forward_hook(make_hook("avgpool")))

    return hooks, features


def extract_features(model, loader, device, max_samples=1000):
    """
    Run forward pass and collect features at all 6 layers.
    """
    model.eval()
    hooks, feat_dict = register_hooks(model)

    all_hospital = []
    all_tumor    = []
    count        = 0

    with torch.no_grad():
        for batch in loader:
            imgs     = batch[0].to(device)
            labels   = batch[1].squeeze().long()
            metadata = batch[2]
            model(imgs)
            all_hospital.append(metadata[:, 0].long())
            all_tumor.append(labels)
            count += imgs.size(0)
            if count >= max_samples:
                break

    for h in hooks:
        h.remove()

    features     = {k: torch.cat(v).numpy() for k, v in feat_dict.items()}
    hospital_ids = torch.cat(all_hospital).numpy()
    tumor_labels = torch.cat(all_tumor).numpy()

    n = min(max_samples, len(hospital_ids))
    features     = {k: v[:n] for k, v in features.items()}
    hospital_ids = hospital_ids[:n]
    tumor_labels = tumor_labels[:n]

    return features, hospital_ids, tumor_labels


def train_probe(X_train, y_train, X_val, y_val):
    """
    Train logistic regression probe on frozen features.
    """
    if len(np.unique(y_train)) < 2:
        return 0.5

    scaler  = StandardScaler()
    X_train = scaler.fit_transform(X_train)
    X_val   = scaler.transform(X_val)

    clf = LogisticRegression(
        max_iter=500,
        C=1.0,
        solver="lbfgs",
        multi_class="auto",
        n_jobs=-1,
    )
    try:
        clf.fit(X_train, y_train)
        return clf.score(X_val, y_val)
    except Exception:
        return float("nan")


# ──────────────────────────────────────────────
# CHECKPOINT DISCOVERY
# ──────────────────────────────────────────────

def find_checkpoints(run_dir: str) -> List[tuple]:
    """
    Find all checkpoints in a run directory.
    Returns list of (epoch, checkpoint_path) sorted by epoch.
    """
    ckpt_dir = os.path.join(run_dir, "checkpoints")
    if not os.path.isdir(ckpt_dir):
        return []

    checkpoints = []

    # Periodic checkpoints: ep050.pt, ep100.pt, etc.
    for f in sorted(glob.glob(os.path.join(ckpt_dir, "ep*.pt"))):
        epoch_str = os.path.basename(f).replace("ep", "").replace(".pt", "")
        try:
            epoch = int(epoch_str)
            checkpoints.append((epoch, f))
        except ValueError:
            continue

    # Final checkpoint
    final = os.path.join(ckpt_dir, "final.pt")
    if os.path.isfile(final):
        hist_path = os.path.join(run_dir, "results", "history.json")
        if os.path.isfile(hist_path):
            try:
                hist  = json.load(open(hist_path))
                epoch = hist[-1]["epoch"] if hist else 9999
            except Exception:
                epoch = 9999
        else:
            epoch = 9999
        checkpoints.append((epoch, final))

    return sorted(checkpoints, key=lambda x: x[0])


def load_model_from_checkpoint(ckpt_path: str, n_classes: int = 2,
                                device: str = "cuda") -> nn.Module:
    model = timm.create_model("resnet18", pretrained=False,
                               num_classes=n_classes)
    state = torch.load(ckpt_path, map_location=device)
    model.load_state_dict(state, strict=True)
    model.eval()
    return model.to(device)


# ──────────────────────────────────────────────
# MAIN PROBE ANALYSIS
# ──────────────────────────────────────────────

def run_probe_analysis(run_dir: str, data_root: str,
                       device: str = "cuda",
                       max_samples: int = 800,
                       latest_only: bool = False) -> Optional[Dict]:
    """
    For each checkpoint in a run, extract features at all 6 layers
    and train hospital + tumor probes.
    """
    from utils.camelyon_data import get_camelyon_subsets

    cfg_path = os.path.join(run_dir, "config.json")
    if not os.path.isfile(cfg_path):
        print(f"  No config.json in {run_dir}, skipping")
        return None

    cfg        = json.load(open(cfg_path))
    condition  = cfg.get("condition", "unknown")
    n_train    = cfg.get("n_train", 300)
    seed       = cfg.get("seed", 42)

    print(f"\n{'='*55}")
    print(f"  M1 Probe Analysis: {os.path.basename(run_dir)}")
    print(f"  condition={condition}, n_train={n_train}, seed={seed}")
    print(f"{'='*55}")

    checkpoints = find_checkpoints(run_dir)
    if not checkpoints:
        print(f"  No checkpoints found — skipping")
        return None

    if latest_only:
        checkpoints = checkpoints[-1:]

    print(f"  Found {len(checkpoints)} checkpoints: "
          f"epochs {[e for e,_ in checkpoints]}")
    print(f"  Hospital probe: fits on training data (H0-H2), "
          f"evaluates on H3 and H4 separately")

    # ── Data ─────────────────────────────────────────────────────────
    transform = transforms.Compose([
        transforms.Resize((96, 96)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225]),
    ])
    train_ds, id_val_ds, ood_test_ds, full_ds = get_camelyon_subsets(
        root_dir=data_root, download=False)

    # Wrap datasets with transform (WILDS returns PIL images)
    class _TransformWrapper:
        def __init__(self, dataset, transform):
            self.dataset = dataset
            self.transform = transform
        def __len__(self):
            return len(self.dataset)
        def __getitem__(self, idx):
            img, label, metadata = self.dataset[idx]
            return self.transform(img), label, metadata

    id_val_t = _TransformWrapper(id_val_ds, transform)
    ood_test_t = _TransformWrapper(ood_test_ds, transform)
    train_t = _TransformWrapper(train_ds, transform)

    torch.manual_seed(seed)
    probe_idx = torch.randperm(len(id_val_t))[:max_samples // 2]
    ood_idx   = torch.randperm(len(ood_test_t))[:max_samples // 2]
    train_idx = torch.randperm(len(train_t))[:max_samples]

    probe_loader = DataLoader(
        Subset(id_val_t, probe_idx),
        batch_size=128, shuffle=False, num_workers=0)
    ood_loader   = DataLoader(
        Subset(ood_test_t, ood_idx),
        batch_size=128, shuffle=False, num_workers=0)
    train_loader = DataLoader(
        Subset(train_t, train_idx),
        batch_size=128, shuffle=False, num_workers=0)

    # ── Results storage ──────────────────────────────────────────────
    results = {
        "run_id":     os.path.basename(run_dir),
        "condition":  condition,
        "n_train":    n_train,
        "seed":       seed,
        "epochs":     [],
        "layers":     LAYER_NAMES,
        "hospital_probe_id":  [],  # Hospital accuracy on H3
        "hospital_probe_ood": [],  # Hospital accuracy on H4
        "tumor_probe_id":     [],  # Tumor accuracy on H3
        "tumor_probe_ood":    [],  # Tumor accuracy on H4
    }

    # ── Per-checkpoint analysis ───────────────────────────────────────
    for epoch, ckpt_path in checkpoints:
        print(f"\n  Epoch {epoch} | {os.path.basename(ckpt_path)}")

        try:
            model = load_model_from_checkpoint(
                ckpt_path, n_classes=2, device=device)
        except Exception as e:
            print(f"    Failed to load checkpoint: {e}")
            continue

        # Extract features from all three datasets
        feats_train, hosp_train, tumor_train = extract_features(
            model, train_loader, device, max_samples=max_samples)

        feats_id, hosp_id, tumor_id = extract_features(
            model, probe_loader, device, max_samples=max_samples // 2)

        feats_ood, hosp_ood, tumor_ood = extract_features(
            model, ood_loader, device, max_samples=max_samples // 2)

        epoch_hosp_id = []
        epoch_hosp_ood = []
        epoch_tumor_id = []
        epoch_tumor_ood = []

        for layer_name in LAYER_NAMES:
            # Fit probes on training features, evaluate on H3 and H4
            X_train_layer = feats_train[layer_name]
            X_id_layer = feats_id[layer_name]
            X_ood_layer = feats_ood[layer_name]

            # Hospital probe: can model distinguish hospitals H0-H2?
            # If yes on H3/H4 → stain is encoded
            h_acc_id = train_probe(X_train_layer, hosp_train,
                                   X_id_layer,  hosp_id)
            h_acc_ood = train_probe(X_train_layer, hosp_train,
                                    X_ood_layer, hosp_ood)

            # Tumor probe: can model distinguish tumor vs normal?
            t_acc_id = train_probe(X_train_layer, tumor_train,
                                   X_id_layer,  tumor_id)
            t_acc_ood = train_probe(X_train_layer, tumor_train,
                                    X_ood_layer, tumor_ood)

            epoch_hosp_id.append(h_acc_id)
            epoch_hosp_ood.append(h_acc_ood)
            epoch_tumor_id.append(t_acc_id)
            epoch_tumor_ood.append(t_acc_ood)

            print(f"    {layer_name:8s}: "
                  f"hosp_H3={h_acc_id:.3f} hosp_H4={h_acc_ood:.3f}  "
                  f"tumor_H3={t_acc_id:.3f} tumor_H4={t_acc_ood:.3f}")

        results["epochs"].append(epoch)
        results["hospital_probe_id"].append(epoch_hosp_id)
        results["hospital_probe_ood"].append(epoch_hosp_ood)
        results["tumor_probe_id"].append(epoch_tumor_id)
        results["tumor_probe_ood"].append(epoch_tumor_ood)

        del model

    # ── Save raw data ─────────────────────────────────────────────────
    out_dir = os.path.join(run_dir, "mechinterp")
    os.makedirs(out_dir, exist_ok=True)

    data_path = os.path.join(out_dir, "m1_probe_data.json")
    with open(data_path, "w") as f:
        json.dump(results, f, indent=2)
    print(f"\n  Probe data → {data_path}")

    # ── Plots ─────────────────────────────────────────────────────────
    _plot_probe_heatmaps(results, out_dir)
    _plot_probe_curves(results, out_dir)

    print(f"  Figures → {out_dir}/")
    return results


def _plot_probe_heatmaps(results: Dict, out_dir: str):
    """
    Epoch (x) × layer (y), color = probe accuracy.

    Hospital probe shown on H3 (held-in held-out hospital, classes overlap
    with training). The H4 version is degenerate by construction since the
    probe is fit on the training-hospital class set and H4 is not in it
    (hospital_probe_ood ≡ 0 across all epochs / layers).

    Tumor probe shown on H4 (truly OOD hospital) since tumor labels are
    binary and shared across hospitals — H4 captures the causal-feature
    transferability we care about.
    """
    epochs = results["epochs"]
    layers = results["layers"]

    if not epochs:
        return

    hosp_matrix  = np.array(results["hospital_probe_id"])    # H3 — has signal
    tumor_matrix = np.array(results["tumor_probe_ood"])      # H4 — true OOD

    fig, axes = plt.subplots(1, 2, figsize=(14, 5))

    for ax, matrix, title, cmap in [
        (axes[0], hosp_matrix,  "Hospital probe on H3 (shortcut recoverability)\nHigh = stain still encoded = BAD", "Reds"),
        (axes[1], tumor_matrix, "Tumor probe on H4 (causal, OOD)\nHigh = causal feature transfers = GOOD",          "Greens"),
    ]:
        im = ax.imshow(
            matrix.T,
            aspect="auto",
            cmap=cmap,
            vmin=0.0, vmax=1.0,
            interpolation="nearest",
            origin="lower",
        )
        ax.set_xticks(range(len(epochs)))
        ax.set_xticklabels(epochs, rotation=45, ha="right", fontsize=8)
        ax.set_yticks(range(len(layers)))
        ax.set_yticklabels(layers, fontsize=9)
        ax.set_xlabel("Training epoch")
        ax.set_ylabel("ResNet layer")
        ax.set_title(title, fontsize=10, fontweight="bold")
        plt.colorbar(im, ax=ax, label="Probe accuracy")

    fig.suptitle(
        f"M1 — Layer-wise Linear Probing: {results['run_id']}\n"
        "Circuit signature: deep-layer hospital-probe drop (Reds) + sustained tumor recoverability (Greens)",
        fontsize=10, y=1.02
    )
    plt.tight_layout()
    out = os.path.join(out_dir, "m1_probe_heatmap.png")
    plt.savefig(out, bbox_inches="tight")
    plt.close()


def _plot_probe_curves(results: Dict, out_dir: str):
    """
    Line plot per-layer: hospital probe (H3 — recoverability) + tumor probe
    (H4 — causal-feature transfer to truly unseen hospital), with OOD accuracy
    from history.json overlaid.
    """
    epochs     = results["epochs"]
    run_dir    = os.path.join(out_dir, "..")
    layers     = results["layers"]
    avgpool_idx = layers.index("avgpool")
    layer2_idx  = layers.index("layer2")

    fig, axes = plt.subplots(1, 2, figsize=(14, 5))

    for ax, layer_idx, layer_label in [
        (axes[0], avgpool_idx, "avgpool (penultimate)"),
        (axes[1], layer2_idx,  "layer2 (early)"),
    ]:
        hosp  = [results["hospital_probe_id"][i][layer_idx]
                 for i in range(len(epochs))]
        tumor = [results["tumor_probe_ood"][i][layer_idx]
                 for i in range(len(epochs))]

        ax.plot(epochs, hosp,  "r-o", markersize=4, lw=2,
                label="Hospital probe on H3 (shortcut recoverability ↓ want)")
        ax.plot(epochs, tumor, "g-s", markersize=4, lw=2,
                label="Tumor probe on H4 (causal transfer ↑ want)")

        hist_path = os.path.join(run_dir, "results", "history.json")
        if os.path.isfile(hist_path):
            try:
                hist     = json.load(open(hist_path))
                hist_eps = [r["epoch"]   for r in hist]
                ood_accs = [r.get("ood_acc", float("nan")) for r in hist]
                ax.plot(hist_eps, ood_accs, "b--", lw=1.5, alpha=0.7,
                        label="OOD accuracy (H4)")
            except Exception:
                pass

        ax.axhline(0.5, color="gray", ls=":", lw=1, alpha=0.5,
                   label="Chance (0.5)")
        ax.set_xlabel("Training epoch")
        ax.set_ylabel("Probe / OOD accuracy")
        ax.set_title(f"Layer: {layer_label}", fontweight="bold")
        ax.legend(fontsize=9)
        ax.set_ylim([0, 1.05])
        ax.grid(alpha=0.3)

    fig.suptitle(
        f"M1 — Probe Curves: {results['run_id']}\n"
        "Hospital recoverability (H3) drops in deep layers + tumor transfers (H4) — circuit signature",
        fontsize=10, y=1.02
    )
    plt.tight_layout()
    out = os.path.join(out_dir, "m1_probe_curves.png")
    plt.savefig(out, bbox_inches="tight")
    plt.close()


def main():
    p = argparse.ArgumentParser(
        description="M1: Layer-wise linear probing for CausalGrok")
    p.add_argument("--run_dir",     default=None,
                   help="Single run directory to analyze")
    p.add_argument("--all_runs",    action="store_true",
                   help="Analyze all camelyon_v2 grokking runs")
    p.add_argument("--data_root",   default="data/wilds")
    p.add_argument("--device",      default="cuda")
    p.add_argument("--max_samples", type=int, default=800)
    p.add_argument("--latest_only", action="store_true",
                   help="Analyze only latest checkpoint (quick check)")
    args = p.parse_args()

    if args.all_runs:
        run_dirs = sorted(glob.glob(
            "experiments/runs/*camelyon_v2*grokking*"))
        print(f"Found {len(run_dirs)} grokking runs")
        all_results = []
        for rd in run_dirs:
            r = run_probe_analysis(rd, args.data_root,
                                   device=args.device,
                                   max_samples=args.max_samples,
                                   latest_only=args.latest_only)
            if r:
                all_results.append(r)

        if all_results:
            os.makedirs("paper_figures", exist_ok=True)
            with open("paper_figures/m1_all_probes.json", "w") as f:
                json.dump(all_results, f, indent=2)
            print(f"\nCombined → paper_figures/m1_all_probes.json")

    elif args.run_dir:
        run_probe_analysis(args.run_dir, args.data_root,
                           device=args.device,
                           max_samples=args.max_samples,
                           latest_only=args.latest_only)
    else:
        print("Specify --run_dir <path> or --all_runs")


if __name__ == "__main__":
    main()