from __future__ import annotations
import argparse
import json
from collections import defaultdict
from pathlib import Path
from typing import Any
import torch
import torch.nn.functional as F
from datasets import load_dataset
from PIL import Image
from tqdm import tqdm
from transformers import CLIPModel, CLIPProcessor


MODELS = [
    ("pretrained", "openai/clip-vit-large-patch14"),
    ("finetuned", "zer0int/CLIP-Regression-ViT-L-14"),
]


def normalize_feature(x: torch.Tensor, eps: float = 1e-12) -> torch.Tensor:
    return x / x.norm(dim=-1, keepdim=True).clamp_min(eps)


def unwrap_feature_output(x: Any, kind: str) -> torch.Tensor:
    """
    Robustly extract CLIP embedding tensors from transformers outputs.
    """
    if isinstance(x, torch.Tensor):
        return x

    preferred_attrs = ["pooler_output", "last_hidden_state"]
    if kind == "image":
        preferred_attrs = ["image_embeds"] + preferred_attrs
    elif kind == "text":
        preferred_attrs = ["text_embeds"] + preferred_attrs

    for attr in preferred_attrs:
        if hasattr(x, attr):
            v = getattr(x, attr)
            if isinstance(v, torch.Tensor):
                if attr == "last_hidden_state":
                    return v[:, 0, :]
                return v

    if isinstance(x, (tuple, list)):
        tensors = [v for v in x if isinstance(v, torch.Tensor)]
        if tensors:
            for t in tensors:
                if t.ndim == 2:
                    return t
            if tensors[0].ndim == 3:
                return tensors[0][:, 0, :]
            return tensors[0]

    raise TypeError(f"Could not unwrap {kind} feature output of type {type(x)}")


def batched(iterable, batch_size: int):
    for start in range(0, len(iterable), batch_size):
        yield start, iterable[start : start + batch_size]


def get_image(x: Any) -> Image.Image:
    """
    datasets.Image usually gives PIL.Image directly.
    Keep fallback for dict/path variants.
    """
    if isinstance(x, Image.Image):
        return x.convert("RGB")
    if isinstance(x, dict):
        if x.get("bytes") is not None:
            import io
            return Image.open(io.BytesIO(x["bytes"])).convert("RGB")
        if x.get("path") is not None:
            return Image.open(x["path"]).convert("RGB")
    if isinstance(x, (str, Path)):
        return Image.open(x).convert("RGB")
    raise TypeError(f"Unsupported image field type: {type(x)}")


@torch.inference_mode()
def evaluate_model(
    ds,
    model_alias: str,
    model_name_or_path: str,
    device: torch.device,
    batch_size: int,
    fp16: bool,
) -> dict[str, Any]:
    print(f"\n[load] {model_alias}: {model_name_or_path}")

    processor = CLIPProcessor.from_pretrained(model_name_or_path)
    model = CLIPModel.from_pretrained(model_name_or_path)
    model = model.eval().to(device)

    if fp16:
        model = model.half()

    totals = defaultdict(int)
    corrects = defaultdict(int)
    margins = defaultdict(list)

    indices = list(range(len(ds)))

    for _, batch_indices in tqdm(list(batched(indices, batch_size)), desc=f"eval {model_alias}"):
        examples = [ds[i] for i in batch_indices]

        images = [get_image(ex["image"]) for ex in examples]
        object_labels = [str(ex["object_label"]) for ex in examples]
        attack_words = [str(ex["attack_word"]) for ex in examples]
        types = [str(ex["type"]) for ex in examples]

        # Two prompts per image: object prompt and attack-word prompt.
        texts = []
        for obj, atk in zip(object_labels, attack_words):
            texts.append(f"a photo of a {obj}")
            texts.append(f"a photo of a {atk}")

        image_inputs = processor(images=images, return_tensors="pt")
        text_inputs = processor(text=texts, return_tensors="pt", padding=True, truncation=True)

        image_inputs = {k: v.to(device) for k, v in image_inputs.items()}
        text_inputs = {k: v.to(device) for k, v in text_inputs.items()}

        if fp16:
            image_inputs = {
                k: (v.half() if torch.is_floating_point(v) else v)
                for k, v in image_inputs.items()
            }

        image_features_raw = model.get_image_features(**image_inputs)
        text_features_raw = model.get_text_features(**text_inputs)

        image_features = unwrap_feature_output(image_features_raw, kind="image")
        text_features = unwrap_feature_output(text_features_raw, kind="text")

        image_features = normalize_feature(image_features.float())
        text_features = normalize_feature(text_features.float())

        text_features = text_features.view(len(examples), 2, -1)
        object_sims = (image_features * text_features[:, 0, :]).sum(dim=-1)
        attack_sims = (image_features * text_features[:, 1, :]).sum(dim=-1)
        batch_margins = object_sims - attack_sims
        batch_preds = batch_margins > 0

        for typ, ok, margin in zip(types, batch_preds.tolist(), batch_margins.tolist()):
            totals[typ] += 1
            corrects[typ] += int(ok)
            margins[typ].append(float(margin))

    total_all = sum(totals.values())
    correct_all = sum(corrects.values())

    results = {
        "model_alias": model_alias,
        "model_name_or_path": model_name_or_path,
        "by_type": {},
        "all": {
            "n": total_all,
            "correct": correct_all,
            "accuracy": correct_all / total_all if total_all else None,
            "mean_margin_object_minus_attack": (
                sum(m for vals in margins.values() for m in vals) / total_all if total_all else None
            ),
        },
    }

    for typ in sorted(totals):
        vals = margins[typ]
        results["by_type"][typ] = {
            "n": totals[typ],
            "correct": corrects[typ],
            "accuracy": corrects[typ] / totals[typ] if totals[typ] else None,
            "mean_margin_object_minus_attack": sum(vals) / len(vals) if vals else None,
            "min_margin": min(vals) if vals else None,
            "max_margin": max(vals) if vals else None,
        }

    del model
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    return results


def print_results(results: list[dict[str, Any]]) -> None:
    print("\n=== ZERO-SHOT RESULTS ===")

    for res in results:
        print(f"\n[{res['model_alias']}] {res['model_name_or_path']}")
        for typ in ["NoRTA", "RTA", "SynthRTA"]:
            item = res["by_type"].get(typ)
            if item is None:
                print(f"  {typ:8s}: missing")
                continue
            print(
                f"  {typ:8s}: "
                f"acc={item['accuracy']:.4f} "
                f"correct={item['correct']}/{item['n']} "
                f"mean_margin={item['mean_margin_object_minus_attack']:+.4f}"
            )

        all_item = res["all"]
        print(
            f"  {'ALL':8s}: "
            f"acc={all_item['accuracy']:.4f} "
            f"correct={all_item['correct']}/{all_item['n']} "
            f"mean_margin={all_item['mean_margin_object_minus_attack']:+.4f}"
        )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="ZS binary-choice eval for local RTA-100-Triplet HF dataset.")
    parser.add_argument("--dataset-path", type=str, default="zer0int/RTA-100-Triplet", help="Local HF dataset repo path or HF repo id.")
    parser.add_argument("--split", default="train")
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument("--fp16", action="store_true")
    parser.add_argument("--output-json", default=None)
    parser.add_argument("--trust-remote-code", action="store_true")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    device = torch.device(args.device)

    print(f"[load_dataset] {args.dataset_path} split={args.split}")
    ds = load_dataset(args.dataset_path, split=args.split, trust_remote_code=args.trust_remote_code)

    print("\n[dataset]")
    print(ds)

    print("[features]")
    print(ds.features)

    results = []
    for alias, model_name_or_path in MODELS:
        res = evaluate_model(
            ds=ds,
            model_alias=alias,
            model_name_or_path=model_name_or_path,
            device=device,
            batch_size=args.batch_size,
            fp16=args.fp16,
        )
        results.append(res)

    print_results(results)

    if args.output_json:
        out = Path(args.output_json)
        out.parent.mkdir(parents=True, exist_ok=True)
        out.write_text(json.dumps(results, ensure_ascii=False, indent=2), encoding="utf-8")
        print(f"\n[wrote] {out}")


if __name__ == "__main__":
    main()