# AUTOGENERATED! DO NOT EDIT! File to edit: metrics.ipynb (unless otherwise specified).

__all__ = ['leave_one_out_F_validation', 'leave_one_out_H_validation', 'get_error_stat_string', 'get_big_errors_string']

# Cell
import torch
import kornia as K
import kornia.geometry as KG
import cv2
import numpy as np

def leave_one_out_F_validation(corrs: np.ndarray, error: str = 'symepi'):
    assert error in ['symepi', 'sampson']
    metric = KG.symmetrical_epipolar_distance if error == 'symepi' else KG.sampson_epipolar_distance
    num = len(corrs)
    out = {"min": None,
           "max": None,
           "mean": None,
           "all": [],
           "max_idx": None}
    if num < 9:
        return out
    errors = []
    for leave_idx in range(num):
        corrs_cur = np.concatenate([
            corrs[:leave_idx],
            corrs[leave_idx+1:]
        ], axis=0)
        pt = torch.from_numpy(corrs[leave_idx].reshape(1, 1, 4))
        F, _ = cv2.findFundamentalMat(corrs_cur[:,:2],
                                      corrs_cur[:,2:], cv2.FM_8POINT)
        errors.append(metric(pt[:,:,:2],
                             pt[:,:,2:],
                             torch.from_numpy(F)[None], squared=False).item())
    errors = np.array(errors)
    out["all"] = errors
    out["min"] = errors.min()
    out["mean"] = errors.mean()
    out["max_idx"] = np.argmax(errors)
    out["max"] = errors[out["max_idx"]]
    return out

def leave_one_out_H_validation(corrs: np.ndarray, error: str = 'symtransfer'):
    assert error in ['symtransfer']
    metric = KG.symmetric_transfer_error
    num = len(corrs)
    out = {"min": None,
           "max": None,
           "mean": None,
           "all": [],
           "max_idx": None}
    if num < 5:
        return out
    errors = []
    for leave_idx in range(num):
        corrs_cur = np.concatenate([
            corrs[:leave_idx],
            corrs[leave_idx+1:]
        ], axis=0)
        pt = torch.from_numpy(corrs[leave_idx].reshape(1, 1, 4))
        H, _ = cv2.findHomography(corrs_cur[:,:2],
                                  corrs_cur[:,2:], 0)
        errors.append(metric(pt[:,:,:2],
                             pt[:,:,2:],
                             torch.from_numpy(H)[None], squared=False).item())
    errors = np.array(errors)
    out["all"] = errors
    out["min"] = errors.min()
    out["mean"] = errors.mean()
    out["max_idx"] = np.argmax(errors)
    out["max"] = errors[out["max_idx"]]
    return out

def get_error_stat_string(errors):
    out =  (f'Errors: min = {errors["min"]:.2f}, max={errors["max"]:.2f}, at {errors["max_idx"]}, mean={errors["mean"]:.2f}')
    return out

def get_big_errors_string(errors, th=3.0):
    mask = errors['all'] > th
    out =  (f'Error > {th} px, idxs: {np.arange(len(errors["all"]))[mask]}')
    return out



