"""Evaluate Argus's heads on their benchmarks.

    python evaluate.py imagenet --root VAL_DIR            # VAL_DIR/<wnid>/*.JPEG
    python evaluate.py ade20k --root ADEChallengeData2016
    python evaluate.py nyu --root NYU_TEST_DIR            # rgb/<id>.* and depth/<id>.png
    python evaluate.py coco --images val2017 --annotations instances_val2017.json

Every command takes --model (default phanerozoic/argus), --device, --batch
and --limit N to evaluate the first N items only.
"""

import argparse
import json
import os
import sys

import numpy as np
import torch
from PIL import Image
from transformers import AutoModel

IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".JPEG", ".JPG", ".PNG")


def _batches(items, size):
    for i in range(0, len(items), size):
        yield items[i:i + size]


def _load(path):
    return Image.open(path).convert("RGB")


def evaluate_imagenet(model, args):
    """Top-1 and top-5 accuracy with classify(), on an ImageNet validation
    folder of <wnid>/<image> files."""
    wnid_to_idx = {w: i for i, w in enumerate(model.class_ids)}
    items = sorted(
        (os.path.join(args.root, w, f), wnid_to_idx[w])
        for w in os.listdir(args.root) if w in wnid_to_idx
        for f in os.listdir(os.path.join(args.root, w)) if f.endswith(IMAGE_EXTS)
    )[: args.limit]
    top1 = top5 = 0
    for chunk in _batches(items, args.batch):
        results = model.classify([_load(p) for p, _ in chunk], top_k=5)
        for (_, label), result in zip(chunk, results):
            predicted = [wnid_to_idx[r["class_id"]] for r in result]
            top1 += predicted[0] == label
            top5 += label in predicted
    return {"images": len(items), "top1": top1 / len(items), "top5": top5 / len(items)}


def evaluate_ade20k(model, args):
    """Mean IoU over the 150 classes with segment(), on the ADEChallengeData2016
    validation split, whose annotations store class k as k + 1 and 0 as
    unlabeled."""
    image_dir = os.path.join(args.root, "images", "validation")
    label_dir = os.path.join(args.root, "annotations", "validation")
    names = sorted(f for f in os.listdir(image_dir) if f.endswith(IMAGE_EXTS))[: args.limit]
    n = model.config.num_seg_classes
    intersection = np.zeros(n, dtype=np.int64)
    union = np.zeros(n, dtype=np.int64)
    for chunk in _batches(names, args.batch):
        preds = model.segment([_load(os.path.join(image_dir, f)) for f in chunk])
        for f, pred in zip(chunk, preds):
            label = np.asarray(Image.open(os.path.join(label_dir, os.path.splitext(f)[0] + ".png")), dtype=np.int64) - 1
            pred = pred.cpu().numpy()
            valid = label >= 0
            p, t = pred[valid], label[valid]
            intersection += np.bincount(t[p == t], minlength=n)[:n]
            union += np.bincount(p, minlength=n)[:n] + np.bincount(t, minlength=n)[:n] - np.bincount(t[p == t], minlength=n)[:n]
    present = union > 0
    return {"images": len(names), "miou": float((intersection[present] / union[present]).mean())}


def evaluate_nyu(model, args):
    """Depth error with depth(), on NYU Depth V2 test images and 16-bit depth
    PNGs in millimetres, inside the Eigen crop and up to 10 m."""
    rgb_dir, depth_dir = os.path.join(args.root, "rgb"), os.path.join(args.root, "depth")
    stems = sorted(os.path.splitext(f)[0] for f in os.listdir(depth_dir) if f.endswith(".png"))[: args.limit]
    rgb_files = {os.path.splitext(f)[0]: f for f in os.listdir(rgb_dir)}
    abs_rel, sq_rmse, delta1, count = 0.0, 0.0, 0.0, 0
    for chunk in _batches(stems, args.batch):
        preds = model.depth([_load(os.path.join(rgb_dir, rgb_files[s])) for s in chunk])
        for s, pred in zip(chunk, preds):
            gt = np.asarray(Image.open(os.path.join(depth_dir, s + ".png")), dtype=np.float64) / 1000.0
            pred = pred.cpu().numpy().astype(np.float64)
            crop = np.zeros_like(gt, dtype=bool)
            h, w = gt.shape
            crop[int(0.4081 * h):int(0.9919 * h), int(0.0359 * w):int(0.9640 * w)] = True
            valid = crop & (gt > 1e-3) & (gt <= 10.0)
            g, p = gt[valid], np.clip(pred[valid], 1e-3, 10.0)
            abs_rel += np.mean(np.abs(p - g) / g)
            sq_rmse += np.sqrt(np.mean((p - g) ** 2))
            delta1 += np.mean(np.maximum(p / g, g / p) < 1.25)
            count += 1
    return {"images": count, "abs_rel": abs_rel / count, "rmse": sq_rmse / count, "delta1": delta1 / count}


def evaluate_coco(model, args):
    """Box mAP with detect(), scored by pycocotools on COCO val2017."""
    from pycocotools.coco import COCO
    from pycocotools.cocoeval import COCOeval

    gt = COCO(args.annotations)
    name_to_cat = {c["name"]: c["id"] for c in gt.loadCats(gt.getCatIds())}
    present = set(os.listdir(args.images))
    image_ids = [i for i in sorted(gt.getImgIds()) if gt.imgs[i]["file_name"] in present][: args.limit]
    detections = []
    for chunk in _batches(image_ids, args.batch):
        results = model.detect([_load(os.path.join(args.images, gt.imgs[i]["file_name"])) for i in chunk])
        for image_id, result in zip(chunk, results):
            for d in result:
                x1, y1, x2, y2 = d["box"]
                detections.append({
                    "image_id": image_id,
                    "category_id": name_to_cat[d["class_name"]],
                    "bbox": [x1, y1, x2 - x1, y2 - y1],
                    "score": d["score"],
                })
    evaluator = COCOeval(gt, gt.loadRes(detections), "bbox")
    evaluator.params.imgIds = image_ids
    evaluator.evaluate()
    evaluator.accumulate()
    evaluator.summarize()
    s = evaluator.stats
    return {"images": len(image_ids), "map": s[0], "map50": s[1], "map75": s[2], "map_small": s[3], "map_medium": s[4], "map_large": s[5]}


TASKS = {"imagenet": evaluate_imagenet, "ade20k": evaluate_ade20k, "nyu": evaluate_nyu, "coco": evaluate_coco}


def main(argv=None):
    parser = argparse.ArgumentParser(description="Evaluate Argus's heads on their benchmarks.")
    parser.add_argument("task", choices=sorted(TASKS))
    parser.add_argument("--root", help="dataset root for imagenet, ade20k and nyu")
    parser.add_argument("--images", help="COCO val2017 image directory")
    parser.add_argument("--annotations", help="COCO instances_val2017.json")
    parser.add_argument("--model", default="phanerozoic/argus")
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    parser.add_argument("--batch", type=int, default=8)
    parser.add_argument("--limit", type=int, default=None)
    args = parser.parse_args(argv)

    model = AutoModel.from_pretrained(args.model, trust_remote_code=True).to(args.device).eval()
    result = TASKS[args.task](model, args)
    print(json.dumps({"task": args.task, **result}, indent=2))
    return 0


if __name__ == "__main__":
    sys.exit(main())