import torch, os, json
import numpy as np
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval

CLASS_NAMES = ["Photograph","Illustration","Map","Comics/Cartoon","Editorial Cartoon","Headline","Advertisement"]

def _build_gt(val_ds):
    cats = [{'id': i+1, 'name': nm} for i, nm in enumerate(CLASS_NAMES)]
    images = []
    anns = []
    ann_id = 1
    for ex in val_ds:
        img_id = int(ex['image_id'])
        images.append({'id': img_id, 'width': ex['width'], 'height': ex['height'], 'file_name': str(img_id)})
        objs = ex['objects']
        for obj in objs:
            x, y, w, h = obj['bbox']
            anns.append({
                'id': ann_id,
                'image_id': img_id,
                'category_id': int(obj['category_id']) + 1,
                'bbox': [float(x), float(y), float(w), float(h)],
                'area': float(obj['area']) if 'area' in obj else float(w*h),
                'iscrowd': int(obj['iscrowd']) if 'iscrowd' in obj else 0,
            })
            ann_id += 1
    data = {'images': images, 'annotations': anns, 'categories': cats}
    return data

def evaluate(model, processor, val_ds, val_dl, device):
    gt_data = _build_gt(val_ds)
    # id->(h,w)
    id2size = {int(ex['image_id']): (ex['height'], ex['width']) for ex in val_ds}

    results = []  # COCO result list
    with torch.no_grad():
        for batch in val_dl:
            pv = batch['pixel_values'].to(device)
            pm = batch['pixel_mask'].to(device)
            out = model(pixel_values=pv, pixel_mask=pm)
            target_sizes = torch.tensor([list(id2size[int(i)]) for i in batch['img_id']], dtype=torch.int64)
            post = processor.post_process_object_detection(out, target_sizes=target_sizes, threshold=0.0)
            for bi, img_id in enumerate(batch['img_id']):
                imid = int(img_id)
                boxes = post[bi]['boxes'].cpu()
                scores = post[bi]['scores'].cpu()
                labels = post[bi]['labels'].cpu()
                for score, lab, box in zip(scores, labels, boxes):
                    x0, y0, x1, y1 = box.tolist()
                    w = max(0.0, x1-x0); h = max(0.0, y1-y0)
                    results.append({'image_id': imid,
                                    'category_id': int(lab)+1,
                                    'bbox': [x0, y0, w, h],
                                    'score': float(score)})

    coco_gt = COCO()
    coco_gt.dataset = gt_data
    coco_gt.createIndex()
    if len(results) == 0:
        return {'mAP': 0.0}
    coco_dt = coco_gt.loadRes(results)
    imgids = [im['id'] for im in gt_data['images']]
    ev = COCOeval(coco_gt, coco_dt, iouType='bbox')
    ev.params.imgIds = imgids
    ev.evaluate()
    ev.accumulate()
    ev.summarize()
    # ev.stats: [AP50..AP75, APsmall, APmedium, APlarge, AR1, AR10, AR100, ARsmall, ARmedium, ARlarge]
    ap50_95 = ev.stats[0]
    ap50 = ev.stats[1]
    ap75 = ev.stats[2]
    ar100 = ev.stats[8]
    prec = ev.eval['precision']  # T,R,K,A,M
    class_ap = {}
    for i, clsnamed in enumerate(CLASS_NAMES):
        p = prec[:,:,i,0,-1]
        p = p[p > -1]
        class_ap[clsnamed] = float(p.mean()) if p.size else 0.0
    for k,v in class_ap.items():
        print(f"  AP {k}: {v:.4f}")
    metrics = {'mean_AP_0.50_0.95': float(ap50_95), 'per_class_AP': class_ap,
               'AP_0.50': float(ap50),
               'AP_0.75': float(ap75),
               'AR_max_100': float(ar100),
               'num_predictions': len(results),
               'pred_file': 'predictions.json'}
    return metrics