#!/usr/bin/env python3
"""Train CLARA on MVSA-Multiple using script workflow converted from notebook."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any, Dict

import torch
from transformers import CLIPProcessor, DebertaV2Tokenizer

import sys

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from src.mvsa_multiple_pipeline import (
    CLARAModel,
    DEFAULT_MVSA_MULTIPLE_CONFIG,
    MVSALoader,
    Trainer,
    create_dataloaders,
    resolve_device,
    set_seed,
    summarize_splits,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Train CLARA on MVSA-Multiple")
    parser.add_argument("--data-root", default="data/MVSA-Multiple")
    parser.add_argument("--text-dir", default=None, help="Default: <data-root>/data")
    parser.add_argument("--label-file", default=None, help="Default: <data-root>/labelResultAll.txt")
    parser.add_argument(
        "--architecture",
        choices=["token", "legacy_global"],
        default=None,
        help="token: revised token-level RAVEL; legacy_global: audited global-vector baseline",
    )
    parser.add_argument("--output-dir", default="outputs/mvsa_multiple")
    parser.add_argument("--checkpoint-name", default="clara_mvsa_multiple.pt")
    parser.add_argument(
        "--init-checkpoint",
        default=None,
        help="Optional checkpoint path to initialize model weights before training.",
    )

    parser.add_argument("--batch-size", type=int, default=48)
    parser.add_argument("--max-epochs", type=int, default=50)
    parser.add_argument("--learning-rate", type=float, default=5e-5)
    parser.add_argument("--weight-decay", type=float, default=0.01)
    parser.add_argument("--warmup-ratio", type=float, default=0.16)
    parser.add_argument("--scheduler-type", choices=["linear", "cosine"], default=None)
    parser.add_argument("--early-stopping-patience", type=int, default=12)
    parser.add_argument("--grad-accum-steps", type=int, default=2)

    parser.add_argument("--max-length", type=int, default=128)
    parser.add_argument("--num-workers", type=int, default=6)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--device", default="auto", help="auto|cuda|cpu")

    parser.add_argument("--train-ratio", type=float, default=0.8)
    parser.add_argument("--val-ratio", type=float, default=0.1)
    parser.add_argument(
        "--preprocessing-mode",
        choices=["paper", "strict"],
        default="paper",
        help="paper: paper-style MVSA preprocessing; strict: unanimous + cross-agree filtering",
    )
    parser.add_argument("--disable-paper-exact-counts", action="store_true")
    parser.add_argument("--allow-non-unanimous", action="store_true")
    parser.add_argument("--allow-cross-disagree", action="store_true")

    parser.add_argument("--negative-class-boost", type=float, default=12.0)
    parser.add_argument("--negative-focal-gamma", type=float, default=4.0)
    parser.add_argument("--min-ratio-negative", type=float, default=0.30)
    parser.add_argument("--label-smoothing", type=float, default=0.02)
    parser.add_argument("--consistency-weight", type=float, default=0.10)
    parser.add_argument("--pred-veri-weight", type=float, default=0.30)
    parser.add_argument("--feedback-tau", type=float, default=0.5)
    parser.add_argument("--disagreement-tau", type=float, default=1.0)

    parser.add_argument("--disable-mixup-negative", action="store_true")
    parser.add_argument("--disable-weighted-sampler", action="store_true")
    parser.add_argument("--mixup-alpha", type=float, default=0.4)
    parser.add_argument("--paper-loss-mode", action="store_true")
    parser.add_argument("--feedback-loss-weight", type=float, default=0.5)
    parser.add_argument("--lambda-primary", type=float, default=0.5)
    parser.add_argument("--lambda-unimodal", type=float, default=0.25)

    parser.add_argument("--disable-clip-lora", action="store_true")
    parser.add_argument("--hidden-dim", type=int, default=512)
    parser.add_argument(
        "--unfreeze-epoch",
        type=int,
        default=0,
        help="0 disables non-LoRA backbone unfreezing; set >0 only for full/unfrozen ablations.",
    )
    parser.add_argument("--unfreeze-vision-backbone", action="store_true")
    parser.add_argument("--lora-rank", type=int, default=8)
    parser.add_argument("--lora-alpha", type=int, default=None)
    parser.add_argument("--lora-dropout", type=float, default=None)
    parser.add_argument("--vision-model-id", default=None)
    parser.add_argument("--text-model-id", default=None)
    parser.add_argument(
        "--text-unfreeze-mode",
        choices=["freeze_all", "freeze_lower", "unfreeze_all"],
        default="freeze_all",
    )
    return parser.parse_args()


def build_config(args: argparse.Namespace) -> Dict[str, Any]:
    cfg = dict(DEFAULT_MVSA_MULTIPLE_CONFIG)

    text_dir = args.text_dir or str(Path(args.data_root) / "data")
    label_file = args.label_file or str(Path(args.data_root) / "labelResultAll.txt")

    cfg.update(
        {
            "text_dir": text_dir,
            "label_file": label_file,
            "architecture": args.architecture or cfg.get("architecture", "token"),
            "hidden_dim": args.hidden_dim,
            "batch_size": args.batch_size,
            "max_epochs": args.max_epochs,
            "learning_rate": args.learning_rate,
            "weight_decay": args.weight_decay,
            "warmup_ratio": args.warmup_ratio,
            "scheduler_type": args.scheduler_type or cfg.get("scheduler_type", "linear"),
            "early_stopping_patience": args.early_stopping_patience,
            "grad_accum_steps": args.grad_accum_steps,
            "max_length": args.max_length,
            "num_workers": args.num_workers,
            "seed": args.seed,
            "train_ratio": args.train_ratio,
            "val_ratio": args.val_ratio,
            "preprocessing_mode": args.preprocessing_mode,
            "paper_exact_counts": not args.disable_paper_exact_counts,
            "require_unanimous": not args.allow_non_unanimous,
            "require_cross_agree": not args.allow_cross_disagree,
            "negative_class_boost": args.negative_class_boost,
            "negative_focal_gamma": args.negative_focal_gamma,
            "min_ratio_negative": args.min_ratio_negative,
            "label_smoothing": args.label_smoothing,
            "consistency_weight": args.consistency_weight,
            "pred_veri_weight": args.pred_veri_weight,
            "feedback_tau": args.feedback_tau,
            "disagreement_tau": args.disagreement_tau,
            "use_mixup_negative": not args.disable_mixup_negative,
            "use_weighted_sampler": not args.disable_weighted_sampler,
            "mixup_alpha": args.mixup_alpha,
            "paper_loss_mode": args.paper_loss_mode,
            "feedback_loss_weight": args.feedback_loss_weight,
            "lambda_primary": args.lambda_primary,
            "lambda_unimodal": args.lambda_unimodal,
            "enable_clip_lora": not args.disable_clip_lora,
            "unfreeze_epoch": int(args.unfreeze_epoch),
            "unfreeze_vision_backbone": bool(args.unfreeze_vision_backbone),
            "text_unfreeze_mode": args.text_unfreeze_mode,
        }
    )
    if args.vision_model_id:
        cfg["vision_model_id"] = args.vision_model_id
    if args.text_model_id:
        cfg["text_model_id"] = args.text_model_id
    rank = int(args.lora_rank)
    alpha = int(args.lora_alpha) if args.lora_alpha is not None else int(rank * 2)
    cfg["lora_clip"] = dict(cfg.get("lora_clip", {}))
    cfg["lora_deb"] = dict(cfg.get("lora_deb", {}))
    cfg["lora_clip"]["r"] = rank
    cfg["lora_deb"]["r"] = rank
    cfg["lora_clip"]["alpha"] = alpha
    cfg["lora_deb"]["alpha"] = alpha
    if args.lora_dropout is not None:
        cfg["lora_clip"]["dropout"] = float(args.lora_dropout)
        cfg["lora_deb"]["dropout"] = float(args.lora_dropout)
    return cfg


def main() -> None:
    args = parse_args()
    cfg = build_config(args)

    set_seed(int(cfg["seed"]))
    device = resolve_device(args.device)

    print(f"Device: {device}")
    print(f"Text dir: {cfg['text_dir']}")
    print(f"Label file: {cfg['label_file']}")

    loader = MVSALoader(cfg["text_dir"], cfg["label_file"])
    loader.load(
        preprocessing_mode=str(cfg.get("preprocessing_mode", "paper")),
        require_unanimous=bool(cfg["require_unanimous"]),
        require_cross_agree=bool(cfg["require_cross_agree"]),
        paper_exact_counts=bool(cfg.get("paper_exact_counts", False)),
    )
    train_samples, val_samples, test_samples = loader.split(
        train_ratio=float(cfg["train_ratio"]),
        val_ratio=float(cfg["val_ratio"]),
        seed=int(cfg["seed"]),
        paper_811=bool(str(cfg.get("preprocessing_mode", "paper")).lower() == "paper"),
    )

    if not train_samples or not val_samples or not test_samples:
        raise RuntimeError("Missing train/val/test split after filtering. Please check dataset/filter setup.")

    split_stats = summarize_splits(train_samples, val_samples, test_samples)
    print("Split stats:")
    print(json.dumps(split_stats, indent=2))

    clip_processor = CLIPProcessor.from_pretrained(cfg["vision_model_id"])
    tokenizer = DebertaV2Tokenizer.from_pretrained(cfg["text_model_id"])

    pin_memory = bool(cfg["pin_memory"] and device.type == "cuda")
    train_loader, val_loader, _ = create_dataloaders(
        train_samples=train_samples,
        val_samples=val_samples,
        test_samples=test_samples,
        clip_processor=clip_processor,
        tokenizer=tokenizer,
        batch_size=int(cfg["batch_size"]),
        max_length=int(cfg["max_length"]),
        num_workers=int(cfg["num_workers"]),
        pin_memory=pin_memory,
        persistent_workers=bool(cfg["persistent_workers"]),
        prefetch_factor=int(cfg["prefetch_factor"]),
        use_mixup_negative=bool(cfg["use_mixup_negative"]),
        mixup_alpha=float(cfg["mixup_alpha"]),
        negative_class_boost=float(cfg["negative_class_boost"]),
        min_ratio_negative=float(cfg["min_ratio_negative"]),
        weighted_train_sampler=bool(cfg.get("use_weighted_sampler", True)),
    )

    model = CLARAModel(cfg).to(device)
    if args.init_checkpoint:
        init_path = Path(args.init_checkpoint)
        if not init_path.exists():
            raise FileNotFoundError(f"Init checkpoint not found: {init_path}")
        init_payload = torch.load(str(init_path), map_location="cpu")
        init_state = init_payload.get("model_state", init_payload)
        missing, unexpected = model.load_state_dict(init_state, strict=False)
        print(f"Initialized from checkpoint: {init_path}")
        print(f"- Missing keys: {len(missing)}")
        print(f"- Unexpected keys: {len(unexpected)}")

    stats = model.parameter_stats()
    pct = 100.0 * stats["trainable"] / max(1, stats["total"])
    print(f"Model params: total={stats['total']:,}, trainable={stats['trainable']:,} ({pct:.2f}%)")

    use_bf16 = bool(device.type == "cuda" and torch.cuda.is_bf16_supported())

    trainer = Trainer(
        model=model,
        train_loader=train_loader,
        val_loader=val_loader,
        train_samples=train_samples,
        cfg=cfg,
        device=device,
        output_dir=args.output_dir,
        checkpoint_name=args.checkpoint_name,
        use_bf16=use_bf16,
    )

    result = trainer.train()

    cfg_path = Path(args.output_dir) / "train_config_used.json"
    cfg_path.parent.mkdir(parents=True, exist_ok=True)
    cfg_path.write_text(json.dumps(cfg, indent=2), encoding="utf-8")

    print("Training complete")
    print(f"Best Val F1-Weighted: {result['best_val_f1_weighted']:.4f}")
    print(f"Checkpoint: {result['checkpoint_path']}")
    print(f"History: {result['history_path']}")
    print(f"Config: {cfg_path}")


if __name__ == "__main__":
    main()