from __future__ import print_function
import unittest
import verif.metric
import verif.metric_type
import numpy as np


def get(filenames):
    inputs = [
        verif.input.get_input(file)
        for file in ["verif/tests/files/" + f for f in filenames]
    ]
    data = verif.data.Data(inputs=inputs)
    return data


interval6 = verif.interval.Interval(5.5, 6.5, True, True)
interval100 = verif.interval.Interval(99.5, 100.5, True, True)


class MetricTest(unittest.TestCase):
    def test_func(self):
        metric = verif.metric.Mae()
        metric.aggregator = verif.aggregator.Mean()
        value = metric.compute_from_obs_fcst(np.array([2, 1, 2]), np.array([2, 3, 1]))
        self.assertEqual(value, 1)
        value = metric.compute_from_obs_fcst(np.array([-2]), np.array([2]))
        self.assertEqual(value, 4)
        value = metric.compute_from_obs_fcst(np.array([]), np.array([]))
        self.assertTrue(np.isnan(value))
        value = metric.compute_from_obs_fcst(
            np.array([2, np.nan, 2]), np.array([2, 3, 1])
        )
        self.assertEqual(value, 0.5)

    def test_get_all(self):
        types = verif.metric.get_all()
        self.assertTrue(verif.metric.Mae in [t[1] for t in types])

    def test_get_all_by_type(self):
        types = verif.metric.get_all_by_type(verif.metric_type.Deterministic())
        self.assertTrue(verif.metric.Mae in [t[1] for t in types])

        types = verif.metric.get_all_by_type(float)
        self.assertEqual(len(types), 0)

    def test_get(self):
        metric = verif.metric.get("mae")
        self.assertTrue(isinstance(metric, verif.metric.Mae))

        # with self.assertRaises(ValueError):
        metric = verif.metric.get("not_exist")
        self.assertEqual(metric, None)

    def test_get_all_obs_fcst_based(self):
        types = verif.metric.get_all_obs_fcst_based()
        self.assertTrue(verif.metric.Mae in [t[1] for t in types])

    def test_name(self):
        mae = verif.metric.Mae()
        self.assertEqual("Mean absolute error", mae.name)

    def test_reample(self):
        metric = verif.metric.Ets()

        obs = np.random.randn(10000)
        fcst = np.random.randn(10000)
        othreshold = verif.interval.Interval(0.1, np.inf, True, True)
        fthresholds = [
            verif.interval.Interval(i, np.inf, True, True) for i in np.random.randn(1)
        ]

        # Check that this runs
        q = np.zeros(len(fthresholds), float)
        for i in range(0, len(fthresholds)):
            q[i] = metric.compute_from_obs_fcst_resample(
                obs, fcst, 1, othreshold, fthresholds[i]
            )
        for i in range(0, len(fthresholds)):
            q[i] = metric.compute_from_obs_fcst(obs, fcst, fthresholds[i])

    def test_label(self):
        variable = verif.variable.Variable("test", "C")
        metric = verif.metric.Mae()
        self.assertEqual(metric.label(variable), "MAE (C)")

    def test_is_valid(self):
        metric = verif.metric.Mae()
        self.assertTrue(metric.is_valid())

        metric = verif.metric.Metric()
        self.assertFalse(metric.is_valid())

        metric = verif.metric.XConditional()
        self.assertFalse(metric.is_valid())

    def test_help_and_label(self):
        """Check that all metrics have a help text and a label"""
        variable = verif.variable.Variable("test", "C")
        for name, metric in verif.metric.get_all():
            if metric.is_valid():
                with self.subTest(metric=name):
                    string = metric.help()
                    self.assertTrue(len(string) > 0)
                    metric = metric()
                    label = metric.label(variable)
                    self.assertTrue(len(label) > 0)


class TestObsFcstBased(unittest.TestCase):
    def test_basic(self):
        # Try ObsFcstBased metrics here only
        metrics = [
            verif.metric.Rmse(),
            verif.metric.Mae(),
            verif.metric.Bias(),
            verif.metric.Corr(),
            verif.metric.Dmb(),
            verif.metric.Mbias(),
            verif.metric.Ef(),
        ]
        obs = [[0, 1.5, 2], [0]]  # Case1, Case 2, ...
        fcst = [[3.1, 1.1, -2.1], [1]]
        expected = [
            [2.976575213, 1],  # Rmse
            [2.5333333333, 1],  # Mae
            [-0.466666666666, 1],  # Bias
            [-0.915724295, np.nan],  # Corr
            [1.0 / 0.6, 0],  # Dmb
            [0.6, np.nan],  # Mbias
            [1.0 / 3, 1],  # Ef
        ]

        for m in range(len(metrics)):
            metric = metrics[m]
            for i in range(0, len(obs)):
                value = metric.compute_from_obs_fcst(
                    np.array(obs[i]), np.array(fcst[i])
                )
                if np.isnan(value):
                    self.assertTrue(np.isnan(expected[m][i]))
                else:
                    self.assertAlmostEqual(expected[m][i], value)

    def test_obs_axis(self):
        data = get(["file1.txt"])
        metric = verif.metric.Mae()
        value = metric.compute_single(data, 0, verif.axis.Obs(), 0, interval6)
        # Check that the MAE when obs == 6 is indeed 7
        self.assertEqual(value, 1)

        value = metric.compute_single(data, 0, verif.axis.Obs(), 0, interval100)
        self.assertTrue(np.isnan(value))

    def test_fcst_axis(self):
        data = get(["file1.txt"])
        metric = verif.metric.Mae()
        value = metric.compute_single(data, 0, verif.axis.Fcst(), 0, interval6)
        # Check that the MAE when fcst == 6 is indeed 3
        self.assertEqual(value, 3)

        value = metric.compute_single(data, 0, verif.axis.Fcst(), 0, interval100)
        self.assertTrue(np.isnan(value))


class TestFromField(unittest.TestCase):
    def test_obs_axis(self):
        data = get(["file1.txt"])
        metric = verif.metric.Fcst()
        value = metric.compute_single(data, 0, verif.axis.Obs(), 0, interval6)
        # Check that the average forecast when obs == 6 is indeed 7
        self.assertEqual(value, 7)

    def test_fcst_axis(self):
        data = get(["file1.txt"])
        metric = verif.metric.Mae()
        value = metric.compute_single(data, 0, verif.axis.Fcst(), 0, interval6)
        # Check that the average obs when fcst == 6 is indeed 3
        self.assertEqual(value, 3)


class TestThresholdBased(unittest.TestCase):
    def test_basic(self):
        """Test the basic threshold-based metrics"""
        metrics = [
            verif.metric.Within(),
            verif.metric.A(),  # Hit
            verif.metric.B(),  # FA
            verif.metric.C(),  # Miss
            verif.metric.D(),  # Correct rejection
            verif.metric.Hit(),
            verif.metric.Threat(),
            verif.metric.Conditional(),
            verif.metric.XConditional(func=np.median),
        ]
        intervals = [
            verif.interval.Interval(-np.inf, 0, True, True),  # [-inf, 0]
            verif.interval.Interval(-np.inf, 1, True, True),
            verif.interval.Interval(-np.inf, 2, True, True),
        ]
        obs = [0, 1.5, 2]
        fcst = [3.1, 1.1, -2.1]
        N = len(obs) * 1.0

        # Each line is one metric (one number for each threshold)
        expected = [
            [0 / N, 100 / N, 100 / N],  # Within
            [0 / N, 0 / N, 2 / N],  # Hit
            [1 / N, 1 / N, 0 / N],  # FA
            [1 / N, 1 / N, 1 / N],  # Miss
            [1 / N, 1 / N, 0 / N],  # Correct rejection
            [0, 0, 2.0 / 3],  # Hit rate
            [0, 0, 2.0 / 3],  # Threat score
            [3.1, 3.1, 0.7],  # Average fcst given obs in interval
            [0, 0, 1.5],  # Average obs given obs in interval
        ]

        for m in range(len(metrics)):
            metric = metrics[m]
            for i in range(len(intervals)):
                value = metric.compute_from_obs_fcst(
                    np.array(obs), np.array(fcst), intervals[i]
                )
                ex = expected[m][i] * 1.0
                if np.isnan(value):
                    self.assertTrue(np.isnan(ex))
                else:
                    self.assertAlmostEqual(ex, value)

    def test_threat(self):
        """An extra test for Threat score"""
        metric = verif.metric.Threat()
        obs = np.array([0, 1, 2, 3])
        fcst = np.array([0, 3, 1, 2])

        # Hits: 1
        # FA: 1
        # Miss: 1
        # CR: 0
        interval = verif.interval.Interval(1.5, np.inf, True, True)
        f_interval = verif.interval.Interval(1.5, np.inf, True, True)
        value = metric.compute_from_obs_fcst(obs, fcst, interval, f_interval)
        self.assertEqual(value, 1.0 / 3)


class NameTest(unittest.TestCase):
    def test_valid(self):
        metrics = [x[1] for x in verif.metric.get_all() if x[1].is_valid]
        for metric in metrics:
            self.assertTrue(metric.name != "")
            self.assertTrue(metric.description != "")


class TestBrierScore(unittest.TestCase):
    def test(self):
        """
        These have been checked against the Brier scores in the verification R package. The 'bins'
        attribute was set to False, and the code had to be changed because it didn't work for cases
        where forecasts were not equal to whole 1/10ths (i.e. 0.1, 0.2).

        To give the same results, there can only be one forecasted value within each bin. I.e. 0.21
        and 0.22 cannot appear in the same test case.
        """
        bs = verif.metric.Bs()
        bsrel = verif.metric.BsRel()
        bsres = verif.metric.BsRes()
        bsunc = verif.metric.BsUnc()
        bss = verif.metric.Bss()
        obs = [
            [0],
            [0],
            [0],
            [1],
            [0, 0, 1, 1, 1],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        ]
        fcst = [
            [0],
            [1],
            [0.3],
            [0.1],
            [0.21, 0.21, 0.21, 0.91, 0.91],
            [
                0.06,
                0.61,
                0.45,
                0.87,
                0.13,
                0.61,
                0.79,
                0.61,
                0.06,
                0.06,
                0.79,
                0.61,
                0.13,
                0.13,
                0.79,
                0.21,
                0.06,
                0.55,
                0.37,
                0.37,
            ],
        ]
        ans = {
            bs: [0, 1, 0.09, 0.81, 0.1457, 0.34928],
            bsrel: [0, 1, 0.09, 0.81, 0.01236667, 0.2076133],
            bsres: [0, 0, 0, 0, 0.1066667, 0.1083333],
            bsunc: [0, 0, 0, 0, 0.24, 0.25],
            bss: [np.nan, np.nan, np.nan, np.nan, 0.3929167, -0.39712],
        }
        for i in range(len(obs)):
            o = np.array(obs[i])
            f = np.array(fcst[i])
            for key in ans:
                calculated = key.compute_from_obs_fcst(o, f)
                expected = ans[key][i]
                if np.isnan(expected):
                    self.assertTrue(np.isnan(expected), np.isnan(calculated))
                else:
                    self.assertAlmostEqual(expected, calculated, places=5)


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