#!/usr/bin/env python3
"""Evaluation script for ALE Pathology Metastasis Screening task.

Compares agent predictions against reference annotations.
Scoring: Tile F1 (60%) + Detection Accuracy (20%) + Tumor Area Fraction (20%).

Usage:
    python main.py <input_dir> <output_dir> <reference_dir>
"""

import json
import os
import sys


def load_json(path):
    """Load a JSON file and return its contents."""
    with open(path, "r") as f:
        return json.load(f)


def tiles_to_set(tile_list):
    """Convert a list of {x, y} dicts to a set of (x, y) tuples."""
    return {(t["x"], t["y"]) for t in tile_list}


def compute_f1(predicted_set, reference_set):
    """Compute precision, recall, and F1 between two sets of tiles."""
    if not predicted_set and not reference_set:
        return 1.0, 1.0, 1.0  # Both empty = perfect agreement (no tumor)
    if not predicted_set or not reference_set:
        return 0.0, 0.0, 0.0

    true_positives = len(predicted_set & reference_set)
    precision = true_positives / len(predicted_set) if predicted_set else 0.0
    recall = true_positives / len(reference_set) if reference_set else 0.0

    if precision + recall == 0:
        f1 = 0.0
    else:
        f1 = 2 * precision * recall / (precision + recall)

    return precision, recall, f1


def score_tile_f1(f1_value):
    """Convert raw F1 to a 0-1 score with thresholds.

    Full credit at F1 >= 0.80, linear from 0.40-0.80, zero below 0.40.
    """
    if f1_value >= 0.80:
        return 1.0
    elif f1_value >= 0.40:
        return (f1_value - 0.40) / 0.40
    else:
        return 0.0


def score_detection(pred_tumor_present, ref_tumor_present):
    """Score binary detection accuracy. 1.0 if match, 0.0 if not."""
    return 1.0 if bool(pred_tumor_present) == bool(ref_tumor_present) else 0.0


def score_area_fraction(pred_fraction, ref_fraction):
    """Score tumor area fraction by absolute tolerance.

    Full credit within +/- 0.02, partial within +/- 0.05, zero otherwise.
    """
    diff = abs(pred_fraction - ref_fraction)
    if diff <= 0.02:
        return 1.0
    elif diff <= 0.05:
        return 0.5
    else:
        return 0.0


def validate_predictions(predictions, reference, input_dir):
    """Validate prediction schema and consistency.

    Returns a list of warning strings (non-fatal) and a list of error
    strings (fatal — will zero the score).
    """
    warnings = []
    errors = []

    # ---- Required fields ----
    required_fields = [
        "slide_id", "tumor_present", "tumor_tiles",
        "tumor_tile_count", "total_tissue_tiles", "tumor_area_fraction",
    ]
    for field in required_fields:
        if field not in predictions:
            errors.append(f"Missing required field: '{field}'")
    if errors:
        return warnings, errors

    # ---- slide_id match ----
    ref_slide_id = reference.get("slide_id", "")
    pred_slide_id = predictions.get("slide_id", "")
    if pred_slide_id != ref_slide_id:
        warnings.append(
            f"slide_id mismatch: predicted '{pred_slide_id}' "
            f"vs reference '{ref_slide_id}'"
        )

    # ---- Load grid bounds from metadata ----
    metadata_path = os.path.join(input_dir, "slide_metadata.json")
    grid_w, grid_h = None, None
    if os.path.exists(metadata_path):
        metadata = load_json(metadata_path)
        grid_w = metadata.get("grid_width")
        grid_h = metadata.get("grid_height")

    # ---- Validate tumor_tiles ----
    tumor_tiles = predictions.get("tumor_tiles", [])
    if not isinstance(tumor_tiles, list):
        errors.append("tumor_tiles must be a list")
        return warnings, errors

    seen = set()
    for i, tile in enumerate(tumor_tiles):
        if not isinstance(tile, dict) or "x" not in tile or "y" not in tile:
            errors.append(
                f"tumor_tiles[{i}] must be an object with 'x' and 'y' fields"
            )
            continue

        x, y = tile["x"], tile["y"]

        # Integer check
        if not isinstance(x, int) or not isinstance(y, int):
            warnings.append(
                f"tumor_tiles[{i}]: coordinates ({x}, {y}) should be integers"
            )

        # Bounds check
        if grid_w is not None and grid_h is not None:
            if x < 0 or x >= grid_w or y < 0 or y >= grid_h:
                warnings.append(
                    f"tumor_tiles[{i}]: ({x}, {y}) out of grid bounds "
                    f"(0-{grid_w - 1}, 0-{grid_h - 1})"
                )

        # Duplicate check
        key = (int(x), int(y))
        if key in seen:
            warnings.append(
                f"tumor_tiles[{i}]: duplicate tile ({x}, {y})"
            )
        seen.add(key)

    # ---- Count consistency ----
    pred_count = predictions.get("tumor_tile_count", 0)
    actual_count = len(set((t.get("x"), t.get("y")) for t in tumor_tiles
                           if isinstance(t, dict)))
    if pred_count != actual_count:
        warnings.append(
            f"tumor_tile_count ({pred_count}) does not match "
            f"actual unique tumor tiles ({actual_count})"
        )

    # ---- Area fraction consistency ----
    pred_fraction = predictions.get("tumor_area_fraction", 0.0)
    total_tissue = predictions.get("total_tissue_tiles", 0)
    if total_tissue > 0:
        expected_fraction = actual_count / total_tissue
        diff = abs(pred_fraction - expected_fraction)
        if diff > 0.001:
            warnings.append(
                f"tumor_area_fraction ({pred_fraction:.6f}) inconsistent "
                f"with tumor_tile_count/total_tissue_tiles "
                f"({expected_fraction:.6f})"
            )
    elif actual_count > 0:
        warnings.append(
            "total_tissue_tiles is 0 but tumor tiles were predicted"
        )

    return warnings, errors


def evaluate(input_dir, output_dir, reference_dir):
    """Run evaluation and return results dict."""
    results = {
        "scores": {},
        "details": {},
        "validation": {},
        "final_score": 0.0,
    }

    # ---- Load predictions ----
    pred_path = os.path.join(output_dir, "predictions.json")
    if not os.path.exists(pred_path):
        results["details"]["error"] = (
            f"predictions.json not found at {pred_path}"
        )
        results["final_score"] = 0.0
        return results

    try:
        predictions = load_json(pred_path)
    except (json.JSONDecodeError, Exception) as e:
        results["details"]["error"] = (
            f"Failed to parse predictions.json: {e}"
        )
        results["final_score"] = 0.0
        return results

    # ---- Load reference ----
    ref_path = os.path.join(reference_dir, "reference_annotations.json")
    if not os.path.exists(ref_path):
        results["details"]["error"] = (
            f"reference_annotations.json not found at {ref_path}"
        )
        results["final_score"] = 0.0
        return results

    reference = load_json(ref_path)

    # ---- Validate predictions ----
    warnings, errors = validate_predictions(predictions, reference, input_dir)
    results["validation"]["warnings"] = warnings
    results["validation"]["errors"] = errors

    if errors:
        results["details"]["error"] = (
            f"Prediction validation failed: {'; '.join(errors)}"
        )
        results["final_score"] = 0.0
        return results

    # ---- Extract fields ----
    pred_tumor_tiles = tiles_to_set(predictions.get("tumor_tiles", []))
    ref_tumor_tiles = tiles_to_set(reference.get("tumor_tiles", []))
    ref_tissue_tiles = tiles_to_set(reference.get("tissue_tiles", []))

    pred_tumor_present = predictions.get("tumor_present", False)
    ref_tumor_present = reference.get("tumor_present", False)

    pred_area_fraction = predictions.get("tumor_area_fraction", 0.0)
    ref_area_fraction = reference.get("tumor_area_fraction", 0.0)

    # Filter predictions to only tissue tiles (ignore background predictions)
    pred_tumor_tiles_in_tissue = pred_tumor_tiles & ref_tissue_tiles

    # ---- Component 1: Tile-Level F1 (60%) ----
    precision, recall, f1 = compute_f1(
        pred_tumor_tiles_in_tissue, ref_tumor_tiles
    )
    tile_f1_score = score_tile_f1(f1)

    results["scores"]["tile_f1"] = {
        "weight": 0.60,
        "raw_f1": round(f1, 4),
        "precision": round(precision, 4),
        "recall": round(recall, 4),
        "score": round(tile_f1_score, 4),
        "weighted": round(0.60 * tile_f1_score, 4),
    }

    results["details"]["predicted_tumor_tiles"] = len(pred_tumor_tiles)
    results["details"]["predicted_in_tissue"] = len(
        pred_tumor_tiles_in_tissue
    )
    results["details"]["reference_tumor_tiles"] = len(ref_tumor_tiles)
    results["details"]["reference_tissue_tiles"] = len(ref_tissue_tiles)

    # ---- Component 2: Detection Accuracy (20%) ----
    detection_score = score_detection(pred_tumor_present, ref_tumor_present)

    results["scores"]["detection"] = {
        "weight": 0.20,
        "predicted": bool(pred_tumor_present),
        "reference": bool(ref_tumor_present),
        "score": detection_score,
        "weighted": round(0.20 * detection_score, 4),
    }

    # ---- Component 3: Tumor Area Fraction (20%) ----
    area_score = score_area_fraction(pred_area_fraction, ref_area_fraction)

    results["scores"]["area_fraction"] = {
        "weight": 0.20,
        "predicted": round(pred_area_fraction, 4),
        "reference": round(ref_area_fraction, 4),
        "absolute_diff": round(
            abs(pred_area_fraction - ref_area_fraction), 4
        ),
        "score": area_score,
        "weighted": round(0.20 * area_score, 4),
    }

    # ---- Final Score ----
    final_score = (
        0.60 * tile_f1_score
        + 0.20 * detection_score
        + 0.20 * area_score
    )
    results["final_score"] = round(final_score, 4)

    return results


def main():
    if len(sys.argv) != 4:
        print(
            "Usage: python main.py <input_dir> <output_dir> <reference_dir>",
            file=sys.stderr,
        )
        sys.exit(1)

    input_dir = sys.argv[1]
    output_dir = sys.argv[2]
    reference_dir = sys.argv[3]

    results = evaluate(input_dir, output_dir, reference_dir)

    # Print results as JSON
    print(json.dumps(results, indent=2))

    # Print validation info
    if results.get("validation", {}).get("warnings"):
        print("\nVALIDATION WARNINGS:")
        for w in results["validation"]["warnings"]:
            print(f"  ⚠ {w}")

    if results.get("validation", {}).get("errors"):
        print("\nVALIDATION ERRORS (score zeroed):")
        for e in results["validation"]["errors"]:
            print(f"  ✗ {e}")

    # Print summary
    print("\n" + "=" * 60)
    print("EVALUATION SUMMARY")
    print("=" * 60)
    for component, data in results.get("scores", {}).items():
        print(
            f"  {component:20s}: "
            f"score={data['score']:.4f}  "
            f"(weight={data['weight']:.0%}, "
            f"weighted={data['weighted']:.4f})"
        )
    print("-" * 60)
    print(f"  {'FINAL SCORE':20s}: {results['final_score']:.4f}")
    print("=" * 60)

    return results["final_score"]


if __name__ == "__main__":
    main()