from __future__ import annotations

"""
I2T / T2I evaluation with RaTEScore-aligned GT.

Uses valid.jsonl for image-text pairs and maps each sample to val_caption.csv
row indices, then slices val_ratescore.npy into a sub-matrix aligned with the
eval set order (same logic as RadIR evaluate_uncon soft_label path).
"""

import argparse
import json
import os

import numpy as np
import pandas as pd
import torch
from torch.utils.data import DataLoader, Subset
from transformers import AutoModel, AutoTokenizer

from Rad_IR.radir import lorentz as L
from Rad_IR.radir.RadIR import RADIR, l2norm
from Rad_IR.radir.data.hmclip_mimic_box import HyCoClipMimicBoxDataset_JPG, collate_hmclip_box
from retrieval_metric import compute_ndcg_uncon, soft_retrieval_uncon
from transformer_maskgit.transformer_maskgit.ctvit import CTViT

ROOT = os.path.dirname(os.path.abspath(__file__))


# ---------------------------------------------------------------------------
# RaTEScore index alignment (valid.jsonl  ->  val_caption.csv)
# ---------------------------------------------------------------------------

def study_key_from_csv_path(file_path: str) -> str:
    """/p10/p10190940/s51351116.txt -> p10/p10190940/s51351116"""
    return os.path.splitext(file_path.lstrip("/"))[0]


def study_key_from_jsonl_img_path(img_path: str) -> str:
    """dataset/p10/p10190940/s51351116/xxx.jpg -> p10/p10190940/s51351116"""
    p = img_path.replace("\\", "/")
    if p.startswith("dataset/"):
        p = p[len("dataset/") :]
    parts = p.split("/")
    if len(parts) < 3:
        raise ValueError(f"Cannot parse study key from img_path: {img_path}")
    return "/".join(parts[:3])


def build_eval_index_mapping(
    jsonl_path: str,
    csv_path: str,
    dedupe_study: bool = False,
):
    """
    Map each eval row in jsonl to a row index in val_caption.csv / val_ratescore.npy.

    Returns:
        keep_row_indices: indices into jsonl/dataset to evaluate (in jsonl order)
        csv_indices: parallel list of val_caption.csv row indices for RaTEScore slice
        stats: dict with alignment diagnostics
    """
    df = pd.read_csv(csv_path)
    csv_keys = [study_key_from_csv_path(p) for p in df["File Path"]]
    key_to_csv_idx = {k: i for i, k in enumerate(csv_keys)}

    jsonl_keys = []
    with open(jsonl_path, encoding="utf-8") as f:
        for line in f:
            item = json.loads(line.strip())
            jsonl_keys.append(study_key_from_jsonl_img_path(item["img_path"]))

    keep_row_indices = []
    csv_indices = []
    seen_study = set()
    unmatched = []

    for row_i, key in enumerate(jsonl_keys):
        if key not in key_to_csv_idx:
            unmatched.append((row_i, key))
            continue
        if dedupe_study and key in seen_study:
            continue
        seen_study.add(key)
        keep_row_indices.append(row_i)
        csv_indices.append(key_to_csv_idx[key])

    stats = {
        "jsonl_rows": len(jsonl_keys),
        "csv_rows": len(csv_keys),
        "eval_rows": len(keep_row_indices),
        "unique_studies": len(seen_study),
        "unmatched_jsonl": len(unmatched),
        "csv_not_in_jsonl": sum(1 for k in csv_keys if k not in set(jsonl_keys)),
        "dedupe_study": dedupe_study,
    }
    if unmatched:
        stats["unmatched_examples"] = unmatched[:5]

    return keep_row_indices, csv_indices, stats


def load_aligned_ratescore_gt(gt_path: str, csv_indices: list[int]) -> torch.Tensor:
    """Slice full RaTEScore matrix to match eval sample order."""
    gt = np.load(gt_path)
    idx = np.asarray(csv_indices, dtype=np.int64)
    gt_sub = gt[np.ix_(idx, idx)].astype(np.float32)
    return torch.from_numpy(gt_sub / 100.0)


# ---------------------------------------------------------------------------
# Feature extraction (aligned with RadIR.forward is_condition=False global path)
# ---------------------------------------------------------------------------

def preprocess_image_like_forward(image: torch.Tensor) -> torch.Tensor:
    if image.ndim == 3:
        image = image.unsqueeze(0)
    if image.ndim == 4:
        if image.shape[1] != 1:
            image = image.mean(dim=1, keepdim=True)
    elif image.ndim == 5:
        if image.shape[1] != 1:
            image = image.mean(dim=1, keepdim=True)
    else:
        raise ValueError(f"Unexpected image ndim={image.ndim}, shape={tuple(image.shape)}")
    return image


def extract_hyperbolic_features(model, dataloader, device, modal_embedding=False):
    """Extract hyperbolic features (project_img/txt after l2norm global encodings)."""
    all_hyp_img, all_hyp_txt = [], []

    model.eval()
    with torch.no_grad():
        for batch in dataloader:
            imgs = preprocess_image_like_forward(batch["imgs"].to(device, non_blocking=True))

            enc_image = model.visual_transformer(
                imgs,
                return_encoded_tokens=True,
                modal_embedding=modal_embedding,
                modal_indexs=None,
                is_condition=False,
            )
            if enc_image.dim() == 2:
                image_global = l2norm(enc_image)
            elif enc_image.dim() == 3:
                image_global = l2norm(enc_image[:, 0, :])
            else:
                raise ValueError(f"Unexpected visual feature shape: {enc_image.shape}")

            text_embeddings = model.text_transformer(
                input_ids=batch["caption_ids"].to(device, non_blocking=True),
                token_type_ids=batch["token_type_ids"].to(device, non_blocking=True),
                attention_mask=batch["attention_mask"].to(device, non_blocking=True),
            )
            enc_text = text_embeddings[0]
            text_global = enc_text[:, 0, :]

            all_hyp_img.append(model.project_img(image_global).cpu())
            all_hyp_txt.append(model.project_txt(text_global).cpu())

    return torch.cat(all_hyp_img, dim=0), torch.cat(all_hyp_txt, dim=0)


# ---------------------------------------------------------------------------
# Metrics
# ---------------------------------------------------------------------------

def paired_recall(sim_matrix: torch.Tensor, ks):
    """Hard paired recall: query i should retrieve candidate i (diagonal)."""
    ranking = torch.argsort(sim_matrix, dim=1, descending=True)
    gt = torch.arange(sim_matrix.size(0), device=sim_matrix.device)
    results = {}
    for k in ks:
        hits = (ranking[:, :k] == gt.unsqueeze(1)).any(dim=1).float()
        results[k] = hits.mean().item()
    return results


def print_metric_block(title, ks, paired_r, soft_r, ndcg, scale_pct=True):
    print(f"\n--- {title} ---")
    for k in ks:
        sr = soft_r[k] if isinstance(soft_r, dict) else soft_r[ks.index(k)]
        if sr == -1:
            sr_str = "N/A"
        else:
            sr_str = f"{sr * 100:.2f}%" if scale_pct else f"{sr:.4f}"
        nd = ndcg[ks.index(k)]
        nd_str = f"{nd * 100:.2f}" if scale_pct else f"{nd:.4f}"
        print(
            f"@{k:>3d}  Paired_R={paired_r[k] * 100:.2f}%  "
            f"Soft_R={sr_str}  NDCG={nd_str}"
        )


def evaluate_i2t_t2i(sim_i2t, gt_ratescore, ks):
    """Run RadIR-style I2T metrics with RaTEScore GT matrix."""
    sim_t2i = sim_i2t.t()
    gt = gt_ratescore.to(sim_i2t.device)

    paired_i2t = paired_recall(sim_i2t, ks)
    paired_t2i = paired_recall(sim_t2i, ks)

    soft_i2t = soft_retrieval_uncon(sim_i2t, gt, k=ks)
    soft_t2i = soft_retrieval_uncon(sim_t2i, gt, k=ks)
    if not isinstance(soft_i2t, list):
        soft_i2t = [soft_i2t]
        soft_t2i = [soft_t2i]

    ndcg_i2t = compute_ndcg_uncon(sim_i2t, gt, k=ks)
    ndcg_t2i = compute_ndcg_uncon(sim_t2i, gt, k=ks)

    soft_i2t = dict(zip(ks, soft_i2t))
    soft_t2i = dict(zip(ks, soft_t2i))

    return {
        "i2t": {"paired_r": paired_i2t, "soft_r": soft_i2t, "ndcg": ndcg_i2t},
        "t2i": {"paired_r": paired_t2i, "soft_r": soft_t2i, "ndcg": ndcg_t2i},
    }


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="I2T/T2I evaluation with RaTEScore-aligned GT (valid.jsonl subset)"
    )
    parser.add_argument("--device", default="cuda", type=str)
    parser.add_argument(
        "--model_path",
        default=os.path.join(
            ROOT, "hf_models/microsoft/BiomedVLP-CXR-BERT-specialized"
        ),
    )
    parser.add_argument("--radir_ckpt", default=os.path.join(ROOT, "outputs/ablation/CXRCT_stage1_no_cross_text_entail_20260703_153255/radir_best.pt"))
    parser.add_argument("--dataset_root", default=os.path.join(ROOT, "dataset"))
    parser.add_argument("--caption_file", default="valid.jsonl")
    parser.add_argument(
        "--csv_file",
        default=os.path.join(ROOT, "dataset/val_caption.csv"),
        help="Original RadIR caption list (defines val_ratescore row order)",
    )
    parser.add_argument(
        "--gt_path",
        default=os.path.join(ROOT, "dataset/val_ratescore.npy"),
    )
    parser.add_argument("--cig_matched_dir", default="dataset/processed/matched")
    parser.add_argument("--split", default="valid", type=str)
    parser.add_argument("--batch_size", default=16, type=int)
    parser.add_argument("--num_workers", default=8, type=int)
    parser.add_argument("--imsize", default=224, type=int)
    parser.add_argument("--max_words", default=128, type=int)
    parser.add_argument(
        "--dedupe_study",
        action="store_true",
        help="Keep one image per study (~2808). Default: all jsonl rows (~3120).",
    )
    parser.add_argument(
        "--modal_embedding",
        action="store_true",
        help="Match training when --modal_embedding was enabled",
    )
    parser.add_argument(
        "--ks",
        type=int,
        nargs="+",
        default=[1, 5, 10, 20, 50, 100],
    )
    args = parser.parse_args()
    device = args.device

    jsonl_path = os.path.join(args.dataset_root, args.caption_file)
    keep_rows, csv_indices, stats = build_eval_index_mapping(
        jsonl_path, args.csv_file, dedupe_study=args.dedupe_study
    )
    if stats["unmatched_jsonl"]:
        raise RuntimeError(
            f"{stats['unmatched_jsonl']} jsonl rows have no match in {args.csv_file}. "
            f"Examples: {stats.get('unmatched_examples')}"
        )

    print("===== RaTEScore alignment =====")
    for k, v in stats.items():
        if k != "unmatched_examples":
            print(f"  {k}: {v}")

    gt_ratescore = load_aligned_ratescore_gt(args.gt_path, csv_indices)
    print(f"  gt_sub shape: {tuple(gt_ratescore.shape)}")

    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
    )
    image_encoder = CTViT(
        dim=768,
        codebook_size=8192,
        image_size=args.imsize,
        patch_size=16,
        temporal_patch_size=10,
        spatial_depth=8,
        temporal_depth=6,
        cls_depth=4,
        dim_head=32,
        heads=8,
        channels=1,
    )
    model = RADIR(
        tokenizer=tokenizer,
        image_encoder=image_encoder,
        text_encoder=text_model,
        dim_text=768,
        dim_image=768,
        dim_latent=768,
        use_mlm=False,
        use_all_token_embeds=False,
    )
    print(f"\nLoading checkpoint: {args.radir_ckpt}")
    model.load(args.radir_ckpt)
    model.to(device)

    full_dataset = HyCoClipMimicBoxDataset_JPG(
        split=args.split,
        tokenizer=tokenizer,
        dataset_root=args.dataset_root,
        cig_matched_dir=args.cig_matched_dir,
        caption_file=args.caption_file,
        imsize=args.imsize,
        max_words=args.max_words,
        transform=None,
    )
    dataset = Subset(full_dataset, keep_rows)
    dataloader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=args.num_workers,
        pin_memory=True,
        collate_fn=collate_hmclip_box,
    )
    print(f"\nEval samples = {len(dataset)} (from {len(full_dataset)} jsonl rows)")

    hyp_img, hyp_txt = extract_hyperbolic_features(
        model,
        dataloader,
        device,
        modal_embedding=args.modal_embedding,
    )

    gt = gt_ratescore.to(device)
    ks = args.ks

    curv = model.get_curv()
    sim_i2t = -L.pairwise_dist(
        hyp_img.to(device), hyp_txt.to(device), curv
    ).float()

    print(f"\n===== Hyperbolic I2T/T2I (curv={curv.item():.4f}) =====")
    res = evaluate_i2t_t2i(sim_i2t, gt, ks)
    print_metric_block(
        "Image -> Text",
        ks,
        res["i2t"]["paired_r"],
        res["i2t"]["soft_r"],
        res["i2t"]["ndcg"],
    )
    print_metric_block(
        "Text -> Image",
        ks,
        res["t2i"]["paired_r"],
        res["t2i"]["soft_r"],
        res["t2i"]["ndcg"],
    )
    print("\n  Summary (mean I2T+T2I):")
    for i, k in enumerate(ks):
        pr = (res["i2t"]["paired_r"][k] + res["t2i"]["paired_r"][k]) / 2
        nd = (res["i2t"]["ndcg"][i] + res["t2i"]["ndcg"][i]) / 2
        si = res["i2t"]["soft_r"][k]
        st = res["t2i"]["soft_r"][k]
        sr = (si + st) / 2 if si != -1 and st != -1 else -1
        sr_str = f"{sr * 100:.2f}%" if sr != -1 else "N/A"
        print(
            f"  @{k:>3d}  Paired_R={pr * 100:.2f}%  "
            f"Soft_R={sr_str}  NDCG={nd * 100:.2f}"
        )


if __name__ == "__main__":
    main()