#!/usr/bin/env python3
"""
build_hierarchy.py – Unified hierarchy construction pipeline.

Subcommands:
  learn     Cosine k-means++ clustering for SAD hierarchy prototypes.
  build_lut Build sparse soft top-k ancestor lookup tables.
  extend    Extend an existing hierarchy with deeper ancestor levels.
  merge     Merge per-level LUT data into a mixed-topk LUT.

Usage:
  python scripts/build_hierarchy.py learn --config configs/sad_owt.yaml
  python scripts/build_hierarchy.py build_lut --config configs/sad_owt.yaml
  python scripts/build_hierarchy.py extend --input data/hierarchy_50257-128.pt \
      --levels 128,32,8 --output data/hierarchy_50257-128-32-8.pt
  python scripts/build_hierarchy.py merge --source 1:data/lut_top3.pt \
      --source 2:data/lut_top2.pt --output data/lut_mixed.pt
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]  # sad/

import torch
import torch.nn.functional as F
import yaml
from tqdm import tqdm

sys.path.insert(0, str(ROOT))

from src.models.sad_model import SADModel


# ── Shared utilities ──────────────────────────────────────────────────────────


def load_gpt2_embeddings(gpt2_path: str) -> torch.Tensor:
    from transformers import AutoModel
    m = AutoModel.from_pretrained(gpt2_path, local_files_only=True)
    wte = m.wte.weight.detach().float()
    del m
    return wte


def load_config(path: str) -> dict:
    with open(path) as f:
        return yaml.safe_load(f)


def build_tokenizer(config: dict):
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained(
        ROOT / "tokenizers" / "gpt2",
        local_files_only=True,
    )
    if tok.pad_token_id is None:
        tok.pad_token = tok.eos_token
    if tok.mask_token_id is None:
        tok.add_special_tokens({"mask_token": "[MASK]"})
    return tok


def get_leaf_embeddings(config: dict, gpt2_path: str | None, device: torch.device):
    model_cfg = config["model"]
    hier_cfg = config.get("hierarchy", {})
    gpt2_path = Path(
        gpt2_path
        or hier_cfg.get("gpt2_path", str(ROOT / "models" / "gpt2"))
    )
    if gpt2_path.exists() and (
        (gpt2_path / "pytorch_model.bin").exists()
        or any(gpt2_path.glob("model.safetensors"))
    ):
        all_emb = load_gpt2_embeddings(str(gpt2_path)).to(device)
        V_tok = model_cfg["vocab_size"]
        if all_emb.shape[0] < V_tok:
            pad = torch.zeros(V_tok - all_emb.shape[0], all_emb.shape[1], device=device)
            all_emb = torch.cat([all_emb, pad], dim=0)
        elif all_emb.shape[0] > V_tok:
            all_emb = all_emb[:V_tok]
        print(f"Leaf embeddings: GPT-2 wte  {tuple(all_emb.shape)}")
    else:
        model = SADModel(
            vocab_size=model_cfg["vocab_size"],
            hidden_size=model_cfg["hidden_size"],
            n_blocks=model_cfg["n_blocks"],
            n_heads=model_cfg["n_heads"],
            cond_dim=model_cfg["cond_dim"],
            max_seq_len=model_cfg["max_seq_len"],
            block_size=model_cfg.get("block_size", 8),
            num_levels=model_cfg.get("num_levels", 2),
            level_sizes=model_cfg["level_sizes"],
        ).to(device)
        with torch.no_grad():
            all_emb = model.get_leaf_embeddings().float().to(device)
        print(f"Leaf embeddings: random init  {tuple(all_emb.shape)}")
    return all_emb


# ── learn: k-means++ clustering ───────────────────────────────────────────────


def semantic_kmeans_cluster(
    embeddings: torch.Tensor,
    cluster_size: int,
    max_iters: int = 100,
    tolerance: float = 1e-4,
    batch_size: int = 2048,
    min_size_ratio: float = 0.1,
    max_size_ratio: float = 5.0,
    seed: int = 42,
) -> torch.Tensor:
    torch.manual_seed(seed)
    device = embeddings.device
    N, d = embeddings.shape

    emb = F.normalize(embeddings, dim=-1)
    avg_size = N / cluster_size
    min_size = max(1, int(avg_size * min_size_ratio))
    max_size = int(avg_size * max_size_ratio)
    print(f"  avg={avg_size:.1f}  min={min_size}  max={max_size} per cluster")

    print("  K-means++ initialisation...")
    centroids = torch.zeros(cluster_size, d, device=device)
    first = torch.randint(0, N, (1,)).item()
    centroids[0] = emb[first]

    for k in tqdm(range(1, cluster_size), desc="  kmeans++ init", leave=False):
        sims = torch.mm(emb, centroids[:k].T)
        max_sim = sims.max(dim=1).values
        dist = (1.0 - max_sim).clamp(min=0.0)
        probs = dist ** 2
        probs = probs / probs.sum()
        idx = torch.multinomial(probs, 1).item()
        centroids[k] = emb[idx]

    centroids = F.normalize(centroids, dim=-1)
    cluster_ids = torch.zeros(N, dtype=torch.long, device=device)

    for iteration in range(max_iters):
        affinities = torch.zeros(N, cluster_size, device=device)
        for i in range(0, N, batch_size):
            j = min(i + batch_size, N)
            affinities[i:j] = torch.mm(emb[i:j], centroids.T)

        cluster_ids = affinities.argmax(dim=1)
        sizes = torch.bincount(cluster_ids, minlength=cluster_size)
        num_reassigned = 0

        oversized = (sizes > max_size).nonzero(as_tuple=False).squeeze(1)
        for c in oversized.tolist():
            excess = sizes[c].item() - max_size
            pts = (cluster_ids == c).nonzero(as_tuple=False).squeeze(1)
            aff_c = affinities[pts, c].clone()
            aff_tmp = affinities[pts].clone()
            aff_tmp[:, c] = -float("inf")
            aff_2nd, _ = aff_tmp.max(dim=1)
            loss = aff_c - aff_2nd
            _, order = loss.topk(excess, largest=False)
            for p in pts[order].tolist():
                new_c = affinities[p].clone()
                new_c[c] = -float("inf")
                cluster_ids[p] = new_c.argmax().item()
                num_reassigned += 1

        sizes = torch.bincount(cluster_ids, minlength=cluster_size)
        undersized = (sizes < min_size).nonzero(as_tuple=False).squeeze(1)
        for c in undersized.tolist():
            needed = min_size - sizes[c].item()
            candidates = (cluster_ids != c).nonzero(as_tuple=False).squeeze(1)
            if candidates.numel() == 0:
                continue
            aff_c = affinities[candidates, c]
            _, order = aff_c.topk(min(needed, candidates.numel()), largest=True)
            pts_to_move = candidates[order]
            cluster_ids[pts_to_move] = c
            num_reassigned += len(pts_to_move)

        new_centroids = torch.zeros_like(centroids)
        for k in range(cluster_size):
            mask = cluster_ids == k
            if mask.any():
                new_centroids[k] = emb[mask].mean(dim=0)
        new_centroids = F.normalize(new_centroids, dim=-1)

        cos_change = 1.0 - (centroids * new_centroids).sum(dim=1).mean().item()
        centroids = new_centroids

        print(f"  iter {iteration + 1:3d}: Δ={cos_change:.6f}  reassigned={num_reassigned}")
        if cos_change < tolerance:
            print(f"  Converged after {iteration + 1} iterations.")
            break

    sizes = torch.bincount(cluster_ids, minlength=cluster_size)
    print(f"  Cluster size  min={sizes.min().item()}  "
          f"max={sizes.max().item()}  "
          f"mean={sizes.float().mean().item():.1f}  "
          f"std={sizes.float().std().item():.1f}")
    return centroids


def cmd_learn(args: argparse.Namespace):
    config = load_config(args.config)
    device = torch.device(args.device)

    tokenizer = build_tokenizer(config)
    model_cfg = config["model"]
    hier_cfg = config.get("hierarchy", {})
    level_sizes = model_cfg["level_sizes"]
    mask_id = tokenizer.mask_token_id

    config["model"]["vocab_size"] = len(tokenizer)
    if "level_sizes" in config["model"]:
        config["model"]["level_sizes"][0] = len(tokenizer)

    gpt2_path = Path(
        args.gpt2_path
        or hier_cfg.get("gpt2_path", str(ROOT / "models" / "gpt2"))
    )
    if gpt2_path.exists() and (
        (gpt2_path / "pytorch_model.bin").exists()
        or any(gpt2_path.glob("model.safetensors"))
    ):
        all_emb = load_gpt2_embeddings(str(gpt2_path)).to(device)
        V_tok = model_cfg["vocab_size"]
        if all_emb.shape[0] < V_tok:
            pad = torch.zeros(V_tok - all_emb.shape[0], all_emb.shape[1], device=device)
            all_emb = torch.cat([all_emb, pad], dim=0)
        elif all_emb.shape[0] > V_tok:
            all_emb = all_emb[:V_tok]
        print(f"Leaf embeddings: GPT-2 wte  {all_emb.shape}")
    else:
        model = SADModel(
            vocab_size=model_cfg["vocab_size"],
            hidden_size=model_cfg["hidden_size"],
            n_blocks=model_cfg["n_blocks"],
            n_heads=model_cfg["n_heads"],
            cond_dim=model_cfg["cond_dim"],
            max_seq_len=model_cfg["max_seq_len"],
            block_size=model_cfg.get("block_size", 8),
            num_levels=model_cfg.get("num_levels", 2),
            level_sizes=level_sizes,
        ).to(device)
        with torch.no_grad():
            all_emb = model.get_leaf_embeddings().float().to(device)
        print(f"Leaf embeddings: random init  {all_emb.shape}")

    valid_ids = [i for i in range(all_emb.shape[0]) if i != mask_id]
    leaf_emb = all_emb[valid_ids].detach()
    print(f"Valid tokens:    {len(valid_ids)}  (excluded [MASK] id={mask_id})")
    print(f"Level sizes:     {level_sizes}")
    print(f"max_iters={args.max_iters}  tol={args.tolerance}  "
          f"min_ratio={args.min_size_ratio}  max_ratio={args.max_size_ratio}")
    print(f"Device:          {device}\n")

    prototypes = []
    prev_emb = leaf_emb
    for l_idx in range(1, len(level_sizes)):
        K = level_sizes[l_idx]
        print(f"Level {l_idx}  (K={K})  input_size={prev_emb.shape[0]}")
        centroids = semantic_kmeans_cluster(
            embeddings=prev_emb,
            cluster_size=K,
            max_iters=args.max_iters,
            tolerance=args.tolerance,
            batch_size=args.batch_size,
            min_size_ratio=args.min_size_ratio,
            max_size_ratio=args.max_size_ratio,
            seed=args.seed + l_idx,
        )
        centroids = F.normalize(centroids, dim=-1)
        prototypes.append(centroids.cpu())
        print(f"  → centroids {centroids.shape}  "
              f"norm min={centroids.norm(dim=-1).min():.4f} "
              f"max={centroids.norm(dim=-1).max():.4f}\n")
        prev_emb = centroids

    cache_path = hier_cfg.get("prototype_cache", None)
    if cache_path is None:
        cache_path = f"data/hierarchy_prototypes_{'-'.join(str(s) for s in level_sizes)}.pt"
    cache_path = Path(cache_path)
    cache_path.parent.mkdir(parents=True, exist_ok=True)
    torch.save({"prototypes": prototypes}, cache_path)
    print(f"Saved to {cache_path}")
    for i, p in enumerate(prototypes):
        print(f"  Level {i + 1}: {p.shape}")


# ── build_lut: top-k soft ancestor tables ─────────────────────────────────────


def build_topk_soft_assignment(
    child_embeddings: torch.Tensor,
    parent_prototypes: torch.Tensor,
    top_k: int,
    temp: float,
    batch_size: int = 2048,
) -> dict:
    device = child_embeddings.device
    child = F.normalize(child_embeddings, dim=-1)
    parent = F.normalize(parent_prototypes.to(device), dim=-1)
    N, d = child.shape
    K = parent.shape[0]
    k = min(top_k, K)

    all_indices = torch.empty(N, k, dtype=torch.long, device="cpu")
    all_probs = torch.empty(N, k, dtype=torch.float32, device="cpu")

    for i in tqdm(range(0, N, batch_size), desc="  top-k batch", leave=False):
        j = min(i + batch_size, N)
        sim = torch.mm(child[i:j], parent.T)
        vals, idxs = sim.topk(k, dim=-1, largest=True)
        probs = (vals / max(temp, 1e-6)).softmax(dim=-1)
        all_indices[i:j] = idxs.cpu()
        all_probs[i:j] = probs.cpu()

    return {"indices": all_indices, "probs": all_probs}


def cmd_build_lut(args: argparse.Namespace):
    config = load_config(args.config)
    device = torch.device(args.device)

    tokenizer = build_tokenizer(config)
    model_cfg = config["model"]
    level_sizes = model_cfg["level_sizes"]
    num_ancestor_levels = len(level_sizes) - 1
    mask_id = tokenizer.mask_token_id

    if args.top_k_per_level is not None:
        vals = [int(x.strip()) for x in args.top_k_per_level.split(",") if x.strip()]
        if len(vals) != num_ancestor_levels:
            raise ValueError(f"--top_k_per_level expects {num_ancestor_levels} values, got {len(vals)}")
        top_k_per_level = vals
    else:
        top_k_per_level = [args.top_k] * num_ancestor_levels

    all_emb = get_leaf_embeddings(config, args.gpt2_path, device)
    valid_ids = [i for i in range(all_emb.shape[0]) if i != mask_id]
    leaf_emb = all_emb[valid_ids].detach()
    print(f"Valid tokens: {len(valid_ids)}  (excluded [MASK] id={mask_id})")

    proto_path = Path(
        args.prototype_path
        or config.get("hierarchy", {}).get("prototype_cache", "")
        or config.get("ancestor", {}).get("proto_path", "")
    )
    if not proto_path.exists():
        raise FileNotFoundError(f"Prototype file not found: {proto_path}")
    data = torch.load(proto_path, map_location="cpu", weights_only=False)
    prototypes = data["prototypes"]
    print(f"Loaded {len(prototypes)} prototype level(s) from {proto_path}")
    for i, p in enumerate(prototypes):
        print(f"  Level {i + 1}: {tuple(p.shape)}")

    out_dict = {"top_k": args.top_k, "top_k_per_level": top_k_per_level, "temp": args.temp}
    prev_emb = leaf_emb
    if args.from_leaf_for_all_levels:
        print("Assignment mode: leaf-direct for all ancestor levels")
    else:
        print("Assignment mode: chained through lower-level prototypes")

    for l_idx in range(1, len(level_sizes)):
        proto = prototypes[l_idx - 1]
        level_top_k = top_k_per_level[l_idx - 1]
        child_emb = leaf_emb if args.from_leaf_for_all_levels else prev_emb
        print(f"\nBuilding level {l_idx} top-{level_top_k} assignments "
              f"({child_emb.shape[0]} -> {proto.shape[0]})...")
        level_data = build_topk_soft_assignment(
            child_embeddings=child_emb,
            parent_prototypes=proto,
            top_k=level_top_k,
            temp=args.temp,
            batch_size=args.batch_size,
        )
        out_dict[f"level_{l_idx}"] = level_data
        print(f"  indices: {level_data['indices'].shape}, probs: {level_data['probs'].shape}")
        prev_emb = proto

    if args.output is None:
        level_tag = "-".join(str(s) for s in level_sizes)
        if args.top_k_per_level is not None:
            mode_tag = "leaf_to_each_level" if args.from_leaf_for_all_levels else "chained"
            topk_tag = "-".join(str(k) for k in top_k_per_level)
            out_path = Path(
                f"data/ancestor_lut_{level_tag}_{mode_tag}_topk_{topk_tag}_t{args.temp}.pt"
            )
        else:
            out_path = Path(f"data/ancestor_lut_{level_tag}_top{args.top_k}_t{args.temp}.pt")
    else:
        out_path = Path(args.output)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(out_dict, out_path)
    print(f"\nSaved ancestor LUT to {out_path}")


# ── extend: add deeper levels to existing prototypes ──────────────────────────


def cmd_extend(args: argparse.Namespace):
    in_path = Path(args.input)
    out_path = Path(args.output)
    target_levels = [int(x) for x in args.levels.split(",") if x.strip()]
    device = torch.device(args.device)

    data = torch.load(in_path, map_location="cpu", weights_only=False)
    if isinstance(data, dict):
        if "prototypes" in data:
            prototypes = [p.float() for p in data["prototypes"]]
        else:
            level_keys = sorted(
                [k for k in data.keys() if k.startswith("level_")],
                key=lambda k: int(k.split("_")[1]),
            )
            prototypes = [data[k].float() for k in level_keys]
    else:
        raise ValueError(f"Unsupported prototype format in {in_path}")

    if not prototypes:
        raise ValueError("No prototypes found in input file")
    if prototypes[0].shape[0] != target_levels[0]:
        raise ValueError(
            f"Input level_1 has {prototypes[0].shape[0]} clusters, expected {target_levels[0]}"
        )

    print(f"Loaded {len(prototypes)} existing level(s) from {in_path}")
    for i, p in enumerate(prototypes, start=1):
        print(f"  level_{i}: {tuple(p.shape)}")

    existing_levels = [p.shape[0] for p in prototypes]
    if existing_levels == target_levels:
        print("Input already matches requested hierarchy; re-saving in dual format.")
    elif existing_levels != target_levels[:len(existing_levels)]:
        raise ValueError(
            f"Existing hierarchy sizes {existing_levels} are not a prefix of requested {target_levels}"
        )

    prev_emb = prototypes[-1].to(device)
    for l_idx in range(len(prototypes) + 1, len(target_levels) + 1):
        K = target_levels[l_idx - 1]
        print(f"\nLearning level {l_idx}  (K={K})  input_size={prev_emb.shape[0]}")
        centroids = semantic_kmeans_cluster(
            embeddings=prev_emb,
            cluster_size=K,
            max_iters=args.max_iters,
            tolerance=args.tolerance,
            batch_size=args.batch_size,
            min_size_ratio=args.min_size_ratio,
            max_size_ratio=args.max_size_ratio,
            seed=args.seed + l_idx,
        )
        centroids = F.normalize(centroids, dim=-1)
        prototypes.append(centroids.cpu())
        print(f"  → centroids {centroids.shape}")
        prev_emb = centroids

    out_path.parent.mkdir(parents=True, exist_ok=True)
    torch.save({"prototypes": prototypes}, out_path)
    print(f"\nSaved extended hierarchy to {out_path}")
    for i, p in enumerate(prototypes, start=1):
        print(f"  level_{i}: {tuple(p.shape)}")


# ── merge: combine per-level LUT files ────────────────────────────────────────


def cmd_merge(args: argparse.Namespace):
    merged = {"temp": args.temp}
    top_k_per_level = []

    for spec in args.source:
        level_str, path_str = spec.split(":", 1)
        level = int(level_str)
        path = Path(path_str)
        data = torch.load(path, map_location="cpu", weights_only=False)
        key = f"level_{level}"
        if key not in data:
            raise KeyError(f"{path} does not contain {key}")
        merged[key] = data[key]
        top_k_per_level.append(int(data[key]["indices"].shape[1]))
        print(f"Merged {key} from {path} with top_k={top_k_per_level[-1]}")

    merged["top_k_per_level"] = top_k_per_level
    out_path = Path(args.output)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(merged, out_path)
    print(f"Saved merged LUT to {out_path}")


# ── CLI entrypoint ────────────────────────────────────────────────────────────


def main():
    p = argparse.ArgumentParser(description="Unified hierarchy construction pipeline")
    sub = p.add_subparsers(dest="cmd", required=True)

    # learn
    sp = sub.add_parser("learn", help="Cosine k-means++ clustering for prototypes")
    sp.add_argument("--config", default="configs/sad_owt.yaml")
    sp.add_argument("--gpt2_path", type=str, default=None)
    sp.add_argument("--max_iters", type=int, default=100)
    sp.add_argument("--tolerance", type=float, default=1e-4)
    sp.add_argument("--batch_size", type=int, default=2048)
    sp.add_argument("--min_size_ratio", type=float, default=0.1)
    sp.add_argument("--max_size_ratio", type=float, default=5.0)
    sp.add_argument("--seed", type=int, default=42)
    sp.add_argument("--device", type=str,
                    default="cuda" if torch.cuda.is_available() else "cpu")
    sp.set_defaults(func=cmd_learn)

    # build_lut
    sp = sub.add_parser("build_lut", help="Build sparse soft top-k ancestor LUT")
    sp.add_argument("--config", type=str, default="configs/sad_owt.yaml")
    sp.add_argument("--prototype_path", type=str, default=None)
    sp.add_argument("--gpt2_path", type=str, default=None)
    sp.add_argument("--output", type=str, default=None)
    sp.add_argument("--top_k", type=int, default=3)
    sp.add_argument("--top_k_per_level", type=str, default=None)
    sp.add_argument("--temp", type=float, default=1.0)
    sp.add_argument("--batch_size", type=int, default=2048)
    sp.add_argument("--from_leaf_for_all_levels", action="store_true")
    sp.add_argument("--device", type=str,
                    default="cuda" if torch.cuda.is_available() else "cpu")
    sp.set_defaults(func=cmd_build_lut)

    # extend
    sp = sub.add_parser("extend", help="Extend existing hierarchy with deeper levels")
    sp.add_argument("--input", type=str, required=True)
    sp.add_argument("--levels", type=str, required=True)
    sp.add_argument("--output", type=str, required=True)
    sp.add_argument("--max_iters", type=int, default=100)
    sp.add_argument("--tolerance", type=float, default=1e-4)
    sp.add_argument("--batch_size", type=int, default=2048)
    sp.add_argument("--min_size_ratio", type=float, default=0.1)
    sp.add_argument("--max_size_ratio", type=float, default=5.0)
    sp.add_argument("--seed", type=int, default=42)
    sp.add_argument("--device", type=str, default="cpu")
    sp.set_defaults(func=cmd_extend)

    # merge
    sp = sub.add_parser("merge", help="Merge per-level LUT files")
    sp.add_argument("--output", type=str, required=True)
    sp.add_argument("--source", action="append", required=True)
    sp.add_argument("--temp", type=float, default=1.0)
    sp.set_defaults(func=cmd_merge)

    args = p.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()