"""Evaluate a DeepLab semantic segmentation checkpoint on a YOLO-polygon dataset."""

from __future__ import annotations

import argparse
import csv
import json
import os
import random
import sys
from pathlib import Path

import torch
from torch.utils.data import DataLoader


PROJECT_ROOT = Path(__file__).resolve().parents[1]
SCRIPT_DIR = Path(__file__).resolve().parent
if str(SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPT_DIR))
os.environ.setdefault("TORCH_HOME", str(PROJECT_ROOT / ".torch"))

from train_deeplab_semseg import YoloSemanticDataset, build_model, evaluate, load_compatible_state_dict
from segmentation_utils import load_dataset_class_names


DEFAULT_DATASET = PROJECT_ROOT / "data" / "processed" / "foodseg103_target_yolo"
DEFAULT_CHECKPOINT = (
    PROJECT_ROOT
    / "runs"
    / "foodseg103_target"
    / "deeplabv3_r50_target_weightprobe2_i384"
    / "best.pt"
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--checkpoint", type=Path, default=DEFAULT_CHECKPOINT)
    parser.add_argument("--dataset", type=Path, default=DEFAULT_DATASET)
    parser.add_argument("--split", choices=["train", "val"], default="val")
    parser.add_argument("--imgsz", type=int, default=384)
    parser.add_argument("--batch", type=int, default=8)
    parser.add_argument("--workers", type=int, default=0)
    parser.add_argument("--limit", type=int)
    parser.add_argument("--sample-seed", type=int)
    parser.add_argument("--backbone", choices=["resnet50", "resnet101"], default="resnet50")
    parser.add_argument("--stems-file", type=Path)
    parser.add_argument("--output", type=Path)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    class_names = load_dataset_class_names(args.dataset)
    num_classes_with_bg = len(class_names) + 1
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    dataset = YoloSemanticDataset(args.dataset, args.split, args.imgsz, None, train=False)
    if args.stems_file:
        with args.stems_file.open(newline="", encoding="utf-8") as f:
            stems = [row["stem"] for row in csv.DictReader(f)]
        dataset.image_paths = [args.dataset / "images" / args.split / f"{stem}.jpg" for stem in stems]
    if args.sample_seed is not None:
        rng = random.Random(args.sample_seed)
        rng.shuffle(dataset.image_paths)
    if args.limit is not None:
        dataset.image_paths = dataset.image_paths[: args.limit]
    loader = DataLoader(
        dataset,
        batch_size=args.batch,
        shuffle=False,
        num_workers=args.workers,
        pin_memory=torch.cuda.is_available(),
    )

    model = build_model(num_classes_with_bg, pretrained=False, backbone=args.backbone).to(device)
    load_compatible_state_dict(model, args.checkpoint, device)
    miou, ious = evaluate(model, loader, device, num_classes_with_bg, class_names)

    print(f"split={args.split}")
    print(f"limit={args.limit if args.limit is not None else 'all'}")
    for name, iou in ious.items():
        print(f"{name:12s} IoU={iou:.4f}")
    print(f"core mIoU={miou:.4f}")

    if args.output:
        payload = {
            "split": args.split,
            "limit": args.limit if args.limit is not None else "all",
            "imgsz": args.imgsz,
            "checkpoint": str(args.checkpoint),
            "stems_file": str(args.stems_file) if args.stems_file else None,
            "ious": ious,
            "core_miou": miou,
        }
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(json.dumps(payload, indent=2), encoding="utf-8")
        print(f"Wrote evaluation result: {args.output}")


if __name__ == "__main__":
    main()