"""Student = Qwen3.5 backbone + LoRA + Clef-style JointSchemaHead (option logits -> softmax per question)
(+ optional BoxHead for Jev-style `box` questions, see glint/box.py)."""
from __future__ import annotations
import json, os
from pathlib import Path
import torch
from safetensors.torch import load_file, save_file
from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration
from peft import LoraConfig, PeftModel, get_peft_model
from .joint_schema_model import JointSchemaHead
from .box import BoxHead
from .glint_model import GlintModel

# LoRAは言語モデル側のみ（ビジョンエンコーダは凍結）。Gated DeltaNet層の射影も含める。
LORA_TARGET = (r".*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj"
               r"|in_proj_qkv|in_proj_z|out_proj)")

def default_device():
    return "cuda" if torch.cuda.is_available() else "cpu"

def parse_dtype(name: str):
    return {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[name]

def load_processor(base: str, max_pixels: int, min_pixels: int = 65536):
    proc = AutoProcessor.from_pretrained(base)
    ip = proc.image_processor
    ip.size = {"shortest_edge": min_pixels, "longest_edge": max_pixels}
    if hasattr(ip, "max_pixels"):
        ip.max_pixels = max_pixels
    if hasattr(ip, "min_pixels"):
        ip.min_pixels = min_pixels
    return proc

def build_student(base="Qwen/Qwen3.5-2B", max_pixels=430080, device=None, dtype="bf16",
                  lora_r=32, lora_alpha=64, lora_dropout=0.05, head_cfg=None, grad_ckpt=False, box_head_cfg=None):
    """box_head_cfg: None = no box head (previous behaviour); a dict (may be empty) = BoxHead kwargs."""
    device = device or default_device()
    td = parse_dtype(dtype)
    proc = load_processor(base, max_pixels)
    backbone = Qwen3_5ForConditionalGeneration.from_pretrained(base, dtype=td).to(device)
    backbone.config.use_cache = False
    if grad_ckpt:
        backbone.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
        backbone.enable_input_require_grads()
    backbone = get_peft_model(backbone, LoraConfig(r=lora_r, lora_alpha=lora_alpha, lora_dropout=lora_dropout,
                                                   target_modules=LORA_TARGET))
    hidden = backbone.config.text_config.hidden_size
    head_cfg = dict(head_cfg or {"width": 512, "routing_layers": 2, "layers": 2, "heads": 8, "feedforward": 2048})
    head_cfg["hidden_size"] = hidden
    head = JointSchemaHead(**head_cfg, dropout=0.1).to(device, td)
    box_head = None
    if box_head_cfg is not None:                                   # built after the joint head: its init RNG is unchanged
        box_head_cfg = {**box_head_cfg, "hidden_size": hidden}
        box_head = BoxHead(**box_head_cfg, dropout=0.1).to(device, td)
    model = GlintModel(backbone, head, box_head)
    meta = {"base_model": base, "max_pixels": max_pixels, "dtype": dtype, "head_config": head_cfg,
            "lora": {"r": lora_r, "alpha": lora_alpha, "dropout": lora_dropout, "target_modules": LORA_TARGET}}
    if box_head_cfg is not None:
        meta["box_head_config"] = box_head_cfg
    return model, proc, meta

def save_checkpoint(model: GlintModel, meta: dict, out_dir: str | os.PathLike, extra: dict | None = None):
    out = Path(out_dir); out.mkdir(parents=True, exist_ok=True)
    model.language_model.save_pretrained(out / "adapter")          # LoRAのみ
    save_file({k: v.detach().cpu().contiguous() for k, v in model.head.state_dict().items()}, out / "head.safetensors")
    if getattr(model, "box_head", None) is not None:
        save_file({k: v.detach().cpu().contiguous() for k, v in model.box_head.state_dict().items()}, out / "box_head.safetensors")
    (out / "glint_config.json").write_text(json.dumps({**meta, **(extra or {})}, indent=2, ensure_ascii=False), encoding="utf-8")

def load_checkpoint(ckpt_dir: str | os.PathLike, device=None, max_pixels: int | None = None, merge=True):
    ckpt = Path(ckpt_dir); device = device or default_device()
    meta = json.loads((ckpt / "glint_config.json").read_text(encoding="utf-8"))
    td = parse_dtype(meta["dtype"])
    proc = load_processor(meta["base_model"], max_pixels or meta["max_pixels"])
    backbone = Qwen3_5ForConditionalGeneration.from_pretrained(meta["base_model"], dtype=td).to(device)
    backbone.config.use_cache = False
    backbone = PeftModel.from_pretrained(backbone, ckpt / "adapter")
    if merge:
        backbone = backbone.merge_and_unload()                     # 推論時はLoRAをマージして高速化
    head = JointSchemaHead(**meta["head_config"]).to(device, td)
    head.load_state_dict(load_file(ckpt / "head.safetensors"), strict=True)
    box_head = None
    if meta.get("box_head_config") is not None:                    # older checkpoints have no box head
        box_head = BoxHead(**meta["box_head_config"]).to(device, td)
        box_head.load_state_dict(load_file(ckpt / "box_head.safetensors"), strict=True)
    return GlintModel(backbone, head, box_head).eval(), proc, meta