import numpy as np
from . import geometry


def _process_inputs(estimate, target, mask):
    if estimate.shape != target.shape:
        raise Exception('estimate and target have to be same shape')
    if mask is None:
        mask = np.ones(estimate.shape, dtype=np.bool)
    else:
        mask = mask != 0
    if estimate.shape != mask.shape:
        raise Exception('estimate and mask have to be same shape')
    return estimate, target, mask


def mse(estimate, target, mask=None):
    estimate, target, mask = _process_inputs(estimate, target, mask)
    m = np.sum((estimate[mask] - target[mask]) ** 2) / mask.sum()
    return m


def rmse(estimate, target, mask=None):
    return np.sqrt(mse(estimate, target, mask))


def mae(estimate, target, mask=None):
    estimate, target, mask = _process_inputs(estimate, target, mask)
    m = np.abs(estimate[mask] - target[mask]).sum() / mask.sum()
    return m


def outlier_fraction(estimate, target, mask=None, threshold=0):
    estimate, target, mask = _process_inputs(estimate, target, mask)
    diff = np.abs(estimate[mask] - target[mask])
    m = (diff > threshold).sum() / mask.sum()
    return m


class Metric(object):
    def __init__(self, str_prefix=''):
        self.str_prefix = str_prefix
        self.reset()

    def reset(self):
        pass

    def add(self, es, ta, ma=None):
        pass

    def get(self):
        return {}

    def items(self):
        return self.get().items()

    def __str__(self):
        return ', '.join([f'{self.str_prefix}{key}={value:.5f}' for key, value in self.get().items()])


class MultipleMetric(Metric):
    def __init__(self, *metrics, **kwargs):
        self.metrics = [*metrics]
        super().__init__(**kwargs)

    def reset(self):
        for m in self.metrics:
            m.reset()

    def add(self, es, ta, ma=None):
        for m in self.metrics:
            m.add(es, ta, ma)

    def get(self):
        ret = {}
        for m in self.metrics:
            vals = m.get()
            for k in vals:
                ret[k] = vals[k]
        return ret

    def __str__(self):
        return '\n'.join([str(m) for m in self.metrics])


class BaseDistanceMetric(Metric):
    def __init__(self, name='', **kwargs):
        super().__init__(**kwargs)
        self.name = name

    def reset(self):
        self.dists = []

    def add(self, es, ta, ma=None):
        pass

    def get(self):
        dists = np.hstack(self.dists)
        return {
            f'dist{self.name}_mean': float(np.mean(dists)),
            f'dist{self.name}_std': float(np.std(dists)),
            f'dist{self.name}_median': float(np.median(dists)),
            f'dist{self.name}_q10': float(np.percentile(dists, 10)),
            f'dist{self.name}_q90': float(np.percentile(dists, 90)),
            f'dist{self.name}_min': float(np.min(dists)),
            f'dist{self.name}_max': float(np.max(dists)),
        }


class DistanceMetric(BaseDistanceMetric):
    def __init__(self, vec_length, p=2, **kwargs):
        super().__init__(name=f'{p}', **kwargs)
        self.vec_length = vec_length
        self.p = p

    def add(self, es, ta, ma=None):
        if es.shape != ta.shape or es.shape[1] != self.vec_length or es.ndim != 2:
            print(es.shape, ta.shape)
            raise Exception('es and ta have to be of shape Nxdim')
        if ma is not None:
            es = es[ma != 0]
            ta = ta[ma != 0]
        dist = np.linalg.norm(es - ta, ord=self.p, axis=1)
        self.dists.append(dist)


class OutlierFractionMetric(DistanceMetric):
    def __init__(self, thresholds, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.thresholds = thresholds

    def get(self):
        dists = np.hstack(self.dists)
        ret = {}
        for t in self.thresholds:
            ma = dists > t
            ret[f'of{t}'] = float(ma.sum() / ma.size)
        return ret


class RelativeDistanceMetric(BaseDistanceMetric):
    def __init__(self, vec_length, p=2, **kwargs):
        super().__init__(name=f'rel{p}', **kwargs)
        self.vec_length = vec_length
        self.p = p

    def add(self, es, ta, ma=None):
        if es.shape != ta.shape or es.shape[1] != self.vec_length or es.ndim != 2:
            raise Exception('es and ta have to be of shape Nxdim')
        dist = np.linalg.norm(es - ta, ord=self.p, axis=1)
        denom = np.linalg.norm(ta, ord=self.p, axis=1)
        dist /= denom
        if ma is not None:
            dist = dist[ma != 0]
        self.dists.append(dist)


class RotmDistanceMetric(BaseDistanceMetric):
    def __init__(self, type='identity', **kwargs):
        super().__init__(name=type, **kwargs)
        self.type = type

    def add(self, es, ta, ma=None):
        if es.shape != ta.shape or es.shape[1] != 3 or es.shape[2] != 3 or es.ndim != 3:
            print(es.shape, ta.shape)
            raise Exception('es and ta have to be of shape Nx3x3')
        if ma is not None:
            raise Exception('mask is not implemented')
        if self.type == 'identity':
            self.dists.append(geometry.rotm_distance_identity(es, ta))
        elif self.type == 'geodesic':
            self.dists.append(geometry.rotm_distance_geodesic_unit_sphere(es, ta))
        else:
            raise Exception('invalid distance type')


class QuaternionDistanceMetric(BaseDistanceMetric):
    def __init__(self, type='angle', **kwargs):
        super().__init__(name=type, **kwargs)
        self.type = type

    def add(self, es, ta, ma=None):
        if es.shape != ta.shape or es.shape[1] != 4 or es.ndim != 2:
            print(es.shape, ta.shape)
            raise Exception('es and ta have to be of shape Nx4')
        if ma is not None:
            raise Exception('mask is not implemented')
        if self.type == 'angle':
            self.dists.append(geometry.quat_distance_angle(es, ta))
        elif self.type == 'mineucl':
            self.dists.append(geometry.quat_distance_mineucl(es, ta))
        elif self.type == 'normdiff':
            self.dists.append(geometry.quat_distance_normdiff(es, ta))
        else:
            raise Exception('invalid distance type')


class BinaryAccuracyMetric(Metric):
    def __init__(self, thresholds=np.linspace(0.0, 1.0, num=101, dtype=np.float64)[:-1], **kwargs):
        self.thresholds = thresholds
        super().__init__(**kwargs)

    def reset(self):
        self.tps = [0 for wp in self.thresholds]
        self.fps = [0 for wp in self.thresholds]
        self.fns = [0 for wp in self.thresholds]
        self.tns = [0 for wp in self.thresholds]
        self.n_pos = 0
        self.n_neg = 0

    def add(self, es, ta, ma=None):
        if ma is not None:
            raise Exception('mask is not implemented')
        es = es.ravel()
        ta = ta.ravel()
        if es.shape[0] != ta.shape[0]:
            raise Exception('invalid shape of es, or ta')
        if es.min() < 0 or es.max() > 1:
            raise Exception('estimate has wrong value range')
        ta_p = (ta == 1)
        ta_n = (ta == 0)
        es_p = es[ta_p]
        es_n = es[ta_n]
        for idx, wp in enumerate(self.thresholds):
            wp = np.asscalar(wp)
            self.tps[idx] += (es_p > wp).sum()
            self.fps[idx] += (es_n > wp).sum()
            self.fns[idx] += (es_p <= wp).sum()
            self.tns[idx] += (es_n <= wp).sum()
        self.n_pos += ta_p.sum()
        self.n_neg += ta_n.sum()

    def get(self):
        tps = np.array(self.tps).astype(np.float32)
        fps = np.array(self.fps).astype(np.float32)
        fns = np.array(self.fns).astype(np.float32)
        tns = np.array(self.tns).astype(np.float32)
        wp = self.thresholds

        ret = {}

        precisions = np.divide(tps, tps + fps, out=np.zeros_like(tps), where=tps + fps != 0)
        recalls = np.divide(tps, tps + fns, out=np.zeros_like(tps), where=tps + fns != 0)  # tprs
        fprs = np.divide(fps, fps + tns, out=np.zeros_like(tps), where=fps + tns != 0)

        precisions = np.r_[0, precisions, 1]
        recalls = np.r_[1, recalls, 0]
        fprs = np.r_[1, fprs, 0]

        ret['auc'] = float(-np.trapz(recalls, fprs))
        ret['prauc'] = float(-np.trapz(precisions, recalls))
        ret['ap'] = float(-(np.diff(recalls) * precisions[:-1]).sum())

        accuracies = np.divide(tps + tns, tps + tns + fps + fns)
        aacc = np.mean(accuracies)
        for t in np.linspace(0, 1, num=11)[1:-1]:
            idx = np.argmin(np.abs(t - wp))
            ret[f'acc{wp[idx]:.2f}'] = float(accuracies[idx])

        return ret