import json
from pathlib import Path
from time import time
from tqdm import tqdm
import cv2
import torch

from mycv.paths import COCO_DIR
from mycv.datasets.constants import COCO_CATEGORY_LIST
from mycv.datasets.mycocoeval import myCOCOeval
from mycv.utils.general import ANSI
import mycv.utils.detection as detutils
import mycv.utils.image as imutils


COCO2017_CLS_IDS = [
    1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21,
    22, 23, 24, 25, 27, 28, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44,
    46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65,
    67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90
]

class DefaultEval(torch.utils.data.Dataset):
    def __init__(self, split='val2017', img_size=640, to_square=False):
        if split == 'val2017':
            ann_path = COCO_DIR / 'annotations/instances_val2017.json'
            img_dir = COCO_DIR / 'val2017'
        else:
            raise NotImplementedError()
        json_data = json.load(open(ann_path, 'r'))

        self.json_data = json_data
        self.img_dir   = img_dir
        self.img_infos = json_data['images']
        self.img_size  = img_size
        self.to_square = to_square

    def __len__(self):
        return len(self.img_infos)

    def __getitem__(self, index):
        img_info = self.img_infos[index]
        imname = img_info['file_name']
        impath = self.img_dir / imname
        # read image
        im = imutils.cv2_imread(impath)
        assert im.shape[:2] == (img_info['height'], img_info['width']), f'{img_info}, {im.shape}'
        # resize image
        im, ratio, pad_info = imutils.letterbox(im, tgt_size=self.img_size, longer=True,
                                                  to_square=self.to_square, div=32)
        # to tensor
        im = torch.from_numpy(im).permute(2, 0, 1).float().div_(255.0)
        # information for evaluation
        eval_info = (img_info['id'], ratio, pad_info)
        return im, eval_info

    @staticmethod
    def collate_fn(batch):
        imgs, eval_infos = zip(*batch) # transposed
        return torch.stack(imgs, 0), eval_infos

def get_default_val2017_loader(img_size, batch_size=32, workers=2):
    to_square = (batch_size != 1)
    dataset = DefaultEval('val2017', img_size, to_square=to_square)
    dataloader = torch.utils.data.DataLoader(
        dataset, batch_size=batch_size, shuffle=False, num_workers=workers,
        pin_memory=True, drop_last=False, collate_fn=DefaultEval.collate_fn
    )
    return dataloader


@torch.no_grad()
def evaluate_val2017(
        model: torch.nn.Module,
        iou_thres, conf_thres,
        dataloader=None, img_size=None, batch_size=None, workers=None,
        verbose=True
    ):
    if model.training:
        print(ANSI.warningstr(f'Warning: Model {type(model)} is not set to the training mode!'))
    device = next(model.parameters()).device
    dtype = next(model.parameters()).dtype
    if dataloader is None:
        dataloader = get_default_val2017_loader(img_size, batch_size, workers)
    else:
        assert all(i is None for i in [img_size, batch_size, workers])

    img_id2info = {iminfo['id']: iminfo for iminfo in dataloader.dataset.img_infos}

    if not hasattr(model, 'forward_cxcywhsccc'):
        _msg = ANSI.warningstr(
            'Recommend to implement forward_cxcywhsc(), ' + \
            'input~(nB,3,H,W), output~(nB,nDet,5+nCls), ' + \
            'where the last dim~(cx,cy,w,h,score,cls0,cls1,...). ' + \
            'Using forward() instead...')
        print(_msg)
    _forward = getattr(model, 'forward_cxcywhsccc', model.forward)

    all_predictions = []
    print(f'Evaluating {type(model)}, device={device}, dtype={dtype}')
    print(f'iou_thres={iou_thres}, conf_thres={conf_thres}, img_size={img_size}')
    pbar = tqdm(dataloader)
    for imgs, eval_infos in pbar:
        imgs = imgs.to(dtype=dtype, device=device)
        batch_preds = _forward(imgs)

        for i, img_preds in enumerate(batch_preds):
            img_id, ratio, pad_tl = eval_infos[i]

            img_preds = detutils.default_post_processing_(img_preds, conf_thres, iou_thres, duplicate=True)
            # shift and scale to the original image domain
            img_preds = detutils.recover_bbox_(img_preds, pad_tl[0], pad_tl[1], ratio)
            # clip xyxy            
            _img_info = img_id2info[img_id]
            detutils.clip_x1y1x2y2_(img_preds, (_img_info['height'], _img_info['width']))
            # x1y1x2y2 to x1y1wh
            img_preds[:, 2:4].sub_(img_preds[:,0:2])

            # remove empty predictions
            _hw_min, _ = img_preds[:, 2:4].min(dim=1)
            img_preds = img_preds[_hw_min > 0, :]

            for instance in img_preds:
                instance = instance.tolist()
                x1, y1, w, h, score, cls_idx = instance
                all_predictions.append({
                    'image_id': img_id,
                    'bbox': [round(f,3) for f in [x1, y1, w, h]],
                    'score': round(score, 5),
                    'category_id': COCO_CATEGORY_LIST[round(cls_idx)]['id']
                })

        mem = torch.cuda.max_memory_allocated(device) / 1e9
        msg = f'{device}: {mem:.3g}G'
        pbar.set_description(msg)

    tmpath = 'runs/tmp.json'
    print(f'\nsaving json file to {tmpath}...')
    with open(tmpath, 'w') as f:
            json.dump(all_predictions, f)

    # compute metrics
    val_json = json.load(open(COCO_DIR / 'annotations/instances_val2017.json', 'r'))
    results = my_evaluate_json_bbox(val_json, all_predictions, verbose=verbose)
    summary = results.pop('summary')
    if verbose: # print results
        # print(summary)
        titles = [
            'Images',
            'mAP@.5:.95', 'mAP@.5', 'mAP@.75',
            'mAP@small', 'mAP@medium', 'mAP@large'
        ]
        header = ''.join([f'{ANSI.headerstr(t):>12}' for t in titles])
        print(header)
        metrics = ['ap', 'ap50', 'ap75', 'aps', 'apm', 'apl']
        _msg = ("{:>12}"*7).format(len(val_json['images']), *[results[m] for m in metrics])
        print(_msg)
    return results


def my_evaluate_json_bbox(gt_json, dts_json, verbose=True):
    cocoEval = myCOCOeval(gt_json, dts_json, 'bbox', verbose=verbose)
    cocoEval.evaluate()
    cocoEval.accumulate()
    cocoEval.summarize()
    mAPs = cocoEval.stats[0:6]
    ap, ap50, ap75, aps, apm, apl = [round(f,4) for f in mAPs]
    results = {
        'summary': cocoEval.summary,
        'ap': ap, 'ap50': ap50, 'ap75': ap75,
        'aps': aps, 'apm': apm, 'apl': apl
    }
    return results


@torch.no_grad()
def evaluate_val2017_vanilla(model, verbose=True):
    val_dir = COCO_DIR / 'val2017'
    val_ann_path = COCO_DIR / 'annotations/instances_val2017.json'
    val_json = json.load(open(val_ann_path, 'r'))

    _msg = ANSI.warningstr('Need to implement cocoeval_image() which accepts an image path.')
    assert hasattr(model, 'cocoeval_image'), f'{_msg}'

    # check dataset
    for img_info in val_json['images']:
        impath = val_dir / img_info['file_name']
        assert impath.is_file(), ANSI.errorstr(f'Cannot find {impath}.')

    # evaluation
    all_predictions = []
    pbar = tqdm(val_json['images'])
    for img_info in pbar:
        impath = str(val_dir / img_info['file_name'])
        img_preds = model.cocoeval_image(impath)
        for p in img_preds:
            p['image_id'] = img_info['id']
        all_predictions.extend(img_preds)
        debug = 1
        pbar.set_description(getattr(model, 'msg', impath))

    tmpath = 'runs/tmp.json'
    print(f'\nsaving json file to {tmpath}...')
    with open(tmpath, 'w') as f:
            json.dump(all_predictions, f, indent=1)

    # compute metrics
    stats = coco_evaluate_json_bbox(val_ann_path, tmpath)
    return stats


def coco_evaluate_json_bbox(annFile, resFile):
    from pycocotools.coco import COCO
    from pycocotools.cocoeval import COCOeval

    cocoGt = COCO(annFile)
    cocoDt = cocoGt.loadRes(resFile)
    cocoEval = COCOeval(cocoGt, cocoDt, 'bbox')
    # cocoEval.params.imgIds = imgIds
    cocoEval.evaluate()
    cocoEval.accumulate()
    cocoEval.summarize()
    stats = cocoEval.stats
    return stats


if __name__ == "__main__":
    # coco_evaluate_json_bbox(
    #     COCO_DIR / 'annotations/instances_val2017.json',
    #     'runs/tmp.json'
    # )
    results = my_evaluate_json_bbox(
        str(COCO_DIR / 'annotations/instances_val2017.json'),
        'runs/tmp.json'
    )
    print(results)
    exit()

    dataloader = get_default_val2017_loader(640, 32, workers=0)
    for data in dataloader:
        debug = 1

    # dataset = COCOEvalDataset()
    annpath = str(COCO_DIR / 'annotations/instances_val2017.json')
    # results = evaluate_json_bbox(annpath, 'runs/tmp.json')
    results = coco_evaluate_json_bbox(annpath, 'runs/val/exp4/uv5s_37.3_predictions.json')
    print(results.pop('summary'))
    print(results)

import json
from pathlib import Path
from time import time
from tqdm import tqdm
import cv2
import torch

from mycv.paths import COCO_DIR
from mycv.datasets.constants import COCO_CATEGORY_LIST
from mycv.datasets.mycocoeval import myCOCOeval
from mycv.utils.general import ANSI
import mycv.utils.detection as detutils
import mycv.utils.image as imutils


COCO2017_CLS_IDS = [
    1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21,
    22, 23, 24, 25, 27, 28, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44,
    46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65,
    67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90
]

class DefaultEval(torch.utils.data.Dataset):
    def __init__(self, split='val2017', img_size=640, to_square=False):
        if split == 'val2017':
            ann_path = COCO_DIR / 'annotations/instances_val2017.json'
            img_dir = COCO_DIR / 'val2017'
        else:
            raise NotImplementedError()
        json_data = json.load(open(ann_path, 'r'))

        self.json_data = json_data
        self.img_dir   = img_dir
        self.img_infos = json_data['images']
        self.img_size  = img_size
        self.to_square = to_square

    def __len__(self):
        return len(self.img_infos)

    def __getitem__(self, index):
        img_info = self.img_infos[index]
        imname = img_info['file_name']
        impath = self.img_dir / imname
        # read image
        im = imutils.cv2_imread(impath)
        assert im.shape[:2] == (img_info['height'], img_info['width']), f'{img_info}, {im.shape}'
        # resize image
        im, ratio, pad_info = imutils.letterbox(im, tgt_size=self.img_size, longer=True,
                                                  to_square=self.to_square, div=32)
        # to tensor
        im = torch.from_numpy(im).permute(2, 0, 1).float().div_(255.0)
        # information for evaluation
        eval_info = (img_info['id'], ratio, pad_info)
        return im, eval_info

    @staticmethod
    def collate_fn(batch):
        imgs, eval_infos = zip(*batch) # transposed
        return torch.stack(imgs, 0), eval_infos

def get_default_val2017_loader(img_size, batch_size=32, workers=2):
    to_square = (batch_size != 1)
    dataset = DefaultEval('val2017', img_size, to_square=to_square)
    dataloader = torch.utils.data.DataLoader(
        dataset, batch_size=batch_size, shuffle=False, num_workers=workers,
        pin_memory=True, drop_last=False, collate_fn=DefaultEval.collate_fn
    )
    return dataloader


@torch.no_grad()
def evaluate_val2017(
        model: torch.nn.Module,
        iou_thres, conf_thres,
        dataloader=None, img_size=None, batch_size=None, workers=None,
        verbose=True
    ):
    if model.training:
        print(ANSI.warningstr(f'Warning: Model {type(model)} is not set to the training mode!'))
    device = next(model.parameters()).device
    dtype = next(model.parameters()).dtype
    if dataloader is None:
        dataloader = get_default_val2017_loader(img_size, batch_size, workers)
    else:
        assert all(i is None for i in [img_size, batch_size, workers])

    img_id2info = {iminfo['id']: iminfo for iminfo in dataloader.dataset.img_infos}

    if not hasattr(model, 'forward_cxcywhsccc'):
        _msg = ANSI.warningstr(
            'Recommend to implement forward_cxcywhsc(), ' + \
            'input~(nB,3,H,W), output~(nB,nDet,5+nCls), ' + \
            'where the last dim~(cx,cy,w,h,score,cls0,cls1,...). ' + \
            'Using forward() instead...')
        print(_msg)
    _forward = getattr(model, 'forward_cxcywhsccc', model.forward)

    all_predictions = []
    print(f'Evaluating {type(model)}, device={device}, dtype={dtype}')
    print(f'iou_thres={iou_thres}, conf_thres={conf_thres}, img_size={img_size}')
    pbar = tqdm(dataloader)
    for imgs, eval_infos in pbar:
        imgs = imgs.to(dtype=dtype, device=device)
        batch_preds = _forward(imgs)

        for i, img_preds in enumerate(batch_preds):
            img_id, ratio, pad_tl = eval_infos[i]

            img_preds = detutils.default_post_processing_(img_preds, conf_thres, iou_thres, duplicate=True)
            # shift and scale to the original image domain
            img_preds = detutils.recover_bbox_(img_preds, pad_tl[0], pad_tl[1], ratio)
            # clip xyxy            
            _img_info = img_id2info[img_id]
            detutils.clip_x1y1x2y2_(img_preds, (_img_info['height'], _img_info['width']))
            # x1y1x2y2 to x1y1wh
            img_preds[:, 2:4].sub_(img_preds[:,0:2])

            # remove empty predictions
            _hw_min, _ = img_preds[:, 2:4].min(dim=1)
            img_preds = img_preds[_hw_min > 0, :]

            for instance in img_preds:
                instance = instance.tolist()
                x1, y1, w, h, score, cls_idx = instance
                all_predictions.append({
                    'image_id': img_id,
                    'bbox': [round(f,3) for f in [x1, y1, w, h]],
                    'score': round(score, 5),
                    'category_id': COCO_CATEGORY_LIST[round(cls_idx)]['id']
                })

        mem = torch.cuda.max_memory_allocated(device) / 1e9
        msg = f'{device}: {mem:.3g}G'
        pbar.set_description(msg)

    tmpath = 'runs/tmp.json'
    print(f'\nsaving json file to {tmpath}...')
    with open(tmpath, 'w') as f:
            json.dump(all_predictions, f)

    # compute metrics
    val_json = json.load(open(COCO_DIR / 'annotations/instances_val2017.json', 'r'))
    results = my_evaluate_json_bbox(val_json, all_predictions, verbose=verbose)
    summary = results.pop('summary')
    if verbose: # print results
        # print(summary)
        titles = [
            'Images',
            'mAP@.5:.95', 'mAP@.5', 'mAP@.75',
            'mAP@small', 'mAP@medium', 'mAP@large'
        ]
        header = ''.join([f'{ANSI.headerstr(t):>12}' for t in titles])
        print(header)
        metrics = ['ap', 'ap50', 'ap75', 'aps', 'apm', 'apl']
        _msg = ("{:>12}"*7).format(len(val_json['images']), *[results[m] for m in metrics])
        print(_msg)
    return results


def my_evaluate_json_bbox(gt_json, dts_json, verbose=True):
    cocoEval = myCOCOeval(gt_json, dts_json, 'bbox', verbose=verbose)
    cocoEval.evaluate()
    cocoEval.accumulate()
    cocoEval.summarize()
    mAPs = cocoEval.stats[0:6]
    ap, ap50, ap75, aps, apm, apl = [round(f,4) for f in mAPs]
    results = {
        'summary': cocoEval.summary,
        'ap': ap, 'ap50': ap50, 'ap75': ap75,
        'aps': aps, 'apm': apm, 'apl': apl
    }
    return results


@torch.no_grad()
def evaluate_val2017_vanilla(model, verbose=True):
    val_dir = COCO_DIR / 'val2017'
    val_ann_path = COCO_DIR / 'annotations/instances_val2017.json'
    val_json = json.load(open(val_ann_path, 'r'))

    _msg = ANSI.warningstr('Need to implement cocoeval_image() which accepts an image path.')
    assert hasattr(model, 'cocoeval_image'), f'{_msg}'

    # check dataset
    for img_info in val_json['images']:
        impath = val_dir / img_info['file_name']
        assert impath.is_file(), ANSI.errorstr(f'Cannot find {impath}.')

    # evaluation
    all_predictions = []
    pbar = tqdm(val_json['images'])
    for img_info in pbar:
        impath = str(val_dir / img_info['file_name'])
        img_preds = model.cocoeval_image(impath)
        for p in img_preds:
            p['image_id'] = img_info['id']
        all_predictions.extend(img_preds)
        debug = 1
        pbar.set_description(getattr(model, 'msg', impath))

    tmpath = 'runs/tmp.json'
    print(f'\nsaving json file to {tmpath}...')
    with open(tmpath, 'w') as f:
            json.dump(all_predictions, f, indent=1)

    # compute metrics
    stats = coco_evaluate_json_bbox(val_ann_path, tmpath)
    return stats


def coco_evaluate_json_bbox(annFile, resFile):
    from pycocotools.coco import COCO
    from pycocotools.cocoeval import COCOeval

    cocoGt = COCO(annFile)
    cocoDt = cocoGt.loadRes(resFile)
    cocoEval = COCOeval(cocoGt, cocoDt, 'bbox')
    # cocoEval.params.imgIds = imgIds
    cocoEval.evaluate()
    cocoEval.accumulate()
    cocoEval.summarize()
    stats = cocoEval.stats
    return stats


if __name__ == "__main__":
    # coco_evaluate_json_bbox(
    #     COCO_DIR / 'annotations/instances_val2017.json',
    #     'runs/tmp.json'
    # )
    results = my_evaluate_json_bbox(
        str(COCO_DIR / 'annotations/instances_val2017.json'),
        'runs/tmp.json'
    )
    print(results)
    exit()

    dataloader = get_default_val2017_loader(640, 32, workers=0)
    for data in dataloader:
        debug = 1

    # dataset = COCOEvalDataset()
    annpath = str(COCO_DIR / 'annotations/instances_val2017.json')
    # results = evaluate_json_bbox(annpath, 'runs/tmp.json')
    results = coco_evaluate_json_bbox(annpath, 'runs/val/exp4/uv5s_37.3_predictions.json')
    print(results.pop('summary'))
    print(results)

# read image # resize image # to tensor # information for evaluation # transposed # shift and scale to the original image domain # clip xyxy # x1y1x2y2 to x1y1wh # remove empty predictions # compute metrics # print results # print(summary) # check dataset # evaluation # compute metrics # cocoEval.params.imgIds = imgIds # coco_evaluate_json_bbox( #     COCO_DIR / 'annotations/instances_val2017.json', #     'runs/tmp.json' # ) # dataset = COCOEvalDataset() # results = evaluate_json_bbox(annpath, 'runs/tmp.json')