#!/usr/bin/env python3
"""
Anima trajectory imitation distillation
========================================

DiffSynth-Studio の Z-Image trajectory imitation を Anima に最小移植したもの。
**critic なし、単一ネットワーク**で前回 5 回失敗の R3GAN 不安定性を完全回避。

アルゴリズム:
  - teacher = Anima base (frozen, no LoRA) を 50-step CFG=2.0 で rollout
  - student = Anima base + wide LoRA (rank 32) を 8-step CFG=1.0 で rollout
  - L_align = MSE(student velocity, teacher segment velocity) on 8 segments
  - L_reg   = LPIPS(student final image, teacher final image)  (任意、--lpips-weight>0 で有効)

warm-start:
  --warm-lora /models/loras/anima_turbo.safetensors  で Civitai 公式 Anima Turbo LoRA を
  student LoRA 初期値に注入できる。format は ComfyUI 形式 (diffusion_model.<...>.lora_A.weight)、
  内部で PEFT 形式に変換してロード。

使い方 (Modal 経由):
  # Smoke test (1 step, sanity check)
  modal run modal_app.py::train_traj_imitation --total-steps 1 --batch-size 1 \\
      --teacher-steps 12 --student-steps 8 --lpips-weight 0.0

  # 本番 (2000 step, ~$45, B200)
  modal run --detach modal_app.py::train_traj_imitation \\
      --total-steps 2000 --batch-size 1 --teacher-steps 50 --student-steps 8 \\
      --warm-lora /models/loras/anima_turbo.safetensors --lpips-weight 0.1
"""
from __future__ import annotations
import argparse
import copy
import json
import os
import sys
import time
from pathlib import Path

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
from safetensors.torch import save_file, load_file

# import-path: 同ディレクトリ前提
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from distill.anima_loader import AnimaPaths, build_anima, AnimaBundle
from distill.dmd2_trainer import attach_wide_lora
from distill.traj_scheduler import make_schedule, snap_to_targets
from distill.traj_loss import (
    fetch_trajectory, align_trajectory, compute_regularization,
)


# ---------------------------------------------------------------------------
# Text-only dataset: caption だけ使う (trajectory imitation は init=noise なので image 不要)
# ---------------------------------------------------------------------------
class TextOnlyDataset(Dataset):
    def __init__(self, root: str | Path):
        self.root = Path(root)
        # *.txt を再帰的に拾う。/dataset/cleaned 配下を想定。
        self.files = sorted(self.root.rglob("*.txt"))
        if not self.files:
            raise RuntimeError(f"No .txt files under {self.root}")

    def __len__(self) -> int:
        return len(self.files)

    def __getitem__(self, idx: int) -> str:
        return self.files[idx].read_text(encoding="utf-8").strip()


def text_collate(batch: list[str]) -> list[str]:
    return list(batch)


# ---------------------------------------------------------------------------
# ComfyUI Anima LoRA -> PEFT 形式変換 (warm-start 用)
# 既存の _convert_peft_to_comfy_lora の逆向き
# ---------------------------------------------------------------------------
def convert_comfy_to_peft_lora(sd: dict) -> dict:
    """ComfyUI(Anima/Cosmos)形式 LoRA を PEFT 形式に変換。
       Comfy:  'diffusion_model.<module>.lora_A.weight'
       PEFT:   'base_model.model.<module>.lora_A.default.weight'
    """
    out = {}
    for k, v in sd.items():
        nk = k
        if nk.startswith("diffusion_model."):
            nk = nk[len("diffusion_model."):]
        # adapter name '.default.' を挿入
        nk = nk.replace(".lora_A.weight", ".lora_A.default.weight")
        nk = nk.replace(".lora_B.weight", ".lora_B.default.weight")
        nk = nk.replace(".lora_A.bias", ".lora_A.default.bias")
        nk = nk.replace(".lora_B.bias", ".lora_B.default.bias")
        nk = "base_model.model." + nk
        out[nk] = v
    return out


def load_warm_lora(student_model, warm_lora_path: str) -> None:
    """Civitai 形式の Anima Turbo LoRA を student の PEFT 形式に変換して注入。
       shape / rank が合わない key は skip して警告。"""
    print(f"[warm] loading {warm_lora_path}")
    sd_comfy = load_file(warm_lora_path)
    print(f"[warm] {len(sd_comfy)} comfy keys, sample: {list(sd_comfy.keys())[:3]}")
    sd_peft = convert_comfy_to_peft_lora(sd_comfy)
    print(f"[warm] converted to {len(sd_peft)} peft keys, sample: {list(sd_peft.keys())[:3]}")

    # student の現 state_dict と shape を照合して filter
    model_sd = student_model.state_dict()
    matched, skipped_shape, missing = 0, 0, 0
    to_load = {}
    for k, v in sd_peft.items():
        if k not in model_sd:
            missing += 1
            continue
        if model_sd[k].shape != v.shape:
            skipped_shape += 1
            continue
        to_load[k] = v.to(dtype=model_sd[k].dtype)
        matched += 1
    print(f"[warm] matched={matched} skipped_shape={skipped_shape} missing={missing}")
    if matched == 0:
        raise RuntimeError(
            "No LoRA keys matched. Check format conversion & target_modules of wide LoRA."
        )
    missing_keys, _ = student_model.load_state_dict(to_load, strict=False)
    # PEFT 内部の LoRA キーだけ載れば OK、その他 missing は base 側なので無視
    lora_missing = [k for k in missing_keys if "lora_" in k]
    if lora_missing:
        print(f"[warm] WARN: {len(lora_missing)} lora keys not loaded (e.g. {lora_missing[:3]})")


# ---------------------------------------------------------------------------
# save / restore
# ---------------------------------------------------------------------------
def save_lora_state(model, path: Path, name: str) -> None:
    path.mkdir(parents=True, exist_ok=True)
    sd = {k: v.detach().cpu() for k, v in model.state_dict().items() if "lora_" in k}
    save_file(sd, str(path / f"{name}.safetensors"))


# ---------------------------------------------------------------------------
# main
# ---------------------------------------------------------------------------
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dataset", required=True, type=str,
                    help="caption (*.txt) を含むディレクトリ")
    ap.add_argument("--out", required=True, type=str,
                    help="LoRA 出力ディレクトリ")
    ap.add_argument("--warm-lora", default="", type=str,
                    help="Civitai Anima Turbo LoRA path (ComfyUI 形式)。空なら cold-start")
    ap.add_argument("--total-steps", type=int, default=2000)
    ap.add_argument("--batch-size", type=int, default=1)
    ap.add_argument("--teacher-steps", type=int, default=50)
    ap.add_argument("--student-steps", type=int, default=8)
    ap.add_argument("--teacher-cfg", type=float, default=2.0)
    ap.add_argument("--student-cfg", type=float, default=1.0)
    ap.add_argument("--sigma-shift", type=float, default=3.0,
                    help="Anima 公式 workflow は 3.0")
    ap.add_argument("--lora-rank", type=int, default=32)
    ap.add_argument("--lr", type=float, default=1e-4)
    ap.add_argument("--weight-decay", type=float, default=0.01)
    ap.add_argument("--grad-clip", type=float, default=1.0)
    ap.add_argument("--lpips-weight", type=float, default=0.0,
                    help="LPIPS regularization weight。0 で disable (start 推奨)、後で 0.1-0.5 に")
    ap.add_argument("--resolution", type=int, default=1024)
    ap.add_argument("--log-every", type=int, default=10)
    ap.add_argument("--sample-every", type=int, default=500)
    ap.add_argument("--num-workers", type=int, default=2)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--weight-mode", default="uniform", choices=["uniform", "inv_sigma"])
    ap.add_argument("--neg-prompt", default="",
                    help="negative prompt for teacher CFG (空文字 = empty conditioning)")
    args = ap.parse_args()

    torch.manual_seed(args.seed)
    device = torch.device("cuda")
    dtype = torch.bfloat16

    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    # ----- load Anima base -----
    print("[load] Anima bundle (DiT + Qwen3 + WanVAE + LLMAdapter)")
    bundle = build_anima(device=device, dtype=dtype)

    # ----- teacher = base のクローン (frozen, no LoRA) -----
    print("[setup] cloning DiT for teacher (frozen) and student (wide LoRA)")
    teacher_transformer = copy.deepcopy(bundle.transformer).to(device=device, dtype=dtype).eval()
    for p in teacher_transformer.parameters():
        p.requires_grad = False

    # ----- student = base に wide LoRA を attach -----
    student_transformer = bundle.transformer  # bundle 側の方を再利用 (メモリ節約)
    student_transformer = attach_wide_lora(student_transformer, rank=args.lora_rank)
    student_transformer.to(device=device, dtype=dtype)
    # base 凍結、LoRA だけ trainable
    for n, p in student_transformer.named_parameters():
        p.requires_grad = ("lora_" in n)
    trainable = sum(p.numel() for p in student_transformer.parameters() if p.requires_grad)
    total = sum(p.numel() for p in student_transformer.parameters())
    print(f"[setup] student trainable: {trainable/1e6:.1f}M / {total/1e6:.1f}M")
    # bundle.transformer は student と同じ object (LoRA wrap 済)。
    # text_encode/vae_*以外には bundle.transformer 不要なので参照だけ残す。
    bundle.transformer = student_transformer

    # ----- warm-start (任意) -----
    if args.warm_lora:
        load_warm_lora(student_transformer, args.warm_lora)

    # ----- LPIPS (任意) -----
    lpips_fn = None
    if args.lpips_weight > 0:
        import lpips as _lpips
        lpips_fn = _lpips.LPIPS(net="alex").to(device).eval()
        for p in lpips_fn.parameters():
            p.requires_grad = False
        print("[setup] LPIPS(alex) loaded for regularization")

    # ----- schedules -----
    # student schedule: 推論時に使う 8-step grid (t=1 → 0)
    student_sched = make_schedule(args.student_steps, args.sigma_shift, device=device, dtype=torch.float32)
    student_ts = student_sched.timesteps  # (N+1,)
    print(f"[schedule] student t = {student_ts.tolist()}")

    # ----- dataset -----
    print(f"[data] loading {args.dataset}")
    dataset = TextOnlyDataset(args.dataset)
    print(f"   {len(dataset)} captions")
    loader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.num_workers,
        collate_fn=text_collate,
        drop_last=True,
    )

    # ----- optimizer -----
    opt = torch.optim.AdamW(
        [p for p in student_transformer.parameters() if p.requires_grad],
        lr=args.lr, betas=(0.9, 0.999), weight_decay=args.weight_decay, eps=1e-8,
    )

    # ----- training loop -----
    print(f"[train] total_steps={args.total_steps} batch={args.batch_size} "
          f"teacher={args.teacher_steps} student={args.student_steps} "
          f"lpips_weight={args.lpips_weight}")
    log_path = out_dir / "traj_log.jsonl"
    log_f = open(log_path, "a", buffering=1)
    t0 = time.time()
    data_iter = iter(loader)

    # student velocity_fn: バッチ次元 latents + (B,) timesteps + cond -> velocity
    def _student_v(x, t, cond):
        return AnimaBundle.dit_forward(student_transformer, x, t, cond)

    def _teacher_v(x, t, cond):
        return AnimaBundle.dit_forward(teacher_transformer, x, t, cond)

    # negative cond は 1 度だけエンコード
    with torch.no_grad():
        cond_neg = bundle.text_encode([args.neg_prompt or ""])  # (1, 512, 1024)
    # latent サイズは VAE 8x down 想定 (Anima WanVAE)
    H_lat = args.resolution // 8
    W_lat = args.resolution // 8

    for step in range(args.total_steps):
        try:
            captions = next(data_iter)
        except StopIteration:
            data_iter = iter(loader)
            captions = next(data_iter)

        B = len(captions)

        # text encode (positive)
        with torch.no_grad():
            cond_pos = bundle.text_encode(captions)  # (B, 512, 1024)
            # neg は B にブロードキャスト
            cond_neg_b = cond_neg.expand(B, -1, -1).contiguous() if B > 1 else cond_neg

        # 初期 noise (t=1 の状態)
        init_noise = torch.randn(
            B, 16, 1, H_lat, W_lat, device=device, dtype=dtype,
        )

        # ----- teacher rollout (no_grad inside fetch_trajectory) -----
        teacher_sched = make_schedule(args.teacher_steps, args.sigma_shift,
                                      device=device, dtype=torch.float32)
        teacher_sched = snap_to_targets(teacher_sched, student_ts)
        teacher_transformer.eval()
        teacher_traj = fetch_trajectory(
            _teacher_v, init_noise, teacher_sched, student_ts,
            cond_pos, cond_neg_b, cfg_scale=args.teacher_cfg,
        )

        # ----- student align loss -----
        student_transformer.train()
        l_align = align_trajectory(
            _student_v, teacher_traj, student_ts,
            cond_pos, cond_neg_b, cfg_scale_student=args.student_cfg,
            weight_mode=args.weight_mode,
        )

        # ----- LPIPS reg (任意) -----
        if lpips_fn is not None:
            l_reg = compute_regularization(
                _student_v, bundle, init_noise,
                teacher_final_latent=teacher_traj[-1],
                student_timesteps=student_ts,
                cond_pos=cond_pos, cond_neg=cond_neg_b,
                lpips_fn=lpips_fn,
                cfg_scale=args.student_cfg,
            )
        else:
            l_reg = torch.zeros((), device=device)

        loss = l_align + args.lpips_weight * l_reg

        opt.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(
            [p for p in student_transformer.parameters() if p.requires_grad],
            args.grad_clip,
        )
        opt.step()

        # ----- log -----
        if step % args.log_every == 0:
            metrics = {
                "step": step,
                "elapsed": time.time() - t0,
                "loss": float(loss.detach()),
                "l_align": float(l_align.detach()),
                "l_reg": float(l_reg.detach()),
            }
            log_f.write(json.dumps(metrics) + "\n")
            msg = " ".join(f"{k}={v:.4f}" if isinstance(v, float) else f"{k}={v}"
                           for k, v in metrics.items() if k != "step")
            print(f"[step {step}/{args.total_steps}] {msg}", flush=True)

        # ----- checkpoint -----
        if step > 0 and step % args.sample_every == 0:
            save_lora_state(student_transformer, out_dir, f"traj_step{step:05d}")
            print(f"[save] traj_step{step:05d}.safetensors", flush=True)
            try:
                import modal
                modal.Volume.from_name("anima-outputs").commit()
                print(f"[save] volume committed at step {step}", flush=True)
            except Exception as e:
                print(f"[save] volume commit failed: {e}", flush=True)

    # final
    print("[done] saving final LoRA")
    save_lora_state(student_transformer, out_dir, "traj_final")
    log_f.close()


if __name__ == "__main__":
    main()