"""Small SFT diagnostics for checkpoint quality and trainability.

This script intentionally bypasses the full trainer so it can answer one narrow
question quickly: can the checkpoint reduce response-only SFT loss on a tiny,
fixed batch?
"""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path
from typing import Any

import torch

from taoTrain.checkpointing.checkpoint import CheckpointManager
from taoTrain.config import TrainingModeEnum, load_config
from taoTrain.core import create_model
from taoTrain.data.sft_utils import build_sft_sequence_tokens, parse_sft_record
try:
    from taoTrain.data.sft_utils import build_response_only_next_token_labels
except ImportError:
    def build_response_only_next_token_labels(input_ids: list[int], mask: list[int]) -> list[int]:
        labels = [token_id if mask_value else -100 for token_id, mask_value in zip(input_ids, mask)]
        return labels[1:] + [-100]
from taoTrain.data.tokenizer import SentencePieceTokenizerWrapper
from taoTrain.utils import set_seed


def load_tokenizer(tokenizer_path: str):
    path = Path(tokenizer_path)
    if path.suffix == ".model":
        import sentencepiece as spm

        sp = spm.SentencePieceProcessor()
        sp.Load(str(path))
        return SentencePieceTokenizerWrapper(sp)

    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
    if getattr(tokenizer, "pad_token", None) is None and getattr(tokenizer, "eos_token", None):
        tokenizer.pad_token = tokenizer.eos_token
    return tokenizer


def read_jsonl_records(path: str, limit: int) -> list[dict[str, Any]]:
    records = []
    with open(path, "r", encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if not line:
                continue
            records.append(json.loads(line))
            if len(records) >= limit:
                break
    return records


def build_batch(config, tokenizer, records: list[dict[str, Any]], device: torch.device) -> dict[str, torch.Tensor]:
    input_rows = []
    attention_rows = []
    label_rows = []
    train_tokens = []

    for record in records:
        turns, _ = parse_sft_record(record, config)
        if not turns:
            continue
        input_ids, attention_mask, mask = build_sft_sequence_tokens(
            turns=turns,
            tokenizer=tokenizer,
            user_token=getattr(config, "user_token", "<user>"),
            assistant_token=getattr(config, "assistant_token", "<assistant>"),
            max_seq_length=config.model.max_seq_length,
        )
        labels = build_response_only_next_token_labels(input_ids, mask)
        input_rows.append(input_ids)
        attention_rows.append(attention_mask)
        label_rows.append(labels)
        train_tokens.append(sum(1 for value in labels if value != -100))

    if not input_rows:
        raise ValueError("No valid SFT records found for the diagnostic batch")

    return {
        "input_ids": torch.tensor(input_rows, dtype=torch.long, device=device),
        "attention_mask": torch.tensor(attention_rows, dtype=torch.long, device=device),
        "labels": torch.tensor(label_rows, dtype=torch.long, device=device),
        "train_tokens": torch.tensor(train_tokens, dtype=torch.long),
    }


@torch.no_grad()
def score_batch(model, batch: dict[str, torch.Tensor], dtype: torch.dtype) -> float:
    model.eval()
    device_type = "cuda" if batch["input_ids"].is_cuda else "cpu"
    enabled = device_type == "cuda" and dtype in (torch.float16, torch.bfloat16)
    with torch.autocast(device_type=device_type, dtype=dtype, enabled=enabled):
        outputs = model(
            input_ids=batch["input_ids"],
            attention_mask=batch["attention_mask"],
            labels=batch["labels"],
        )
    return float(outputs["loss"].detach().cpu())


def grad_l2_norm(parameters) -> float:
    total = 0.0
    for parameter in parameters:
        if parameter.grad is None:
            continue
        grad = parameter.grad.detach()
        total += float(torch.sum(grad.float() * grad.float()).cpu())
    return math.sqrt(total)


def grad_summary(named_parameters, max_items: int = 12) -> dict[str, Any]:
    groups: dict[str, dict[str, Any]] = {}
    worst = []
    nonfinite = []

    for name, parameter in named_parameters:
        if parameter.grad is None:
            continue
        grad = parameter.grad.detach().float()
        finite = torch.isfinite(grad)
        finite_count = int(finite.sum().cpu())
        numel = grad.numel()
        finite_abs_max = float(grad[finite].abs().max().cpu()) if finite_count else float("inf")
        has_nonfinite = finite_count != numel
        if has_nonfinite:
            nonfinite.append(name)

        if ".layers." in name:
            parts = name.split(".")
            try:
                idx = parts.index("layers")
                group = "layer_" + parts[idx + 1]
            except (ValueError, IndexError):
                group = "layers"
        else:
            group = name.split(".", 1)[0]

        entry = groups.setdefault(group, {
            "numel": 0,
            "finite": 0,
            "nonfinite_tensors": 0,
            "max_abs_grad": 0.0,
        })
        entry["numel"] += numel
        entry["finite"] += finite_count
        entry["nonfinite_tensors"] += int(has_nonfinite)
        entry["max_abs_grad"] = max(entry["max_abs_grad"], finite_abs_max)
        worst.append((finite_abs_max, name))

    worst.sort(reverse=True, key=lambda item: item[0])
    return {
        "groups": groups,
        "worst_tensors": [{"name": name, "max_abs_grad": value} for value, name in worst[:max_items]],
        "nonfinite_tensors": nonfinite[:max_items],
        "nonfinite_tensor_count": len(nonfinite),
    }


def freeze_ssm_core_parameters(model) -> int:
    frozen = 0
    markers = (
        ".ssm_lanes.",
        ".ssm.",
    )
    for name, parameter in model.named_parameters():
        if any(marker in name for marker in markers):
            parameter.requires_grad_(False)
            frozen += parameter.numel()
    return frozen


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", required=True)
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--samples", type=int, default=2)
    parser.add_argument("--steps", type=int, default=80)
    parser.add_argument("--lr", type=float, default=3e-4)
    parser.add_argument("--log-every", type=int, default=10)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--dtype", choices=["config", "float32", "float16", "bfloat16"], default="config")
    parser.add_argument("--no-clip", action="store_true")
    parser.add_argument("--freeze-ssm-core", action="store_true")
    parser.add_argument("--ssm-branch-rms-norm", action="store_true")
    parser.add_argument("--ssm-branch-clip-value", type=float, default=None)
    parser.add_argument("--block-residual-rms-norm", action="store_true")
    parser.add_argument("--block-residual-rms-target", type=float, default=None)
    parser.add_argument("--seed", type=int, default=123)
    args = parser.parse_args()

    set_seed(args.seed)
    config = load_config(args.config, TrainingModeEnum.SFT)
    if args.ssm_branch_rms_norm:
        config.model.ssm_branch_rms_norm = True
    if args.ssm_branch_clip_value is not None:
        config.model.ssm_branch_clip_value = args.ssm_branch_clip_value
    if args.block_residual_rms_norm:
        config.model.block_residual_rms_norm = True
    if args.block_residual_rms_target is not None:
        config.model.block_residual_rms_target = args.block_residual_rms_target
    device = torch.device(args.device if args.device == "cpu" or torch.cuda.is_available() else "cpu")
    if args.dtype == "float32":
        dtype = torch.float32
    elif args.dtype == "float16":
        dtype = torch.float16
    elif args.dtype == "bfloat16":
        dtype = torch.bfloat16
    else:
        dtype = torch.bfloat16 if str(config.dtype) == "DataTypeEnum.BFLOAT16" or str(config.dtype) == "bfloat16" else torch.float32

    tokenizer = load_tokenizer(config.dataset.tokenizer_path)
    records = read_jsonl_records(config.dataset.jsonl_path, args.samples)
    batch = build_batch(config, tokenizer, records, device)

    model = create_model(config, device)
    checkpoint = CheckpointManager(config.checkpoint_dir).load(args.checkpoint, device=device)
    model.load_state_dict(checkpoint["model_state"], strict=False)
    frozen_params = freeze_ssm_core_parameters(model) if args.freeze_ssm_core else 0

    initial_loss = score_batch(model, batch, dtype)

    trainable_params = [parameter for parameter in model.parameters() if parameter.requires_grad]
    optimizer = torch.optim.AdamW(trainable_params, lr=args.lr, weight_decay=0.0)
    history = []
    device_type = "cuda" if device.type == "cuda" else "cpu"
    autocast_enabled = device_type == "cuda" and dtype in (torch.float16, torch.bfloat16)

    model.train()
    for step in range(1, args.steps + 1):
        optimizer.zero_grad(set_to_none=True)
        with torch.autocast(device_type=device_type, dtype=dtype, enabled=autocast_enabled):
            outputs = model(
                input_ids=batch["input_ids"],
                attention_mask=batch["attention_mask"],
                labels=batch["labels"],
            )
            loss = outputs["loss"]
        loss.backward()
        grad_norm = grad_l2_norm(trainable_params)
        stats = None
        if step == 1 or step % args.log_every == 0 or step == args.steps:
            stats = grad_summary(model.named_parameters())
        if not args.no_clip:
            torch.nn.utils.clip_grad_norm_(trainable_params, 1.0)
        optimizer.step()

        if step == 1 or step % args.log_every == 0 or step == args.steps:
            item = {
                "step": step,
                "loss": float(loss.detach().cpu()),
                "grad_l2_norm": grad_norm,
            }
            if stats is not None:
                item["grad_summary"] = stats
            history.append(item)

    final_loss = score_batch(model, batch, dtype)
    result = {
        "checkpoint": str(Path(args.checkpoint)),
        "config": str(Path(args.config)),
        "dataset": config.dataset.jsonl_path,
        "samples": len(records),
        "sequence_length": config.model.max_seq_length,
        "train_tokens_per_sample": batch["train_tokens"].tolist(),
        "lr": args.lr,
        "steps": args.steps,
        "clip_grad_norm": not args.no_clip,
        "freeze_ssm_core": args.freeze_ssm_core,
        "ssm_branch_rms_norm": config.model.ssm_branch_rms_norm,
        "ssm_branch_clip_value": config.model.ssm_branch_clip_value,
        "block_residual_rms_norm": config.model.block_residual_rms_norm,
        "block_residual_rms_target": config.model.block_residual_rms_target,
        "frozen_params": frozen_params,
        "trainable_params": sum(parameter.numel() for parameter in trainable_params),
        "initial_loss": initial_loss,
        "final_loss": final_loss,
        "loss_delta": final_loss - initial_loss,
        "history": history,
        "device": str(device),
        "dtype": str(dtype),
    }

    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(json.dumps(result, indent=2), encoding="utf-8")
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()