import os
import argparse
import torch
import wandb
from datetime import datetime

from transformer_maskgit.transformer_maskgit.ctvit import CTViT
from transformers import AutoTokenizer, AutoModel

from Rad_IR.radir import RADIR
from Rad_IR.radir.data.hmclip_mimic_box import HyCoClipMimicBoxDataset_JPG

from scripts.RadIR_Trainer_CXR_CT import RadIRTrainer


# 当前 ablation 标识，改 loss 时同步更新此处
ABLATION_TAG = "no_i2t_ce"


def build_run_name(args):
    if args.run_name:
        return args.run_name
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    return f"CXRCT_stage1_{ABLATION_TAG}_{timestamp}"


def build_results_folder(args, run_name):
    if args.results_folder:
        return args.results_folder
    return os.path.join("./outputs/ablation", run_name)


# --------------------------------------------------------
# Args
# --------------------------------------------------------
def get_args():
    parser = argparse.ArgumentParser("RADIR Training with Trainer")

    parser.add_argument('--resume', type=str, default=None, help="checkpoint 路径，填了就断点续训")
    parser.add_argument('--batch', type=int, default=16)
    parser.add_argument('--num_train_steps', type=int, default=100000)
    parser.add_argument('--warmup_steps', type=int, default=1000)

    parser.add_argument('--lr', type=float, default=5e-5)
    parser.add_argument('--wd', type=float, default=0.0)

    parser.add_argument('--channels', type=int, default=1)
    parser.add_argument('--height', type=int, default=224)
    parser.add_argument('--width', type=int, default=224)

    parser.add_argument('--modal_embedding', action='store_true')

    parser.add_argument(
        '--run_name',
        type=str,
        default=None,
        help='实验名称；默认 CXRCT_stage1_{ablation_tag}_{timestamp}',
    )
    parser.add_argument(
        '--results_folder',
        type=str,
        default=None,
        help='输出目录；默认 ./outputs/ablation/{run_name}',
    )
    parser.add_argument(
        '--wandb_project',
        type=str,
        default='RadIR_Ablation',
        help='wandb project 名称',
    )
    parser.add_argument(
        '--wandb_group',
        type=str,
        default='ablation_no_i2t_ce',
        help='wandb group 名称',
    )

    # Dataset
    parser.add_argument(
        '--dataset_root',
        type=str,
        default="/testessfs10/users/zeyu.zhang/HMClip/RadIR/RadIR/dataset",
    )
    parser.add_argument(
        '--matched_dir',
        type=str,
        default="/testessfs10/users/zeyu.zhang/HMClip/RadIR/RadIR/dataset/processed/matched",
    )

    # Local HF model
    parser.add_argument(
        '--model_path',
        type=str,
        default="/testessfs10/users/zeyu.zhang/HMClip/RadIR/RadIR/hf_models/microsoft/BiomedVLP-CXR-BERT-specialized",
    )

    parser.add_argument('--num_workers', type=int, default=4)

    return parser.parse_args()


# --------------------------------------------------------
# Main
# --------------------------------------------------------
def main():
    args = get_args()
    run_name = build_run_name(args)
    args.results_folder = build_results_folder(args, run_name)
    args.run_name = run_name

    if not torch.cuda.is_available():
        raise RuntimeError("需要 GPU（CUDA）")

    os.makedirs(args.results_folder, exist_ok=True)

    print(f">>> Run name: {run_name}")
    print(f">>> Results folder: {args.results_folder}")

    # --------------------------------------------------------
    # Load HF Text Encoder
    # --------------------------------------------------------
    print(">>> Loading tokenizer & text encoder...")
    tokenizer = AutoTokenizer.from_pretrained(
        args.model_path, trust_remote_code=True, local_files_only=True
    )

    text_model = AutoModel.from_pretrained(
        args.model_path, trust_remote_code=True, local_files_only=True
    )

    # --------------------------------------------------------
    # Load Image Encoder
    # --------------------------------------------------------
    print(">>> Loading CTViT...")
    image_encoder = CTViT(
        dim=768,
        codebook_size=8192,
        image_size=args.height,
        patch_size=16,
        temporal_patch_size=10,
        spatial_depth=8,
        temporal_depth=6,
        cls_depth=4,
        dim_head=32,
        heads=8,
        channels=args.channels,
    )

    # --------------------------------------------------------
    # Build RADIR
    # --------------------------------------------------------
    print(">>> Building RADIR...")
    radir = RADIR(
        tokenizer=tokenizer,
        image_encoder=image_encoder,
        text_encoder=text_model,
        dim_text=768,
        dim_image=512,
        dim_latent=512,
        use_mlm=False,
        use_all_token_embeds=False,
    )

    # --------------------------------------------------------
    # Dataset
    # --------------------------------------------------------
    print(">>> Loading datasets...")
    train_dataset = HyCoClipMimicBoxDataset_JPG(
        split="train",
        tokenizer=tokenizer,
        dataset_root=args.dataset_root,
        cig_matched_dir=args.matched_dir,
        caption_file="train.jsonl",
        transform=None,
    )

    valid_dataset = HyCoClipMimicBoxDataset_JPG(
        split="valid",
        tokenizer=tokenizer,
        dataset_root=args.dataset_root,
        cig_matched_dir=args.matched_dir,
        caption_file="valid.jsonl",
        transform=None,
    )

    if int(os.environ.get("RANK", 0)) == 0:
        wandb.init(
            project=args.wandb_project,
            name=run_name,
            group=args.wandb_group,
            tags=["stage1", "CXRCT", "ablation", ABLATION_TAG],
            dir=args.results_folder,
            config=vars(args),
        )

    # --------------------------------------------------------
    # Trainer
    # --------------------------------------------------------
    print(">>> Building Trainer...")
    trainer = RadIRTrainer(
        RADIR=radir,
        tokenizer=tokenizer,
        train_dataset=train_dataset,
        valid_dataset=valid_dataset,
        num_train_steps=args.num_train_steps,
        warmup_steps=args.warmup_steps,
        batch_size=args.batch,
        lr=args.lr,
        wd=args.wd,
        results_folder=args.results_folder,
        num_workers=args.num_workers,
        modal_embedding=args.modal_embedding,
        debug=False,
    )

    # --------------------------------------------------------
    # Start Training
    # --------------------------------------------------------
    resume_path = None
    if args.resume is not None:
        if args.resume == "auto":
            resume_path = os.path.join(args.results_folder, "radir_last.pt")
        else:
            resume_path = args.resume

    print("✅ Start Training...")
    trainer.train_model(resume_path=resume_path)


if __name__ == "__main__":
    main()