# Copyright 2021 Huawei Technologies Co., Ltd

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

#     http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.



# Model validation metrics



import matplotlib.pyplot as plt

import numpy as np





def fitness(x):

    # Model fitness as a weighted combination of metrics

    w = [0.0, 0.0, 0.1, 0.9]  # weights for [P, R, mAP@0.5, mAP@0.5:0.95]

    return (x[:, :4] * w).sum(1)





def fitness_p(x):

    # Model fitness as a weighted combination of metrics

    w = [1.0, 0.0, 0.0, 0.0]  # weights for [P, R, mAP@0.5, mAP@0.5:0.95]

    return (x[:, :4] * w).sum(1)





def fitness_r(x):

    # Model fitness as a weighted combination of metrics

    w = [0.0, 1.0, 0.0, 0.0]  # weights for [P, R, mAP@0.5, mAP@0.5:0.95]

    return (x[:, :4] * w).sum(1)





def fitness_ap50(x):

    # Model fitness as a weighted combination of metrics

    w = [0.0, 0.0, 1.0, 0.0]  # weights for [P, R, mAP@0.5, mAP@0.5:0.95]

    return (x[:, :4] * w).sum(1)





def fitness_ap(x):

    # Model fitness as a weighted combination of metrics

    w = [0.0, 0.0, 0.0, 1.0]  # weights for [P, R, mAP@0.5, mAP@0.5:0.95]

    return (x[:, :4] * w).sum(1)





def fitness_f(x):

    # Model fitness as a weighted combination of metrics

    #w = [0.0, 0.0, 0.0, 1.0]  # weights for [P, R, mAP@0.5, mAP@0.5:0.95]

    return ((x[:, 0]*x[:, 1])/(x[:, 0]+x[:, 1]))





def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, fname='precision-recall_curve.png'):

    """ Compute the average precision, given the recall and precision curves.

    Source: https://github.com/rafaelpadilla/Object-Detection-Metrics.

    # Arguments

        tp:  True positives (nparray, nx1 or nx10).

        conf:  Objectness value from 0-1 (nparray).

        pred_cls:  Predicted object classes (nparray).

        target_cls:  True object classes (nparray).

        plot:  Plot precision-recall curve at mAP@0.5

        fname:  Plot filename

    # Returns

        The average precision as computed in py-faster-rcnn.

    """



    # Sort by objectness

    i = np.argsort(-conf)

    tp, conf, pred_cls = tp[i], conf[i], pred_cls[i]



    # Find unique classes

    unique_classes = np.unique(target_cls)



    # Create Precision-Recall curve and compute AP for each class

    px, py = np.linspace(0, 1, 1000), []  # for plotting

    pr_score = 0.1  # score to evaluate P and R https://github.com/ultralytics/yolov3/issues/898

    s = [unique_classes.shape[0], tp.shape[1]]  # number class, number iou thresholds (i.e. 10 for mAP0.5...0.95)

    ap, p, r = np.zeros(s), np.zeros(s), np.zeros(s)

    for ci, c in enumerate(unique_classes):

        i = pred_cls == c

        n_l = (target_cls == c).sum()  # number of labels

        n_p = i.sum()  # number of predictions



        if n_p == 0 or n_l == 0:

            continue

        else:

            # Accumulate FPs and TPs

            fpc = (1 - tp[i]).cumsum(0)

            tpc = tp[i].cumsum(0)



            # Recall

            recall = tpc / (n_l + 1e-16)  # recall curve

            r[ci] = np.interp(-pr_score, -conf[i], recall[:, 0])  # r at pr_score, negative x, xp because xp decreases



            # Precision

            precision = tpc / (tpc + fpc)  # precision curve

            p[ci] = np.interp(-pr_score, -conf[i], precision[:, 0])  # p at pr_score



            # AP from recall-precision curve

            for j in range(tp.shape[1]):

                ap[ci, j], mpre, mrec = compute_ap(recall[:, j], precision[:, j])

                if j == 0:

                    py.append(np.interp(px, mrec, mpre))  # precision at mAP@0.5



    # Compute F1 score (harmonic mean of precision and recall)

    f1 = 2 * p * r / (p + r + 1e-16)



    if plot:

        py = np.stack(py, axis=1)

        fig, ax = plt.subplots(1, 1, figsize=(5, 5))

        ax.plot(px, py, linewidth=0.5, color='grey')  # plot(recall, precision)

        ax.plot(px, py.mean(1), linewidth=2, color='blue', label='all classes %.3f mAP@0.5' % ap[:, 0].mean())

        ax.set_xlabel('Recall')

        ax.set_ylabel('Precision')

        ax.set_xlim(0, 1)

        ax.set_ylim(0, 1)

        plt.legend()

        fig.tight_layout()

        fig.savefig(fname, dpi=200)



    return p, r, ap, f1, unique_classes.astype('int32')





def compute_ap(recall, precision):

    """ Compute the average precision, given the recall and precision curves.

    Source: https://github.com/rbgirshick/py-faster-rcnn.

    # Arguments

        recall:    The recall curve (list).

        precision: The precision curve (list).

    # Returns

        The average precision as computed in py-faster-rcnn.

    """



    # Append sentinel values to beginning and end

    mrec = np.concatenate(([0.0], recall, [1.0]))

    mpre = np.concatenate(([1.0], precision, [0.0]))



    # Compute the precision envelope

    mpre = np.flip(np.maximum.accumulate(np.flip(mpre)))



    # Integrate area under curve

    method = 'interp'  # methods: 'continuous', 'interp'

    if method == 'interp':

        x = np.linspace(0, 1, 101)  # 101-point interp (COCO)

        ap = np.trapz(np.interp(x, mrec, mpre), x)  # integrate

    else:  # 'continuous'

        i = np.where(mrec[1:] != mrec[:-1])[0]  # points where x axis (recall) changes

        ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])  # area under curve



    return ap, mpre, mrec