"""
Warm-starting your own model from a released checkpoint.

This is the piece most people actually want: the pretrained checkpoints take
8 input channels, and almost certainly your data does not have the beam-geometry
and plan-prompt channels the source domain had. `load_pretrained_into` copies
every tensor whose shape matches and leaves the rest at their fresh
initialisation, which for these three architectures means only the input stem is
re-initialised:

    C3D         170 / 172 tensors transferred (net_A + net_B first conv skipped)
    MedNeXt     228 / 229 tensors transferred (stem.weight skipped)
    SwinUNETR   164 / 167 tensors transferred (patch_embed.proj, encoder1 conv1/conv3)

This is exactly how the released *_finetuned checkpoints were produced. No
channel slicing, averaging or repetition was applied to the input stem -- it was
simply retrained from scratch.

Example
-------
    from build_model import build_model
    from transfer import load_pretrained_into

    model = build_model("c3d", in_channels=5)          # your channel count
    report = load_pretrained_into(model, "weights/c3d_pretrained.pt")
    print(report)
    # ... then train as usual; a low LR (3e-5, cosine) worked best for us.
"""

from dataclasses import dataclass, field
from typing import List

import torch


@dataclass
class TransferReport:
    transferred: List[str] = field(default_factory=list)
    shape_mismatch: List[str] = field(default_factory=list)
    missing_in_checkpoint: List[str] = field(default_factory=list)
    unused_in_checkpoint: List[str] = field(default_factory=list)

    def __str__(self):
        n_ok = len(self.transferred)
        n_total = n_ok + len(self.shape_mismatch) + len(self.missing_in_checkpoint)
        lines = [f"transferred {n_ok}/{n_total} tensors"]
        for label, items in (
            ("re-initialised (shape mismatch)", self.shape_mismatch),
            ("re-initialised (absent from checkpoint)", self.missing_in_checkpoint),
            ("ignored (not in your model)", self.unused_in_checkpoint),
        ):
            if items:
                lines.append(f"  {label}: {len(items)}")
                lines.extend(f"    - {k}" for k in items)
        return "\n".join(lines)


def load_pretrained_into(model, ckpt_path, verbose=True):
    """
    Copy every shape-compatible tensor from a released checkpoint into `model`.

    Tensors that do not match (different input-channel count, different output
    head, ...) keep whatever `model` initialised them to. Returns a
    TransferReport naming exactly what was and was not transferred -- read it,
    do not assume.
    """
    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=True)
    src = ckpt["state_dict"] if "state_dict" in ckpt else ckpt
    src = {k[len("module."):] if k.startswith("module.") else k: v
           for k, v in src.items()}

    own = model.state_dict()
    report = TransferReport()
    staged = {}

    for name, tensor in own.items():
        if name not in src:
            report.missing_in_checkpoint.append(name)
        elif src[name].shape != tensor.shape:
            report.shape_mismatch.append(
                f"{name}: checkpoint {tuple(src[name].shape)} vs model {tuple(tensor.shape)}")
        else:
            staged[name] = src[name]
            report.transferred.append(name)

    report.unused_in_checkpoint = [k for k in src if k not in own]

    own.update(staged)
    model.load_state_dict(own)

    if verbose:
        print(f"[transfer] {ckpt_path} -> {type(model).__name__}")
        print(report)
    return report