"""Evaluate DeepLab with multi-scale and flip test-time augmentation."""

from __future__ import annotations

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

import torch
import torch.nn.functional as F
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 segmentation_utils import load_dataset_class_names
from train_deeplab_semseg import YoloSemanticDataset, build_model, load_compatible_state_dict, compute_miou


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("--base-imgsz", type=int, default=384)
    parser.add_argument("--scales", nargs="+", type=float, default=[1.0])
    parser.add_argument("--flip", action="store_true")
    parser.add_argument("--batch", type=int, default=1)
    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 tta_logits(model, images: torch.Tensor, base_size: int, scales: list[float], flip: bool) -> torch.Tensor:
    original_size = images.shape[-2:]
    logits_sum = None
    count = 0

    for scale in scales:
        size = max(32, int(round(base_size * scale)))
        scaled = F.interpolate(images, size=(size, size), mode="bilinear", align_corners=False)
        logits = model(scaled)["out"]
        logits = F.interpolate(logits, size=original_size, mode="bilinear", align_corners=False)
        logits_sum = logits if logits_sum is None else logits_sum + logits
        count += 1

        if flip:
            flipped = torch.flip(scaled, dims=[-1])
            flip_logits = model(flipped)["out"]
            flip_logits = torch.flip(flip_logits, dims=[-1])
            flip_logits = F.interpolate(flip_logits, size=original_size, mode="bilinear", align_corners=False)
            logits_sum = logits_sum + flip_logits
            count += 1

    return logits_sum / count


@torch.no_grad()
def evaluate_tta(model, loader, device, class_names, base_size: int, scales: list[float], flip: bool) -> tuple[float, dict[str, float]]:
    num_classes_with_bg = len(class_names) + 1
    confusion = torch.zeros((num_classes_with_bg, num_classes_with_bg), dtype=torch.int64, device=device)
    model.eval()

    for images, masks in loader:
        images = images.to(device)
        masks = masks.to(device)
        logits = tta_logits(model, images, base_size, scales, flip)
        preds = logits.argmax(dim=1)
        valid = (masks >= 0) & (masks < num_classes_with_bg)
        indices = masks[valid] * num_classes_with_bg + preds[valid]
        confusion += torch.bincount(indices, minlength=num_classes_with_bg**2).reshape(num_classes_with_bg, num_classes_with_bg)

    ious, miou = compute_miou(confusion.cpu(), class_names)
    return miou, ious


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.base_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_tta(model, loader, device, class_names, args.base_imgsz, args.scales, args.flip)
    print(f"split={args.split}")
    print(f"limit={args.limit if args.limit is not None else 'all'}")
    print(f"base_imgsz={args.base_imgsz}")
    print(f"scales={args.scales}")
    print(f"flip={args.flip}")
    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",
            "base_imgsz": args.base_imgsz,
            "scales": args.scales,
            "flip": args.flip,
            "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()