#   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
#
# 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.

import unittest

import numpy as np
from eager_op_test import OpTest


def compute_mean_iou(
    predictions, labels, num_classes, in_wrongs, in_corrects, in_mean_ious
):
    assert predictions.shape == labels.shape
    predictions = predictions.flatten()
    labels = labels.flatten()

    out_wrong = np.zeros([num_classes]).astype("int32")
    for _, wrong in in_wrongs:
        out_wrong += wrong
    out_correct = np.zeros([num_classes]).astype("int32")
    for _, correct in in_corrects:
        out_correct += correct

    for pred, label in zip(predictions, labels):
        if pred == label:
            out_correct[pred] += 1
        else:
            out_wrong[pred] += 1
            out_wrong[label] += 1

    denominator = out_wrong + out_correct
    valid_count = (denominator != 0).sum()
    denominator = np.where(
        denominator > 0, denominator, np.ones(denominator.shape)
    )
    mean_iou = (out_correct / denominator).sum() / valid_count

    for _, in_mean_iou in in_mean_ious:
        mean_iou += float(in_mean_iou)
    return mean_iou, out_wrong, out_correct


class TestMeanIOUOp(OpTest):
    def setUp(self):
        self.config()
        self.op_type = "mean_iou"
        predictions = np.random.randint(
            0, self.num_classes, self.image_size
        ).astype("int32")
        labels = np.random.randint(0, self.num_classes, self.image_size).astype(
            "int32"
        )

        in_wrongs = []
        for i in range(self.in_wrong_num):
            in_wrongs.append(
                (
                    "in_wrong_%d" % i,
                    np.random.randint(0, 10, [self.num_classes]).astype(
                        "int32"
                    ),
                )
            )

        in_corrects = []
        for i in range(self.in_correct_num):
            in_corrects.append(
                (
                    "in_correct_%d" % i,
                    np.random.randint(0, 10, [self.num_classes]).astype(
                        "int32"
                    ),
                )
            )

        self.inputs = {
            'Predictions': predictions,
            'Labels': labels,
            'InWrongs': in_wrongs,
            'InCorrects': in_corrects,
            'InMeanIou': self.in_mean_ious,
        }
        self.attrs = {'num_classes': int(self.num_classes)}
        mean_iou, out_wrong, out_correct = compute_mean_iou(
            predictions,
            labels,
            self.num_classes,
            in_wrongs,
            in_corrects,
            self.in_mean_ious,
        )
        self.outputs = {
            'OutMeanIou': mean_iou,
            'OutWrong': out_wrong,
            'OutCorrect': out_correct,
        }

    def config(self):
        self.num_classes = 10
        self.image_size = [128, 128]
        self.in_wrong_num = 0
        self.in_correct_num = 0
        self.in_mean_ious = []

    def test_check_output(self):
        self.check_output()


class TestCase1(TestMeanIOUOp):
    def config(self):
        self.num_classes = 5
        self.image_size = [100, 128]
        self.in_wrong_num = 2
        self.in_correct_num = 2
        self.in_mean_ious = []
        for i in range(2):
            self.in_mean_ious.append(
                (
                    "in_mean_iou_%d" % i,
                    np.random.uniform(0, 1, []).astype("float32"),
                )
            )

    # NOTE(dev): Skip check_dygraph becuase Python API doesn't expose
    # in_wrong_num/in_correct_num/in_mean_iou_num argument
    def test_check_output(self):
        self.check_output(check_dygraph=False)


class TestCase2(TestCase1):
    def config(self):
        self.num_classes = 5
        self.image_size = [100, 128]
        self.in_wrong_num = 2
        self.in_correct_num = 2
        self.in_mean_ious = []
        for i in range(2):
            self.in_mean_ious.append(
                (
                    "in_mean_iou_%d" % i,
                    np.random.uniform(0, 1, [1]).astype("float32"),
                )
            )


if __name__ == '__main__':
    unittest.main()