# (C) British Crown Copyright 2015, Met Office
#
# This file is part of Iris.
#
# Iris is free software: you can redistribute it and/or modify it under
# the terms of the GNU Lesser General Public License as published by the
# Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Iris is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with Iris.  If not, see <http://www.gnu.org/licenses/>.
"""Unit tests for the :data:`iris.analysis.PERCENTILE` aggregator."""

from __future__ import (absolute_import, division, print_function)
from six.moves import range

# Import iris.tests first so that some things can be initialised before
# importing anything else.
import iris.tests as tests

import numpy as np
import numpy.ma as ma

from iris.analysis import WPERCENTILE


class Test_aggregate(tests.IrisTest):
    def test_missing_mandatory_kwargs(self):
        emsg = "weighted_percentile aggregator requires " \
               ".* keyword argument 'percent'"
        with self.assertRaisesRegexp(ValueError, emsg):
            WPERCENTILE.aggregate('dummy', axis=0, weights=None)
        emsg = "weighted_percentile aggregator requires " \
               ".* keyword argument 'weights'"
        with self.assertRaisesRegexp(ValueError, emsg):
            WPERCENTILE.aggregate('dummy', axis=0, percent=50)

    def test_wrong_weights_shape(self):
        data = np.arange(11)
        weights = np.ones(10)
        emsg = "_weighted_percentile: weights wrong shape."
        with self.assertRaisesRegexp(ValueError, emsg):
            WPERCENTILE.aggregate(data, axis=0, percent=50, weights=weights)

    def test_1d_single(self):
        data = np.arange(11)
        weights = np.ones(data.shape)
        actual = WPERCENTILE.aggregate(data, axis=0, percent=50,
                                       weights=weights)
        expected = 5
        self.assertTupleEqual(actual.shape, ())
        self.assertEqual(actual, expected)

    def test_1d_single_unequal(self):
        data = np.arange(12)
        weights = np.ones(data.shape)
        weights[0:3] = 3
        actual, weight_total = WPERCENTILE.aggregate(
            data, axis=0, percent=50, weights=weights, returned=True)
        expected = 2.75
        self.assertTupleEqual(actual.shape, ())
        self.assertEqual(actual, expected)
        self.assertEqual(weight_total, 18)

    def test_masked_1d_single(self):
        data = ma.arange(11)
        weights = np.ones(data.shape)
        data[3:7] = ma.masked
        actual = WPERCENTILE.aggregate(data, axis=0, percent=50,
                                       weights=weights)
        expected = 7
        self.assertTupleEqual(actual.shape, ())
        self.assertEqual(actual, expected)

    def test_1d_multi(self):
        data = np.arange(11)
        weights = np.ones(data.shape)
        percent = np.array([20, 50, 90])
        actual = WPERCENTILE.aggregate(data, axis=0, percent=percent,
                                       weights=weights)
        expected = [1.7, 5, 9.4]
        self.assertTupleEqual(actual.shape, percent.shape)
        self.assertArrayAlmostEqual(actual, expected)

    def test_1d_multi_unequal(self):
        data = np.arange(13)
        weights = np.ones(data.shape)
        weights[1::2] = 3
        percent = np.array([20, 50, 96])
        actual = WPERCENTILE.aggregate(data, axis=0, percent=percent,
                                       weights=weights)
        expected = [2.25, 6, 11.75]
        self.assertTupleEqual(actual.shape, percent.shape)
        self.assertArrayAlmostEqual(actual, expected)

    def test_masked_1d_multi(self):
        data = ma.arange(11)
        weights = np.ones(data.shape)
        data[3:9] = ma.masked
        percent = np.array([25, 50, 75])
        actual = WPERCENTILE.aggregate(data, axis=0, percent=percent,
                                       weights=weights)
        expected = [0.75, 2, 9.25]
        self.assertTupleEqual(actual.shape, percent.shape)
        self.assertArrayAlmostEqual(actual, expected)

    def test_2d_single(self):
        shape = (2, 11)
        data = np.arange(np.prod(shape)).reshape(shape)
        weights = np.ones(shape)
        actual = WPERCENTILE.aggregate(data, axis=0, percent=50,
                                       weights=weights)
        self.assertTupleEqual(actual.shape, shape[-1:])
        expected = np.arange(shape[-1]) + 5.5
        self.assertArrayEqual(actual, expected)

    def test_masked_2d_single(self):
        shape = (2, 11)
        data = ma.arange(np.prod(shape)).reshape(shape)
        data[0, ::2] = ma.masked
        data[1, 1::2] = ma.masked
        weights = np.ones(shape)
        actual = WPERCENTILE.aggregate(data, axis=0, percent=50,
                                       weights=weights)
        self.assertTupleEqual(actual.shape, shape[-1:])
        expected = np.empty(shape[-1:])
        expected[1::2] = data[0, 1::2]
        expected[::2] = data[1, ::2]
        self.assertArrayEqual(actual, expected)

    def test_2d_multi(self):
        shape = (2, 10)
        data = np.arange(np.prod(shape)).reshape(shape)
        weights = np.ones(shape)
        percent = np.array([10, 50, 70, 100])
        actual = WPERCENTILE.aggregate(data, axis=0, percent=percent,
                                       weights=weights)
        self.assertTupleEqual(actual.shape, (shape[-1], percent.size))
        expected = np.tile(np.arange(shape[-1]), percent.size)
        expected = expected.reshape(percent.size, shape[-1]).T
        expected[:, 1:-1] += (percent[1:-1]-25)*0.2
        expected[:, -1] += 10.
        self.assertArrayAlmostEqual(actual, expected)

    def test_masked_2d_multi(self):
        shape = (3, 10)
        data = ma.arange(np.prod(shape)).reshape(shape)
        weights = np.ones(shape)
        data[1] = ma.masked
        percent = np.array([10, 50, 70, 80])
        actual = WPERCENTILE.aggregate(data, axis=0, percent=percent,
                                       weights=weights)
        self.assertTupleEqual(actual.shape, (shape[-1], percent.size))
        expected = np.tile(np.arange(shape[-1]), percent.size)
        expected = expected.reshape(percent.size, shape[-1]).T
        expected[:, 1:-1] += (percent[1:-1]-25)*0.4
        expected[:, -1] += 20.
        self.assertArrayAlmostEqual(actual, expected)

    def test_masked_2d_multi_unequal(self):
        shape = (3, 10)
        data = ma.arange(np.prod(shape)).reshape(shape)
        weights = np.ones(shape)
        weights[0] = 3
        data[1] = ma.masked
        percent = np.array([30, 50, 75, 80])
        actual, weight_total = WPERCENTILE.aggregate(
            data, axis=0, percent=percent, weights=weights, returned=True)
        self.assertTupleEqual(actual.shape, (shape[-1], percent.size))
        expected = np.tile(np.arange(shape[-1]), percent.size)
        expected = expected.reshape(percent.size, shape[-1]).T
        expected[:, 1:] = 2.0 * (
            (0.875 - percent[1:]/100.0) * data[0, np.newaxis].T +
            (percent[1:]/100.0 - 0.375) * data[-1, np.newaxis].T
            )
        self.assertArrayAlmostEqual(actual, expected)
        self.assertTupleEqual(weight_total.shape, (shape[-1],))
        self.assertArrayEqual(weight_total, np.repeat(4, shape[-1]))


class Test_name(tests.IrisTest):
    def test(self):
        self.assertEqual(WPERCENTILE.name(), 'weighted_percentile')


class Test_aggregate_shape(tests.IrisTest):
    def test_missing_mandatory_kwarg(self):
        emsg_pc = "weighted_percentile aggregator requires " \
                  ".* keyword argument 'percent'"
        emsg_wt = "weighted_percentile aggregator requires " \
                  ".* keyword argument 'weights'"
        with self.assertRaisesRegexp(ValueError, emsg_pc):
            WPERCENTILE.aggregate_shape(weights=None)
        with self.assertRaisesRegexp(ValueError, emsg_pc):
            kwargs = dict(weights=None)
            WPERCENTILE.aggregate_shape(**kwargs)
        with self.assertRaisesRegexp(ValueError, emsg_pc):
            kwargs = dict(point=10)
            WPERCENTILE.aggregate_shape(**kwargs)
        with self.assertRaisesRegexp(ValueError, emsg_wt):
            WPERCENTILE.aggregate_shape(percent=50)
        with self.assertRaisesRegexp(ValueError, emsg_wt):
            kwargs = dict(percent=50)
            WPERCENTILE.aggregate_shape(**kwargs)
        with self.assertRaisesRegexp(ValueError, emsg_wt):
            kwargs = dict(percent=50, weight=None)
            WPERCENTILE.aggregate_shape(**kwargs)

    def test_mandatory_kwarg_no_shape(self):
        kwargs = dict(percent=50, weights=None)
        self.assertTupleEqual(WPERCENTILE.aggregate_shape(**kwargs), ())
        kwargs = dict(percent=[50], weights=None)
        self.assertTupleEqual(WPERCENTILE.aggregate_shape(**kwargs), ())

    def test_mandatory_kwarg_shape(self):
        kwargs = dict(percent=(10, 20), weights=None)
        self.assertTupleEqual(WPERCENTILE.aggregate_shape(**kwargs), (2,))
        kwargs = dict(percent=range(13), weights=None)
        self.assertTupleEqual(WPERCENTILE.aggregate_shape(**kwargs), (13,))


class Test_cell_method(tests.IrisTest):
    def test(self):
        self.assertIsNone(WPERCENTILE.cell_method)


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