from __future__ import annotations
import argparse
import json
import logging
import random
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
import torch
from transformers import AutoTokenizer
from riemann.growth_ladder import get_stage
from riemann.modeling_riemann import RiemannForCausalLM
from riemann.device import main_device, cpu_device
from riemann.paths import DATA_DIR, CHECKPOINT_ROOT, CONTROLLER_ROOT

logger = logging.getLogger("riemann.fit_ntk")
TOKENIZER_NAME = "gpt2"


def load_examples(path: Path, limit: int | None, seed: int) -> list:
    from ntkmirror import load_jsonl_examples

    examples = list(load_jsonl_examples(str(path)))
    if limit is not None and len(examples) > limit:
        rng = random.Random(seed)
        rng.shuffle(examples)
        examples = examples[:limit]
    return examples


def split_validation(examples: list, fraction: float, seed: int) -> tuple[list, list]:
    if fraction <= 0 or len(examples) < 20:
        return examples, []
    rng = random.Random(seed)
    idx = list(range(len(examples)))
    rng.shuffle(idx)
    n_val = max(8, int(len(examples) * fraction))
    val_idx = set(idx[:n_val])
    train = [e for i, e in enumerate(examples) if i not in val_idx]
    val = [e for i, e in enumerate(examples) if i in val_idx]
    return train, val


def warm_start_from(tuner, prior_path: Path, device, examples: list, args) -> dict:
    from ntkmirror.controller import SignedLogMaskState, _SignedLogMaskModule

    prior = SignedLogMaskState.load(str(prior_path), map_location="cpu")
    n_layers = len(tuner.decoder_layers)
    hidden = tuner.hidden_size
    keep = (prior.layer_indices < n_layers) & (prior.channel_indices < hidden)
    layer_idx = prior.layer_indices[keep]
    chan_idx = prior.channel_indices[keep]
    raw = prior.raw[keep]

    fresh_layer_idx, fresh_chan_idx = tuner._score_gates(
        examples,
        score_batches=min(16, max(1, len(examples) // max(1, args.batch_size))),
        batch_size=args.batch_size,
        max_length=args.max_length,
    )
    fresh_layer_idx = fresh_layer_idx.cpu()
    fresh_chan_idx = fresh_chan_idx.cpu()
    existing = set(zip(layer_idx.tolist(), chan_idx.tolist()))
    new_pairs = [
        (l, c)
        for l, c in zip(fresh_layer_idx.tolist(), fresh_chan_idx.tolist())
        if (l, c) not in existing and (l >= prior.n_layers or c >= prior.hidden_size)
    ]
    budget = max(0, args.gates - layer_idx.numel())
    new_pairs = new_pairs[:budget]
    if new_pairs:
        nl = torch.tensor([p[0] for p in new_pairs], dtype=torch.long)
        nc = torch.tensor([p[1] for p in new_pairs], dtype=torch.long)
        layer_idx = torch.cat([layer_idx, nl])
        chan_idx = torch.cat([chan_idx, nc])
        raw = torch.cat([raw, torch.zeros(len(new_pairs), dtype=raw.dtype)])

    if tuner.controller is not None:
        tuner.controller.remove()
    tuner.controller = _SignedLogMaskModule(
        tuner.decoder_layers,
        layer_idx.to(device),
        chan_idx.to(device),
        hidden_size=hidden,
        max_log_gate=tuner.max_log_gate,
        hook_site=tuner.hook_site,
        raw_init=raw.to(device),
    ).to(device)

    def _keep_warm_controller(*_a, **_k):
        return {
            "selected_gates": float(tuner.controller.raw.numel()),
            "select_seconds": 0.0,
            "hook_site": tuner.hook_site,
        }

    tuner.initialize_controller = _keep_warm_controller
    return {
        "prior_gates": int(prior.n_gates),
        "carried_gates": int(keep.sum()),
        "new_gates_for_grown_dims": len(new_pairs),
        "total_gates": int(layer_idx.numel()),
        "prior_layers": int(prior.n_layers),
        "prior_hidden": int(prior.hidden_size),
    }


def main() -> None:
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s"
    )
    parser = argparse.ArgumentParser()
    parser.add_argument("--stage", default="stage0_124m")
    parser.add_argument("--gates", type=int, default=4096)
    parser.add_argument("--steps", type=int, default=240)
    parser.add_argument("--lr", type=float, default=5e-3)
    parser.add_argument("--batch-size", type=int, default=4)
    parser.add_argument("--max-length", type=int, default=512)
    parser.add_argument("--max-examples", type=int, default=20000)
    parser.add_argument("--val-fraction", type=float, default=0.05)
    parser.add_argument("--val-interval", type=int, default=20)
    parser.add_argument("--early-stop-patience", type=int, default=3)
    parser.add_argument("--kl-to-base", type=float, default=0.0)
    parser.add_argument("--seed", type=int, default=1234)
    parser.add_argument("--init-from", default=None)
    parser.add_argument("--cpu", action="store_true")
    parser.add_argument("--train-jsonl", default=str(DATA_DIR / "ntk_finetune.jsonl"))
    parser.add_argument("--checkpoint-root", default=str(CHECKPOINT_ROOT))
    parser.add_argument("--controller-root", default=str(CONTROLLER_ROOT))
    parser.add_argument("--out", default=None)
    args = parser.parse_args()

    stage = get_stage(args.stage)
    checkpoint_dir = Path(args.checkpoint_root) / stage.name
    if not (checkpoint_dir / "config.json").exists():
        raise FileNotFoundError(
            f"No trained checkpoint at {checkpoint_dir}. Run pretrain.py for {stage.name} first."
        )
    controller_root = Path(args.controller_root)
    controller_root.mkdir(parents=True, exist_ok=True)
    out_path = Path(args.out) if args.out else controller_root / f"{stage.name}.pt"
    train_path = Path(args.train_jsonl)
    if not train_path.exists():
        raise FileNotFoundError(f"{train_path} not found. Run riemann.data first.")

    try:
        from ntkmirror import ForwardFineTuner
    except ImportError as exc:
        raise RuntimeError(
            "ntkmirror is required: pip install git+https://github.com/leochlon/ntkmirror.git"
        ) from exc

    device = cpu_device() if args.cpu else main_device()
    torch.manual_seed(args.seed)
    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    logger.info("Loading %s from %s onto %s", stage.name, checkpoint_dir, device)
    model = RiemannForCausalLM.from_pretrained(str(checkpoint_dir))
    model.config.use_cache = False
    model.to(device)
    model.eval()

    examples = load_examples(train_path, args.max_examples, args.seed)
    train_examples, val_examples = split_validation(
        examples, args.val_fraction, args.seed
    )
    logger.info(
        "NTK fit data: %d train / %d validation examples (from %s)",
        len(train_examples),
        len(val_examples),
        train_path,
    )

    tuner = ForwardFineTuner(
        model,
        tokenizer,
        gates=args.gates,
        model_name=f"my-name-is-riemann/{stage.name}",
        tokenizer_name=TOKENIZER_NAME,
    )
    warm = None
    if args.init_from:
        prior_path = Path(args.init_from)
        if not prior_path.exists():
            prior_path = controller_root / f"{args.init_from}.pt"
        if not prior_path.exists():
            raise FileNotFoundError(f"No prior controller at {prior_path}")
        warm = warm_start_from(tuner, prior_path, device, train_examples, args)
        logger.info(
            "Warm start from %s: carried %d of %d gates, added %d for grown dims, %d total (%d layers, hidden %d -> %d layers, hidden %d)",
            prior_path.name,
            warm["carried_gates"],
            warm["prior_gates"],
            warm["new_gates_for_grown_dims"],
            warm["total_gates"],
            warm["prior_layers"],
            warm["prior_hidden"],
            len(tuner.decoder_layers),
            tuner.hidden_size,
        )

    base_nll = None
    if val_examples:
        base_eval = tuner.evaluate_nll(
            val_examples, batch_size=args.batch_size, max_length=args.max_length
        )
        base_nll = float(base_eval["nll"])
        logger.info(
            "Validation before fitting: nll %.4f, token acc %.4f",
            base_nll,
            base_eval["token_acc"],
        )

    t0 = time.time()
    stats = tuner.fit(
        train_examples,
        steps=args.steps,
        lr=args.lr,
        batch_size=args.batch_size,
        max_length=args.max_length,
        validation_examples=val_examples or None,
        validation_interval=args.val_interval if val_examples else None,
        early_stop_patience=args.early_stop_patience if val_examples else 0,
        select_best_on_validation=bool(val_examples),
        retain_examples=train_examples if args.kl_to_base > 0 else None,
        kl_to_base=args.kl_to_base,
        verbose=True,
    )
    elapsed = time.time() - t0

    final_nll = None
    if val_examples:
        final_eval = tuner.evaluate_nll(
            val_examples, batch_size=args.batch_size, max_length=args.max_length
        )
        final_nll = float(final_eval["nll"])
        logger.info(
            "Validation after fitting: nll %.4f (base %.4f, delta %+.4f), token acc %.4f",
            final_nll,
            base_nll,
            final_nll - base_nll,
            final_eval["token_acc"],
        )

    metadata = {
        "stage": stage.name,
        "checkpoint": str(checkpoint_dir),
        "train_examples": len(train_examples),
        "validation_examples": len(val_examples),
        "steps": args.steps,
        "lr": args.lr,
        "gates_requested": args.gates,
        "fit_seconds": elapsed,
        "validation_nll_base": base_nll,
        "validation_nll_final": final_nll,
        "warm_start": warm,
    }
    tuner.save(str(out_path), metadata=metadata)
    summary_path = out_path.with_suffix(".json")
    summary = {
        k: (v if not isinstance(v, torch.Tensor) else v.tolist())
        for k, v in stats.items()
        if not isinstance(v, list)
    }
    summary.update(metadata)
    summary_path.write_text(
        json.dumps(summary, indent=2, default=str), encoding="utf-8"
    )
    logger.info("Saved gate controller to %s (summary at %s)", out_path, summary_path)


if __name__ == "__main__":
    main()