"""Grow the pretrained 7.8M baseline. Verified path: DEPTH (identity-init new blocks).

Per skill tiny-model-phase2 (measured Aug 2026 on this tablet):
  - width upscaling 320->512 does NOT transfer (val loss 2.58 -> 6.1-7.7) because
    RMSNorm/rope/groupnorm/recurrence all depend on d_model.
  - depth growth with identity blocks DOES preserve baseline exactly (2.567 vs 2.578).

Usage:
  PYTHONPATH=$PWD .venv/bin/python train/grow_weights.py \
      --base ckpt/nlp_full --config tiny13m --ckpt ckpt/tiny13m_grown --verify
"""
import argparse
import copy
import random

import torch
import torch.nn.functional as F

from model.config import TinyLiquidConfig, CONFIGS
from model.tiny_liquid import TinyLiquid
from model.utils import latest_ckpt
from data.tokenizer import load_tokenizer

STD = 0.02
NEW_BLOCK_SCALE = 0.1


def pad_noise(old, rows_new=None, cols_new=None):
    old = old.float()
    r0, c0 = old.shape
    rn = rows_new if rows_new is not None else r0
    cn = cols_new if cols_new is not None else c0
    out = torch.empty(rn, cn)
    out.fill_(0.0)
    out[:r0, :c0] = old
    mask = torch.ones_like(out, dtype=torch.bool)
    mask[:r0, :c0] = False
    out[mask] = torch.normal(0.0, STD, size=(int(mask.sum()),))
    return out


def grow_block_width(block, d_new, basis_rows_new, h_new):
    nb = copy.deepcopy(block)
    nb["norm1.weight"] = torch.cat([block["norm1.weight"].float(),
                                    torch.ones(d_new - block["norm1.weight"].shape[0])])
    nb["basis.w"] = pad_noise(block["basis.w"], rows_new=basis_rows_new, cols_new=d_new)
    nb["basis.w_forget"] = pad_noise(block["basis.w_forget"], rows_new=basis_rows_new, cols_new=d_new)
    nb["basis.gn.weight"] = torch.cat([block["basis.gn.weight"].float(),
                                       torch.ones(basis_rows_new - block["basis.gn.weight"].shape[0])])
    nb["basis.gn.bias"] = torch.cat([block["basis.gn.bias"].float(),
                                     torch.zeros(basis_rows_new - block["basis.gn.bias"].shape[0])])
    nb["norm2.weight"] = torch.cat([block["norm2.weight"].float(),
                                    torch.ones(d_new - block["norm2.weight"].shape[0])])
    nb["mlp.up.weight"] = pad_noise(block["mlp.up.weight"], rows_new=h_new, cols_new=d_new)
    nb["mlp.gate.weight"] = pad_noise(block["mlp.gate.weight"], rows_new=h_new, cols_new=d_new)
    nb["mlp.forget.weight"] = pad_noise(block["mlp.forget.weight"], rows_new=h_new, cols_new=d_new)
    nb["mlp.down.weight"] = pad_noise(block["mlp.down.weight"], rows_new=d_new, cols_new=h_new)
    return nb


def identity_block(cfg, prefix):
    """A block that is an exact identity at init (output == input)."""
    d = cfg.d_model
    e = cfg.basis_n * cfg.basis_b
    h = cfg.mlp_ratio * d
    sd = {}
    sd[f"{prefix}.norm1.weight"] = torch.ones(d)
    sd[f"{prefix}.basis.w"] = torch.zeros(e, d)
    sd[f"{prefix}.basis.w_forget"] = torch.zeros(e, d)
    sd[f"{prefix}.basis.gn.weight"] = torch.ones(e)
    sd[f"{prefix}.basis.gn.bias"] = torch.zeros(e)
    sd[f"{prefix}.norm2.weight"] = torch.ones(d)
    sd[f"{prefix}.mlp.up.weight"] = torch.zeros(h, d)
    sd[f"{prefix}.mlp.gate.weight"] = torch.zeros(h, d)
    sd[f"{prefix}.mlp.forget.weight"] = torch.zeros(h, d)
    sd[f"{prefix}.mlp.down.weight"] = torch.zeros(d, h)
    return sd


def grow_depth(sd_old, cfg_old, cfg_new):
    """Copy trunk exactly; append identity blocks. Baseline loss preserved."""
    assert cfg_old.d_model == cfg_new.d_model
    assert cfg_old.basis_n == cfg_new.basis_n and cfg_old.basis_b == cfg_new.basis_b
    assert cfg_old.mlp_ratio == cfg_new.mlp_ratio
    assert cfg_new.n_blocks >= cfg_old.n_blocks
    grown = {}
    for k, v in sd_old["model"].items():
        grown[k] = v.clone()
    for j in range(cfg_old.n_blocks, cfg_new.n_blocks):
        grown.update(identity_block(cfg_new, f"blocks.{j}"))
    return grown


def grow_width(sd_old, cfg_old, cfg_new):
    """Width upscaling -- EXPERIMENTAL, does NOT transfer on this architecture."""
    old = sd_old["model"]
    d_old, d_new = cfg_old.d_model, cfg_new.d_model
    basis_old = cfg_old.basis_n * cfg_old.basis_b
    basis_new = cfg_new.basis_n * cfg_new.basis_b
    h_old, h_new = cfg_old.mlp_ratio * d_old, cfg_new.mlp_ratio * d_new
    grown = {}
    grown["tok_emb.weight"] = pad_noise(old["tok_emb.weight"], cols_new=d_new)
    grown["persona_emb.weight"] = pad_noise(old["persona_emb.weight"], cols_new=d_new)
    grown["norm_out.weight"] = torch.cat([old["norm_out.weight"].float(),
                                          torch.ones(d_new - d_old)])
    for i in range(cfg_old.n_blocks):
        block = {k[len(f"blocks.{i}."):]: v for k, v in old.items() if k.startswith(f"blocks.{i}.")}
        gb = grow_block_width(block, d_new, basis_new, h_new)
        for k, v in gb.items():
            grown[f"blocks.{i}.{k}"] = v
    last = {k[len(f"blocks.{cfg_old.n_blocks-1}."):]: v for k, v in old.items()
            if k.startswith(f"blocks.{cfg_old.n_blocks-1}.")}
    seed = grow_block_width(last, d_new, basis_new, h_new)
    for j in range(cfg_old.n_blocks, cfg_new.n_blocks):
        for k, v in seed.items():
            grown[f"blocks.{j}.{k}"] = (v * NEW_BLOCK_SCALE).clone()
    return grown


@torch.no_grad()
def val_loss(model, tok, data_path, batch=16, seq=256, n_batches=20):
    import numpy as np
    tokens = torch.from_numpy(np.fromfile(data_path, dtype="uint16")).long()
    rng = random.Random(42)
    total = 0.0
    for _ in range(n_batches):
        pos = rng.randrange(0, max(1, len(tokens) - seq - 1))
        x = tokens[pos:pos + batch * seq].view(batch, seq)
        y = tokens[pos + 1:pos + 1 + batch * seq].view(batch, seq)
        logits = model(x)
        total += F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1)).item()
    return total / n_batches


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", default="ckpt/nlp_full")
    ap.add_argument("--config", default="tiny13m")
    ap.add_argument("--mode", choices=["depth", "width", "tower"], default="depth")
    ap.add_argument("--ckpt", default="ckpt/tiny13m_grown")
    ap.add_argument("--verify", action="store_true")
    ap.add_argument("--threads", type=int, default=4)
    args = ap.parse_args()

    torch.set_num_threads(args.threads)
    tok = load_tokenizer("data/tokenizer.json")
    cfg_new = TinyLiquidConfig(vocab_size=tok.get_vocab_size(), **CONFIGS[args.config])

    base_path = latest_ckpt(args.base)
    assert base_path, f"no checkpoint in {args.base}"
    sd = torch.load(base_path, map_location="cpu")
    cfg_old = TinyLiquidConfig(vocab_size=tok.get_vocab_size(),
                               **{k: v for k, v in sd["config"].items() if k != "vocab_size"})
    print(f"base: {base_path}  {cfg_old.params_estimate()/1e6:.2f}M -> {cfg_new.params_estimate()/1e6:.2f}M ({args.mode})")

    if args.mode == "depth":
        grown = grow_depth(sd, cfg_old, cfg_new)
    elif args.mode == "width":
        grown = grow_width(sd, cfg_old, cfg_new)
    else:
        # tower: keep the trained trunk, add identity-init wide tower.
        # If tower_d changed, handle by loading trunk + padding tower projection.
        assert cfg_new.tower_d and cfg_new.tower_blocks, "tower mode needs tower_d/tower_blocks in config"
        model_old = TinyLiquid(cfg_old)
        model_old.load_state_dict(sd["model"])
        model_old.eval()
        model = TinyLiquid(cfg_new)
        old_sd = sd["model"]
        new_sd = dict(model.state_dict())
        for k, v in old_sd.items():
            if not k.startswith("up_proj") and not k.startswith("down_proj") and not k.startswith("tower."):
                if k in new_sd:
                    new_sd[k] = v.clone()
                else:
                    print(f"WARNING: {k} in old but not in new model")
        if "up_proj" in old_sd:
            old_up = old_sd["up_proj"]; old_down = old_sd["down_proj"]
            new_up = new_sd["up_proj"]; new_down = new_sd["down_proj"]
            if old_up.shape == new_up.shape:
                new_sd["up_proj"] = old_up.clone()
                new_sd["down_proj"] = old_down.clone()
            else:
                new_sd["up_proj"] = new_sd["up_proj"].clone()
                new_sd["down_proj"] = new_sd["down_proj"].clone()
                min_td = min(old_up.shape[0], new_up.shape[0])
                min_d = min(old_up.shape[1], new_up.shape[1])
                new_sd["up_proj"][:min_td, :min_d] = old_up[:min_td, :min_d]
                min_d2 = min(old_down.shape[0], new_down.shape[0])
                min_td2 = min(old_down.shape[1], new_down.shape[1])
                new_sd["down_proj"][:min_d2, :min_td2] = old_down[:min_d2, :min_td2]
        grown = new_sd

    model = TinyLiquid(cfg_new)
    missing, unexpected = model.load_state_dict(grown, strict=False)
    assert not missing and not unexpected, (missing, unexpected)

    if args.verify:
        base_model = TinyLiquid(cfg_old)
        base_model.load_state_dict(sd["model"])
        base_model.eval(); model.eval()
        print(f"baseline val loss (20 bat): {val_loss(base_model, tok, 'data/valid.bin', n_batches=20):.4f}")
        print(f"grown    val loss (20 bat): {val_loss(model, tok, 'data/valid.bin', n_batches=20):.4f}")

    out_dir = __import__("pathlib").Path(args.ckpt)
    out_dir.mkdir(parents=True, exist_ok=True)
    out = out_dir / ("model_final.pt" if args.mode == "tower" else "model_grown.pt")
    torch.save({"config": cfg_new.__dict__, "model": model.state_dict(),
                "step": 0, "best_val": float("inf")}, out)
    print("saved", out)


if __name__ == "__main__":
    main()