# -*- coding: utf-8 -*-
"""
Standalone evaluation of an AnomSeer/Qwen2.5-VL model on the Time-RA RATs-Uni test
set, reporting the **same metrics** as Time-RA / ITFormer for direct comparison:

  num_samples, num_valid, num_invalid, valid_rate,
  binary_{accuracy,precision_macro,recall_macro,f1_macro},
  type_{accuracy,precision_macro,recall_macro,f1_macro},
  thought_rouge_l, thought_bleu

The metric computation (`compute_rats40k_metrics`, `_rouge_l`, `_bleu`) is copied
verbatim from /mnt/share01/sqk/ITFormer/inference_rats40k.py so the numbers line up
exactly (15-class macro over labels 0-14, invalid predictions counted as -1, etc.).

Inference is done with vLLM. ``--model`` is a HF model directory:
  * the base Qwen2.5-VL (zero-shot baseline), or
  * a LoRA checkpoint already merged to HF (see tools/merge_lora_ckpt.py).

Usage
-----
    python eval_rats_uni.py \
        --model /mnt/share01/sqk/models/Qwen2.5-VL-3B-Instruct \
        --data  /mnt/share01/sqk/datasets/RATs40K/RATs-Uni-TSImage_Reason.json \
        --out   ./eval_results/rats_uni_zeroshot.json \
        --tp 2
"""

import os
import io
import re
import json
import argparse
from typing import Any, Dict, List, Optional

from PIL import Image

# Reuse the exact training prompt + taxonomy so eval matches training.
from multimodal_data_processing.rats_uni import build_prompt, _resolve_figure_path

# ---------------------------------------------------------------------------
# Metric helpers — copied verbatim from ITFormer/inference_rats40k.py so the
# reported numbers are computed identically (do not "improve" these).
# ---------------------------------------------------------------------------
ACTION_ID_MAP: Dict[int, str] = {
    0: "Normal Sequence", 1: "Point Anomaly", 2: "Periodic Change Anomaly",
    3: "Trend Change Anomaly", 4: "Change Point Anomaly", 5: "Distributional Change Anomaly",
    6: "Amplitude Anomaly", 7: "Pattern Change Anomaly", 8: "Sparse Anomaly",
    9: "Repeated Value Anomaly", 10: "Sudden Flatline Anomaly", 11: "Drift Anomaly",
    12: "Sudden Spike Anomaly", 13: "Continuous Segment Anomaly", 14: "Nonlinear Pattern Anomaly",
}

from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score
try:
    from rouge_score import rouge_scorer as _rs_mod
    _ROUGE = True
except ImportError:
    _ROUGE = False
try:
    from nltk.translate.bleu_score import SmoothingFunction, corpus_bleu as _cb
    _NLTK = True
except ImportError:
    _NLTK = False


def _rouge_l(preds: List[str], refs: List[str]) -> float:
    if not preds or not _ROUGE:
        return 0.0
    scorer = _rs_mod.RougeScorer(["rougeL"], use_stemmer=False)
    return sum(scorer.score(p, r)["rougeL"].fmeasure for p, r in zip(preds, refs)) / len(preds)


def _bleu(preds: List[str], refs: List[str]) -> float:
    if not preds or not _NLTK:
        return 0.0
    try:
        return float(_cb([[r.split()] for r in refs], [p.split() for p in preds],
                         smoothing_function=SmoothingFunction().method1))
    except Exception:
        return 0.0


def compute_rats40k_metrics(results: List[Dict[str, Any]]) -> Dict[str, float]:
    """Compute all evaluation metrics from a list of per-sample result dicts."""
    gt_type, pred_type, gt_thoughts, pred_thoughts = [], [], [], []
    invalid = 0

    for r in results:
        pred_id = r.get("pred_action_id")
        gt_id = r.get("gt_action_id")
        if gt_id is None:
            invalid += 1
            continue
        gt_id = int(gt_id)
        gt_type.append(gt_id)
        if pred_id is None:
            invalid += 1
            pred_type.append(-1)
        else:
            pred_type.append(int(pred_id))
        gt_thoughts.append(str(r.get("gt_thought", "") or ""))
        pred_thoughts.append(str(r.get("pred_thought", "") or ""))

    metrics: Dict[str, Any] = {
        "num_samples": len(results),
        "num_valid": len(results) - invalid,
        "num_invalid": invalid,
        "valid_rate": (len(results) - invalid) / len(results) if results else 0.0,
    }

    if not gt_type:
        for k in ("binary_accuracy", "binary_precision_macro", "binary_recall_macro",
                  "binary_f1_macro", "type_accuracy", "type_precision_macro",
                  "type_recall_macro", "type_f1_macro", "thought_rouge_l", "thought_bleu"):
            metrics[k] = 0.0
        return metrics

    gt_bin = [0 if x == 0 else 1 for x in gt_type]
    pred_bin = [0 if x == 0 else (1 if x > 0 else -1) for x in pred_type]

    metrics.update({
        "binary_accuracy":        accuracy_score(gt_bin, pred_bin),
        "binary_precision_macro": precision_score(gt_bin, pred_bin, labels=[0, 1], average="macro", zero_division=0),
        "binary_recall_macro":    recall_score(gt_bin, pred_bin, labels=[0, 1], average="macro", zero_division=0),
        "binary_f1_macro":        f1_score(gt_bin, pred_bin, labels=[0, 1], average="macro", zero_division=0),
        "type_accuracy":          accuracy_score(gt_type, pred_type),
        "type_precision_macro":   precision_score(gt_type, pred_type, labels=list(ACTION_ID_MAP), average="macro", zero_division=0),
        "type_recall_macro":      recall_score(gt_type, pred_type, labels=list(ACTION_ID_MAP), average="macro", zero_division=0),
        "type_f1_macro":          f1_score(gt_type, pred_type, labels=list(ACTION_ID_MAP), average="macro", zero_division=0),
    })
    metrics["thought_rouge_l"] = _rouge_l(pred_thoughts, gt_thoughts)
    metrics["thought_bleu"] = _bleu(pred_thoughts, gt_thoughts)
    return metrics


def print_metrics(metrics: Dict[str, Any]) -> None:
    print("\n" + "=" * 50)
    print("RATs-Uni Evaluation Results")
    print("=" * 50)
    for k, v in metrics.items():
        print(f"  {k:<28} {v:.4f}" if isinstance(v, float) else f"  {k:<28} {v}")


# ---------------------------------------------------------------------------
# Response parsing (AnomSeer output format: <think>...</think> <class>...</class>)
# ---------------------------------------------------------------------------
_ID_BY_KEY = {re.sub(r"[^a-z0-9]+", "", n.lower()): i for i, n in ACTION_ID_MAP.items()}
_ALIASES = {
    "normal": 0, "point": 1, "periodic": 2, "trend": 3, "changepoint": 4,
    "distributional": 5, "amplitude": 6, "pattern": 7, "sparse": 8, "repeated": 9,
    "flatline": 10, "drift": 11, "spike": 12, "continuous": 13, "nonlinear": 14,
}


def _norm(t: str) -> str:
    return re.sub(r"[^a-z0-9]+", "", str(t).lower())


def parse_pred_action_id(text: str) -> Optional[int]:
    """Parse <class>...</class> -> action id; falls back to id/alias/substring."""
    m = re.search(r"<class>(.*?)</class>", text, re.DOTALL | re.IGNORECASE)
    cand = m.group(1).strip() if m else text
    # exact integer id when the field is essentially just a number
    mid = re.search(r"\b(\d{1,2})\b", cand)
    if mid and re.fullmatch(r"[^a-zA-Z]*\d{1,2}[^a-zA-Z]*", cand):
        if int(mid.group(1)) in ACTION_ID_MAP:
            return int(mid.group(1))
    key = _norm(cand)
    if key in _ID_BY_KEY:
        return _ID_BY_KEY[key]
    if key in _ALIASES:
        return _ALIASES[key]
    best = None
    for i, name in ACTION_ID_MAP.items():
        if _norm(name) and _norm(name) in key:
            if best is None or len(name) > len(ACTION_ID_MAP[best]):
                best = i
    if best is not None:
        return best
    if mid and int(mid.group(1)) in ACTION_ID_MAP:
        return int(mid.group(1))
    return None


def parse_pred_thought(text: str) -> str:
    """Extract the <think>...</think> reasoning (fallback: text minus the <class> tag)."""
    m = re.search(r"<think>(.*?)</think>", text, re.DOTALL | re.IGNORECASE)
    if m:
        return m.group(1).strip()
    return re.sub(r"<class>.*?</class>", "", text, flags=re.DOTALL | re.IGNORECASE).strip()


# ---------------------------------------------------------------------------
# Data loading
# ---------------------------------------------------------------------------
def load_test_samples(json_path: str, split: str, data_root: str,
                      max_samples: Optional[int]) -> List[dict]:
    with open(json_path) as f:
        data = json.load(f)
    records = data[split]
    keys = sorted(records.keys(), key=lambda k: int(k) if str(k).isdigit() else k)
    samples = []
    for key in keys:
        if max_samples is not None and len(samples) >= max_samples:
            break
        e = records[key]
        if e is None:
            continue
        img_path = _resolve_figure_path(data_root, e.get("FigurePath", ""))
        if not os.path.isfile(img_path):
            continue
        samples.append({
            "index": key,
            "length": len(e.get("Observation") or []),
            "source": e.get("Source", "unknown"),
            "image_path": img_path,
            "gt_action_id": e.get("ActionID"),
            "gt_thought": e.get("Thought", "") or "",
        })
    return samples


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main() -> None:
    ap = argparse.ArgumentParser(description="Eval an AnomSeer/Qwen2.5-VL model on RATs-Uni (Time-RA metrics).")
    ap.add_argument("--model", required=True, help="HF model dir (base, or a merged-LoRA model).")
    ap.add_argument("--data", default="/mnt/share01/sqk/datasets/RATs40K/RATs-Uni-TSImage_Reason.json")
    ap.add_argument("--data_root", default=None, help="Figure root (defaults to the JSON's dir).")
    ap.add_argument("--split", default="TSAD_test")
    ap.add_argument("--out", default="./eval_results/rats_uni_metrics.json")
    ap.add_argument("--tp", type=int, default=2, help="vLLM tensor-parallel size.")
    ap.add_argument("--gpu_mem_util", type=float, default=0.85)
    ap.add_argument("--max_model_len", type=int, default=2048)
    ap.add_argument("--max_tokens", type=int, default=512)
    ap.add_argument("--max_samples", type=int, default=None)
    ap.add_argument("--save_predictions", action="store_true", help="Also dump per-sample predictions.")
    args = ap.parse_args()

    data_root = args.data_root or os.path.dirname(os.path.abspath(args.data))
    samples = load_test_samples(args.data, args.split, data_root, args.max_samples)
    print(f"Loaded {len(samples)} test samples from {args.split}")

    from vllm import LLM, SamplingParams
    from transformers import AutoProcessor

    processor = AutoProcessor.from_pretrained(args.model, trust_remote_code=True)
    llm = LLM(
        model=args.model,
        tensor_parallel_size=args.tp,
        gpu_memory_utilization=args.gpu_mem_util,
        max_model_len=args.max_model_len,
        limit_mm_per_prompt={"image": 1},
        seed=0,
        trust_remote_code=True,
    )
    sampling = SamplingParams(temperature=0.0, max_tokens=args.max_tokens)  # greedy

    # Build vLLM inputs (chat template inserts the vision tokens; image via multi_modal_data).
    vllm_inputs, metas = [], []
    for s in samples:
        text = build_prompt(s["length"], s["source"])
        text = re.sub(r"^<image>\s*", "", text)  # the chat template handles the image placeholder
        messages = [{"role": "user", "content": [
            {"type": "image"},
            {"type": "text", "text": text},
        ]}]
        prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        try:
            image = Image.open(s["image_path"]).convert("RGB")
        except Exception as exc:  # noqa: BLE001
            print(f"[WARN] bad image {s['image_path']}: {exc}")
            continue
        vllm_inputs.append({"prompt": prompt, "multi_modal_data": {"image": image}})
        metas.append(s)

    print(f"Generating for {len(vllm_inputs)} prompts (greedy)...")
    outputs = llm.generate(vllm_inputs, sampling)

    results, predictions = [], []
    for s, out in zip(metas, outputs):
        resp = out.outputs[0].text
        pred_id = parse_pred_action_id(resp)
        pred_thought = parse_pred_thought(resp)
        results.append({
            "gt_action_id": s["gt_action_id"],
            "pred_action_id": pred_id,
            "gt_thought": s["gt_thought"],
            "pred_thought": pred_thought,
        })
        if args.save_predictions:
            predictions.append({**{k: s[k] for k in ("index", "source", "gt_action_id")},
                                "pred_action_id": pred_id, "response": resp})

    metrics = compute_rats40k_metrics(results)
    print_metrics(metrics)

    os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True)
    with open(args.out, "w") as f:
        json.dump(metrics, f, indent=2)
    print(f"\nSaved metrics -> {args.out}")
    if args.save_predictions:
        pred_path = os.path.splitext(args.out)[0] + "_predictions.json"
        with open(pred_path, "w") as f:
            json.dump(predictions, f, indent=2, ensure_ascii=False)
        print(f"Saved predictions -> {pred_path}")


if __name__ == "__main__":
    main()