"""
load_akn.py — self-contained loader for the Antaḥkaraṇa-Net scaled checkpoint.

No repo needed: this file carries the full model definition, so you can load the
released `.pt` and run inference / inspect the saved consolidation state (Ω, θ*)
with just PyTorch.

    from load_akn import load
    model, ckpt = load("antahkarana_cifar100_wrn28-10.pt")
    # model(x, task) -> logits for that task's head ; task in [0..n_tasks-1]

The checkpoint dict contains:
    model_state  — the trained WideResNet trunk + per-task heads
    config       — {depth, width, n_tasks, classes_per_task}
    results      — the honest metrics (forgetting, accuracy, pramāṇa gating, per-task)
    omega / theta_star (optional) — saṃskāra importance Ω and reference weights θ*
"""
from __future__ import annotations
import torch, torch.nn as nn, torch.nn.functional as F


# --------------------------------------------------------------- WideResNet (CIFAR)
class _BasicBlock(nn.Module):
    def __init__(self, in_p, out_p, stride, drop=0.0):
        super().__init__()
        self.bn1 = nn.BatchNorm2d(in_p); self.conv1 = nn.Conv2d(in_p, out_p, 3, stride, 1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_p); self.conv2 = nn.Conv2d(out_p, out_p, 3, 1, 1, bias=False)
        self.drop = drop
        self.equalInOut = (in_p == out_p and stride == 1)
        self.shortcut = None if self.equalInOut else nn.Conv2d(in_p, out_p, 1, stride, 0, bias=False)

    def forward(self, x):
        if not self.equalInOut:
            x = F.relu(self.bn1(x)); out = F.relu(self.bn2(self.conv1(x)))
        else:
            out = F.relu(self.bn1(x)); out = F.relu(self.bn2(self.conv1(out)))
        if self.drop > 0:
            out = F.dropout(out, self.drop, self.training)
        out = self.conv2(out)
        return out + (x if self.equalInOut else self.shortcut(x))


class _WideResNet(nn.Module):
    def __init__(self, depth=28, widen=10, drop=0.0):
        super().__init__()
        assert (depth - 4) % 6 == 0
        n = (depth - 4) // 6
        ch = [16, 16 * widen, 32 * widen, 64 * widen]
        self.conv1 = nn.Conv2d(3, ch[0], 3, 1, 1, bias=False)
        self.block1 = self._make(ch[0], ch[1], n, 1, drop)
        self.block2 = self._make(ch[1], ch[2], n, 2, drop)
        self.block3 = self._make(ch[2], ch[3], n, 2, drop)
        self.bn1 = nn.BatchNorm2d(ch[3]); self.nChannels = ch[3]

    def _make(self, in_p, out_p, n, stride, drop):
        return nn.Sequential(*[_BasicBlock(in_p if i == 0 else out_p, out_p,
                                           stride if i == 0 else 1, drop) for i in range(n)])

    def features(self, x):
        out = self.conv1(x); out = self.block1(out); out = self.block2(out); out = self.block3(out)
        out = F.relu(self.bn1(out)); out = F.avg_pool2d(out, 8)
        return out.view(-1, self.nChannels)


class AntahkaranaWRN(nn.Module):
    """Shared WRN trunk (φ) + one linear head per task (buddhi)."""
    def __init__(self, depth, widen, n_tasks, classes_per_task):
        super().__init__()
        self.backbone = _WideResNet(depth, widen)
        self.heads = nn.ModuleList([nn.Linear(self.backbone.nChannels, classes_per_task)
                                    for _ in range(n_tasks)])

    def features(self, x):
        return self.backbone.features(x)

    def forward(self, x, task: int):
        return self.heads[task](self.features(x))


def load(path: str, map_location="cpu"):
    """Return (model_in_eval_mode, full_checkpoint_dict)."""
    ck = torch.load(path, map_location=map_location, weights_only=False)
    c = ck["config"]
    model = AntahkaranaWRN(c["depth"], c["width"], c["n_tasks"], c["classes_per_task"])
    model.load_state_dict(ck["model_state"])
    model.eval()
    return model, ck


if __name__ == "__main__":
    import sys
    p = sys.argv[1] if len(sys.argv) > 1 else "antahkarana_cifar100_wrn28-10.pt"
    model, ck = load(p)
    c, r = ck["config"], ck["results"]
    n_params = sum(t.numel() for t in model.parameters())
    print(f"loaded {p}")
    print(f"  WRN-{c['depth']}-{c['width']}  |  {n_params/1e6:.1f}M params  |  "
          f"{c['n_tasks']} tasks × {c['classes_per_task']} classes")
    ag, na = r["agent"], r["naive"]
    print(f"  forgetting : naive {na['forgetting']:.3f}  ->  agent {ag['forgetting']:.4f}"
          + (f"  ({r['forgetting_reduction']:.1f}x lower)" if r.get("forgetting_reduction") else ""))
    print(f"  accuracy   : naive {na['avg_acc']:.3f}  ->  agent {ag['avg_acc']:.3f}")
    print(f"  pramāṇa    : gated acc {ag['gated_accuracy']:.3f} @ coverage {ag['gated_coverage']:.2f}")
    print(f"  saṃskāra Ω/θ* present: {('omega' in ck and 'theta_star' in ck)}")
    # tiny forward smoke
    x = torch.randn(2, 3, 32, 32)
    print(f"  forward(x, task=0) -> logits {tuple(model(x, 0).shape)}")