#!/usr/bin/env python3
"""Score PRIMA predictions with explicit annotations and metric definitions."""

from __future__ import annotations

import argparse
import json
import re
import subprocess
import tempfile
from pathlib import Path

import numpy as np
import torch
from pycocotools import mask as mask_utils
from tools.predictions import index_prediction_files, sample_identifier

DEFAULT_TEXT_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
DEFAULT_TEXT_MODEL_REVISION = "1110a243fdf4706b3f48f1d95db1a4f5529b4d41"



def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("prediction_dirs", nargs="+", type=Path)
    parser.add_argument("--gt-json", type=Path, required=True)
    parser.add_argument("--ann-data", type=Path,
                        help="Optional legacy image-size metadata; new predictions store image sizes.")
    parser.add_argument("--output", type=Path)
    parser.add_argument(
        "--subset-mode",
        choices=("available", "intersection"),
        default="available",
        help=(
            "Evaluate each directory's available files, or use the common file "
            "intersection for every directory."
        ),
    )
    parser.add_argument("--with-text", action="store_true")
    parser.add_argument(
        "--with-meteor",
        action="store_true",
        help="Also run the paper's PTB-tokenized COCO METEOR computation.",
    )
    parser.add_argument("--text-model", default=DEFAULT_TEXT_MODEL)
    parser.add_argument(
        "--penalize-missing",
        action=argparse.BooleanOptionalAction,
        default=True,
        help=(
            "Count missing prediction files as empty outputs. The published API "
            "baseline rows used --no-penalize-missing (available outputs only)."
        ),
    )
    return parser.parse_args()


def decode_mask(segmentation, height: int, width: int):
    if isinstance(segmentation, list):
        rle = mask_utils.merge(mask_utils.frPyObjects(segmentation, height, width))
    elif isinstance(segmentation["counts"], list):
        rle = mask_utils.frPyObjects(segmentation, height, width)
    else:
        rle = segmentation
    return mask_utils.decode(rle).astype(bool)


def binary_iou(left: np.ndarray, right: np.ndarray) -> float:
    # Counterfactual irrelevant-image replacement can put a distractor mask at
    # the original image slot with a different native resolution.  Such masks
    # cannot describe the original target and are therefore exact non-matches.
    # Ordinary evaluation masks have matching shapes, so this leaves the paper
    # metric unchanged while making the replacement diagnostic well-defined.
    if left.shape != right.shape:
        return 0.0
    intersection = np.logical_and(left, right).sum()
    union = np.logical_or(left, right).sum()
    return float(intersection / union) if union else 0.0


def greedy_mask_ious(pred_masks, gt_masks):
    matrix = np.zeros((len(pred_masks), len(gt_masks)), dtype=np.float64)
    for pred_index, pred_mask in enumerate(pred_masks):
        for gt_index, gt_mask in enumerate(gt_masks):
            matrix[pred_index, gt_index] = binary_iou(pred_mask, gt_mask)
    paired = []
    while matrix.size and matrix.max() > 0:
        pred_index, gt_index = np.unravel_index(matrix.argmax(), matrix.shape)
        paired.append(float(matrix[pred_index, gt_index]))
        matrix = np.delete(matrix, pred_index, axis=0)
        matrix = np.delete(matrix, gt_index, axis=1)
    return paired


def greedy_mask_miou(pred_masks, gt_masks):
    paired = greedy_mask_ious(pred_masks, gt_masks)
    return float(np.mean(paired)) if paired else 0.0


def extract_caption(saved) -> str:
    output = saved.get("output", [])
    raw = output if isinstance(output, str) else (output[0] if output else "")
    raw = "" if raw is None else str(raw)
    if "ASSISTANT:" in raw:
        raw = raw.split("ASSISTANT:", 1)[1]
    return " ".join(raw.strip().replace("</s>", "").split())


def clean_caption(caption: str) -> str:
    for markup in (
        "<p> ",
        " </p> ",
        "[SEG]",
        "(IMAGE1)",
        "(IMAGE2)",
        "(IMAGE3)",
    ):
        caption = caption.replace(markup, "")
    return re.sub(r"_\d+\b", "", caption).replace("_", " ").strip()


def clean_reference_caption(caption: str) -> str:
    return re.sub(r"_\d+\b", "", caption).replace("_", " ").strip()


def ptb_tokenize(captions):
    """Run the same PTB tokenizer as COCOEvalCap with temp files in /tmp.

    The upstream package writes beside its JAR, which fails on read-only
    installations.  Changing only the temporary-file location leaves the
    tokenizer command and post-processing identical.
    """
    from pycocoevalcap.tokenizer import ptbtokenizer

    image_ids = [key for key, values in captions.items() for _ in values]
    sentences = "\n".join(
        value["caption"].replace("\n", " ")
        for values in captions.values()
        for value in values
    )
    jar_dir = Path(ptbtokenizer.__file__).resolve().parent
    jar_path = jar_dir / ptbtokenizer.STANFORD_CORENLP_3_4_1_JAR
    with tempfile.NamedTemporaryFile(delete=False, dir="/tmp") as handle:
        handle.write(sentences.encode())
        temporary_path = Path(handle.name)
    try:
        command = [
            "java",
            "-cp",
            str(jar_path),
            "edu.stanford.nlp.process.PTBTokenizer",
            "-preserveLines",
            "-lowerCase",
            str(temporary_path),
        ]
        completed = subprocess.run(
            command,
            cwd=jar_dir,
            check=True,
            stdout=subprocess.PIPE,
        )
    finally:
        temporary_path.unlink(missing_ok=True)
    lines = completed.stdout.decode().split("\n")
    tokenized = {}
    for image_id, line in zip(image_ids, lines):
        words = [
            word
            for word in line.rstrip().split(" ")
            if word not in ptbtokenizer.PUNCTUATIONS
        ]
        tokenized.setdefault(image_id, []).append(" ".join(words))
    return tokenized


def coco_meteor(references, predictions) -> float:
    from pycocoevalcap.meteor.meteor import Meteor

    ground_truth = ptb_tokenize(references)
    results = ptb_tokenize(predictions)
    scorer = Meteor()
    score, _ = scorer.compute_score(ground_truth, results)
    return float(score)


def convert_name(name: str) -> str:
    """Normalize grounded labels using the paper's evaluation rules."""
    name = re.sub(r"\d+", "", name)
    name = name.replace("_", " ").strip()
    name = re.sub(r"_\d+\b", "", name)
    possessive = re.search(r"(?:'s\s*|’s\s*)(.*)$", name, flags=re.IGNORECASE)
    if possessive:
        return possessive.group(1).strip()
    name = re.sub(r"\d+", "", name)
    return name.replace("_", "").strip()


def extract_labels(caption: str, num_images: int):
    labels = [[] for _ in range(num_images)]
    pattern = re.compile(
        r"<p>\s*(.*?)\s*</p>((?:\s*\[SEG\]\s*\(IMAGE\d+\))+)",
    )
    for phrase, seg_block in pattern.findall(caption):
        phrase = convert_name(phrase.strip().lower())
        for image_number in re.findall(
            r"\[SEG\]\s*\(IMAGE(\d+)\)", seg_block
        ):
            index = int(image_number) - 1
            if 0 <= index < num_images:
                labels[index].append(phrase)
    return labels


def normalize_pred_masks(saved, num_images: int):
    raw = list(saved.get("pred_masks", []))
    # The model's saved structure has a sentinel at index zero.
    raw = raw[1 : num_images + 1]
    raw.extend([None] * (num_images - len(raw)))
    normalized = []
    for masks in raw:
        if masks is None or len(masks) == 0:
            normalized.append([])
            continue
        # Older cached runs wrap the per-image [N,H,W] tensor in a singleton
        # batch list, whereas newer evaluation saves that tensor directly.
        if isinstance(masks, (list, tuple)) and len(masks) == 1:
            masks = masks[0]
        if torch.is_tensor(masks):
            masks = masks.detach().cpu().numpy()
        else:
            masks = np.stack(
                [
                    mask.detach().cpu().numpy() if torch.is_tensor(mask) else np.asarray(mask)
                    for mask in masks
                ]
            )
        normalized.append([(mask > 0) for mask in masks])
    return normalized


def text_siou(left: str, right: str) -> float:
    left_words, right_words = set(left.split()), set(right.split())
    union = left_words | right_words
    return len(left_words & right_words) / len(union) if union else 0.0


class TextEmbedder:
    def __init__(self, name: str):
        from sentence_transformers import SentenceTransformer

        revision = DEFAULT_TEXT_MODEL_REVISION if name == DEFAULT_TEXT_MODEL else None
        self.model = SentenceTransformer(name, revision=revision)
        self.revision = getattr(self.model[0].auto_model.config, '_commit_hash', None)
        self.cache = {}

    def similarity(self, left: str, right: str) -> float:
        missing = [value for value in (left, right) if value not in self.cache]
        if missing:
            encoded = self.model.encode(missing, convert_to_numpy=True, normalize_embeddings=True)
            self.cache.update(zip(missing, encoded))
        return float(np.dot(self.cache[left], self.cache[right]))


def grounded_matches(gt_masks, gt_labels, pred_masks, pred_labels, embedder):
    ious = np.zeros((len(gt_masks), len(pred_masks)), dtype=np.float64)
    for gt_index, gt_mask in enumerate(gt_masks):
        for pred_index, pred_mask in enumerate(pred_masks):
            ious[gt_index, pred_index] = binary_iou(gt_mask, pred_mask)
    similarities = np.zeros((len(gt_labels), len(pred_labels)), dtype=np.float64)
    sious = np.zeros_like(similarities)
    for gt_index, gt_label in enumerate(gt_labels):
        for pred_index, pred_label in enumerate(pred_labels):
            similarities[gt_index, pred_index] = embedder.similarity(gt_label, pred_label)
            sious[gt_index, pred_index] = text_siou(gt_label, pred_label)
    matches = 0
    similarity_sum = 0.0
    siou_sum = 0.0
    while ious.size:
        gt_index, pred_index = np.unravel_index(ious.argmax(), ious.shape)
        # The current evaluator treats a mask/label count mismatch as an
        # unmatched remainder instead of indexing outside the text matrix.
        if gt_index >= similarities.shape[0] or pred_index >= similarities.shape[1]:
            break
        if ious[gt_index, pred_index] < 0.5 or similarities[gt_index, pred_index] < 0.5:
            break
        matches += 1
        similarity_sum += similarities[gt_index, pred_index]
        siou_sum += sious[gt_index, pred_index]
        ious[gt_index, :] = 0
        ious[:, pred_index] = 0
        similarities[gt_index, :] = 0
        similarities[:, pred_index] = 0
    denom = max(matches, 1)
    return matches, similarity_sum / denom, siou_sum / denom


def text_only_matches(gt_labels, pred_labels, embedder):
    similarities = np.zeros((len(gt_labels), len(pred_labels)), dtype=np.float64)
    sious = np.zeros_like(similarities)
    for gt_index, gt_label in enumerate(gt_labels):
        for pred_index, pred_label in enumerate(pred_labels):
            similarities[gt_index, pred_index] = embedder.similarity(gt_label, pred_label)
            sious[gt_index, pred_index] = text_siou(gt_label, pred_label)
    matches = 0
    similarity_sum = 0.0
    siou_sum = 0.0
    while similarities.size:
        gt_index, pred_index = np.unravel_index(similarities.argmax(), similarities.shape)
        if similarities[gt_index, pred_index] < 0.5:
            break
        matches += 1
        similarity_sum += similarities[gt_index, pred_index]
        siou_sum += sious[gt_index, pred_index]
        similarities[gt_index, :] = 0
        similarities[:, pred_index] = 0
    denom = max(matches, 1)
    return matches, similarity_sum / denom, siou_sum / denom


def image_size(saved, image_index, groundings, image_info=None):
    """Resolve native dimensions without requiring a separate annotation cache."""
    if image_info is not None:
        return int(image_info['height']), int(image_info['width'])
    sizes = saved.get('image_sizes', [])
    if image_index < len(sizes) and sizes[image_index] is not None:
        return tuple(map(int, sizes[image_index]))
    for grounding in groundings:
        segmentation = grounding['segmentation']
        if isinstance(segmentation, dict) and 'size' in segmentation:
            return tuple(map(int, segmentation['size']))
    paths = saved.get('image_paths', [])
    if image_index < len(paths) and Path(paths[image_index]).is_file():
        from PIL import Image
        with Image.open(paths[image_index]) as image:
            width, height = image.size
        return height, width
    return None


def evaluate(
    prediction_dir: Path,
    prediction_files,
    records,
    selected_ids,
    ann_data=None,
    embedder=None,
    penalize_missing=False,
    with_meteor=False,
):
    ann_data = ann_data or {}
    image_mious = []
    all_matched_pair_ious = []
    actual_positives = 0
    true_positives = 0
    image_matched_similarity = 0.0
    image_matched_siou = 0.0
    image_text_similarity = 0.0
    image_text_siou = 0.0
    samples = 0
    annotated_images = 0
    images_with_predictions = 0
    predicted_masks = 0
    load_errors = 0
    meteor_references = {}
    meteor_predictions = {}
    for record in records:
        pair_id = sample_identifier(record)
        if pair_id not in selected_ids:
            continue
        path = prediction_files.get(pair_id)
        if path is None and not penalize_missing:
            continue
        if path is None:
            saved = {"output": ["ASSISTANT: No prediction available. </s>"], "pred_masks": []}
        else:
            try:
                saved = torch.load(path, map_location="cpu")
            except (EOFError, OSError, RuntimeError):
                load_errors += 1
                if not penalize_missing:
                    continue
                saved = {"output": ["ASSISTANT: No prediction available. </s>"], "pred_masks": []}
        num_images = len(record["img_ids"])
        caption = extract_caption(saved)
        labels_by_image = extract_labels(caption, num_images)
        masks_by_image = normalize_pred_masks(saved, num_images)
        samples += 1
        for image_index in range(num_images):
            groundings = record["groundings"].get(f"img_{image_index + 1}", [])
            if not groundings:
                continue
            pred_masks = masks_by_image[image_index]
            pred_labels = labels_by_image[image_index]
            gt_masks = []
            gt_labels = []
            dataset_name = str(record["dataset_name"])
            image_id = str(record["img_ids"][image_index])
            image_info = ann_data.get(dataset_name, {}).get(image_id)
            native_size = image_size(saved, image_index, groundings, image_info)
            if native_size is None:
                if pred_masks:
                    raise ValueError(
                        f'Native image size is missing for {pair_id}, image {image_index + 1}. '
                        'Use predictions with image_sizes, restore the source images, or provide --ann-data.'
                    )
                # An absent prediction scores zero without decoding polygon masks.
                image_mious.append(0.0)
                actual_positives += len(groundings)
                annotated_images += 1
                continue
            native_height, native_width = native_size
            for grounding in groundings:
                segmentation = grounding["segmentation"]
                gt_masks.append(decode_mask(segmentation, native_height, native_width))
                gt_label = str(grounding["name"]).lower()
                gt_labels.append(convert_name(gt_label))
            if len(pred_labels) < len(pred_masks):
                pred_labels.extend([""] * (len(pred_masks) - len(pred_labels)))
            paired_ious = greedy_mask_ious(pred_masks, gt_masks)
            image_mious.append(float(np.mean(paired_ious)) if paired_ious else 0.0)
            all_matched_pair_ious.extend(paired_ious)
            actual_positives += len(gt_labels)
            predicted_masks += len(pred_masks)
            annotated_images += 1
            images_with_predictions += int(bool(pred_masks))
            if with_meteor and pred_masks:
                caption_id = f"{pair_id}_{image_index + 1}"
                meteor_references[caption_id] = [
                    {
                        "caption": clean_reference_caption(
                            str(record["a"])
                        )
                    }
                ]
                meteor_predictions[caption_id] = [
                    {"caption": clean_caption(caption)}
                ]
            if embedder is not None and pred_masks and pred_labels:
                matched, similarity, siou = grounded_matches(
                    gt_masks, gt_labels, pred_masks, pred_labels, embedder
                )
                true_positives += matched
                image_matched_similarity += similarity
                image_matched_siou += siou
                _, similarity, siou = text_only_matches(gt_labels, pred_labels, embedder)
                image_text_similarity += similarity
                image_text_siou += siou
    result = {
        "prediction_dir": str(prediction_dir),
        "samples": samples,
        "load_errors": load_errors,
        "annotated_images": annotated_images,
        "ground_truth_masks": actual_positives,
        "predicted_masks": predicted_masks,
        "miou": float(np.mean(image_mious)) if image_mious else 0.0,
        "matched_pairs": len(all_matched_pair_ious),
        "miou_over_all_matched_pairs": (
            float(np.mean(all_matched_pair_ious)) if all_matched_pair_ious else 0.0
        ),
        "images_with_predictions": images_with_predictions,
        "miou_over_images_with_predictions": (
            float(np.sum(image_mious) / images_with_predictions)
            if images_with_predictions
            else 0.0
        ),
        "images_with_positive_iou": sum(value > 0 for value in image_mious),
        "miou_over_images_with_positive_iou": (
            float(np.mean([value for value in image_mious if value > 0]))
            if any(value > 0 for value in image_mious)
            else 0.0
        ),
    }
    if embedder is not None:
        result.update(
            {
                "recall": true_positives / max(actual_positives, 1),
                "iou_matched_semantic_similarity": image_matched_similarity
                / max(annotated_images, 1),
                "iou_matched_text_siou": image_matched_siou / max(annotated_images, 1),
                "semantic_similarity": image_text_similarity / max(annotated_images, 1),
                "text_siou": image_text_siou / max(annotated_images, 1),
            }
        )
    if with_meteor:
        result["meteor_evaluated_images"] = len(meteor_predictions)
        result["METEOR"] = (
            coco_meteor(meteor_references, meteor_predictions)
            if meteor_predictions
            else 0.0
        )
    return result


def score_predictions(prediction_dirs, records, ann_data=None, *,
                      subset_mode="available", penalize_missing=True, with_text=False,
                      with_meteor=False, text_model=DEFAULT_TEXT_MODEL, ground_truth=None):
    """Score M4SEG predictions for the standalone CLI or post-generation workflow."""
    manifest_ids = {sample_identifier(record) for record in records}
    if len(manifest_ids) != len(records):
        raise ValueError('Evaluation annotations must have unique example identifiers.')
    prediction_files = []
    available = []
    for directory in prediction_dirs:
        files = index_prediction_files(directory)
        prediction_files.append(files)
        available.append(set(files) & manifest_ids)
    if subset_mode == "intersection":
        selected_sets = [set.intersection(*available)] * len(prediction_dirs)
    elif penalize_missing:
        selected_sets = [manifest_ids] * len(prediction_dirs)
    else:
        selected_sets = available
    embedder = TextEmbedder(text_model) if with_text else None
    report = {
        "ground_truth": str(ground_truth),
        "subset_mode": subset_mode,
        "text_model": text_model if with_text else None,
        "text_model_revision": getattr(embedder, 'revision', None),
        "penalize_missing": penalize_missing,
        "metric_definition": {
            "miou": "mean over annotated images of greedy matched-pair mean IoU",
            "missing_or_extra_masks": "unmatched masks are not added to the matched-pair denominator",
            "text_threshold": 0.5,
            "mask_iou_threshold": 0.5,
        },
        "results": [
            evaluate(
                directory,
                files,
                records,
                selected,
                ann_data,
                embedder,
                penalize_missing,
                with_meteor,
            )
            for directory, files, selected in zip(
                prediction_dirs, prediction_files, selected_sets
            )
        ],
    }
    return report


def main():
    args = parse_args()
    import gzip
    opener = gzip.open if str(args.gt_json).endswith('.gz') else open
    with opener(args.gt_json, 'rt', encoding='utf-8') as handle:
        records = json.load(handle)
    ann_data = torch.load(args.ann_data, map_location="cpu", weights_only=False) if args.ann_data else None
    report = score_predictions(
        args.prediction_dirs, records, ann_data,
        subset_mode=args.subset_mode, penalize_missing=args.penalize_missing,
        with_text=args.with_text, with_meteor=args.with_meteor,
        text_model=args.text_model, ground_truth=args.gt_json,
    )
    output = json.dumps(report, indent=2, sort_keys=True)
    print(output)
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(output + "\n")


if __name__ == "__main__":
    main()