# AUTOGENERATED! DO NOT EDIT! File to edit: ../evaluation.ipynb.

# %% auto 0
__all__ = ['evaluate_Fs', 'evaluate_corrs', 'evaluate_Hs']

# %% ../evaluation.ipynb 3
import numpy as np
from .dataset import *
from .metrics import *

def evaluate_Fs(Fs = [], subset = 'test'):
    dset = WxBSDataset('.WxBS', subset=subset, download=True)
    ths = np.arange(20)
    gt_corrs = []
    names = []
    for data_dict in dset:
        corrs = data_dict['pts']
        pairname = data_dict['name']
        names.append(pairname)
        gt_corrs.append(corrs)
    assert len(Fs) == len(gt_corrs)
    per_pair_results = {}
    all_res = []
    for (F, pts, pairname) in zip(Fs, gt_corrs, names):
        res = fraction_of_gt_corrs_consisent_with_F(F, pts, ths)
        per_pair_results[pairname] = res
        all_res.append(res)
    per_pair_results['average'] = np.stack(all_res, axis=1).mean(axis=1)
    return per_pair_results, ths

# %% ../evaluation.ipynb 7
def evaluate_corrs(estimated_right = [],
                   estimated_left = [],
                   subset = 'test'):
    dset = WxBSDataset('.WxBS', subset=subset, download=True)
    ths = np.arange(20)
    gt_corrs = []
    names = []
    for data_dict in dset:
        corrs = data_dict['pts']
        pairname = data_dict['name']
        names.append(pairname)
        gt_corrs.append(corrs)
    assert len(estimated_right) == len(gt_corrs)
    assert len(estimated_left) == len(gt_corrs)
    all_res = []
    per_pair_results = {}
    for (est_right, est_left, gt_pts, pairname) in zip(estimated_right,
                                                       estimated_left, 
                                                       gt_corrs, names):
        res = 0.5 * (PCK(est_right, gt_pts[:, 2:4], ths) + 
                     PCK(est_left,  gt_pts[:, :2], ths))
        per_pair_results[pairname] = res
        all_res.append(res)
    per_pair_results['average'] = np.stack(all_res, axis=1).mean(axis=1)
    return per_pair_results, ths

# %% ../evaluation.ipynb 10
def evaluate_Hs(Hs = []):
    dset = EVDDataset('.EVD', download=True)
    ths = np.logspace(np.log2(1.0), np.log2(20), 10, base=2.0)
    gt_homos = []
    names = []
    shapes1 = []
    shapes2 = []
    
    for data_dict in dset:
        gt_homo = data_dict['H']
        pairname = data_dict['name']
        shapes1.append(data_dict['img1_shape'])
        shapes2.append(data_dict['img2_shape'])
        names.append(pairname)
        gt_homos.append(gt_homo)
    assert len(Hs) == len(gt_homos)
    per_pair_results = {}
    all_res = []
    for (H, Hgt, sh1, sh2, pairname) in zip(Hs, gt_homos, shapes1, shapes2, names):
        mae = get_visible_part_mean_absolute_reprojection_error(sh1, sh2, Hgt, H)
        per_pair_results[pairname] = mae
        all_res.append(mae)
        
    per_pair_results['average'] = get_mAA(all_res, ths)
    return per_pair_results, ths
