"""
Test cate/ops/coregistration.py

Test coregistration, checks if the values seem as expected
when using default upsampling/downsampling methods.

"""

from unittest import TestCase

import numpy as np
import xarray as xr
from numpy.testing import assert_almost_equal, assert_array_equal

from cate.core.op import OP_REGISTRY
from cate.util.misc import object_to_qualified_name

from cate.ops import coregister
from cate.ops.coregistration import _find_intersection
from ..util.test_monitor import RecordingMonitor

class TestCoregistration(TestCase):
    """
    Test coregistration
    """
    def test_nominal(self):
        """
        Test nominal execution
        """
        ds_fine = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.array([np.eye(4, 8), np.eye(4, 8)])),
            'second': (['time', 'lat', 'lon'], np.array([np.eye(4, 8), np.eye(4, 8)])),
            'lat': np.linspace(-67.5, 67.5, 4),
            'lon': np.linspace(-157.5, 157.5, 8),
            'time': np.array([1, 2])}).chunk(chunks={'lat': 2, 'lon': 4})

        ds_coarse = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.array([np.eye(3, 6), np.eye(3, 6)])),
            'second': (['time', 'lat', 'lon'], np.array([np.eye(3, 6), np.eye(3, 6)])),
            'lat': np.linspace(-60, 60, 3),
            'lon': np.linspace(-150, 150, 6),
            'time': np.array([1, 2])}).chunk(chunks={'lat': 3, 'lon': 3})

        # Test that the coarse dataset has been resampled onto the grid
        # of the finer dataset.
        rm = RecordingMonitor()
        ds_coarse_resampled = coregister(ds_fine, ds_coarse, monitor=rm)
        self.assertEqual([('start', 'coregister dataset', 2),
                          ('progress', 0.0, 'coregister dataarray', 0),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 0),
                          ('progress', 0.125, None, 6),
                          ('progress', 0.125, None, 13),
                          ('progress', 0.125, None, 19),
                          ('progress', 0.125, None, 25),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 25),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 25),
                          ('progress', 0.125, None, 31),
                          ('progress', 0.125, None, 38),
                          ('progress', 0.125, None, 44),
                          ('progress', 0.125, None, 50),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 50),
                          ('progress', 0.0, 'coregister dataarray', 50),
                          ('progress', 0.0, 'coregister dataarray', 50),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 50),
                          ('progress', 0.125, None, 56),
                          ('progress', 0.125, None, 63),
                          ('progress', 0.125, None, 69),
                          ('progress', 0.125, None, 75),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 75),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 75),
                          ('progress', 0.125, None, 81),
                          ('progress', 0.125, None, 88),
                          ('progress', 0.125, None, 94),
                          ('progress', 0.125, None, 100),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 100),
                          ('progress', 0.0, 'coregister dataarray', 100),
                          ('done',)], rm.records)

        expected = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.array([[[1., 0.28571429, 0., 0., 0., 0., 0., 0.],
                                                         [0.33333333, 0.57142857, 0.38095238, 0., 0., 0., 0., 0.],
                                                         [0., 0.47619048, 0.52380952, 0.28571429, 0.04761905, 0., 0.,
                                                          0.],
                                                         [0., 0., 0.42857143, 0.85714286, 0.14285714, 0., 0., 0.]],
                                                        [[1., 0.28571429, 0., 0., 0., 0., 0., 0.],
                                                         [0.33333333, 0.57142857, 0.38095238, 0., 0., 0., 0., 0.],
                                                         [0., 0.47619048, 0.52380952, 0.28571429, 0.04761905, 0., 0.,
                                                          0.],
                                                         [0., 0., 0.42857143, 0.85714286, 0.14285714, 0., 0., 0.]]])),
            'second': (['time', 'lat', 'lon'], np.array([[[1., 0.28571429, 0., 0., 0., 0., 0., 0.],
                                                          [0.33333333, 0.57142857, 0.38095238, 0., 0., 0., 0., 0.],
                                                          [0., 0.47619048, 0.52380952, 0.28571429, 0.04761905, 0., 0.,
                                                           0.],
                                                          [0., 0., 0.42857143, 0.85714286, 0.14285714, 0., 0., 0.]],
                                                         [[1., 0.28571429, 0., 0., 0., 0., 0., 0.],
                                                          [0.33333333, 0.57142857, 0.38095238, 0., 0., 0., 0., 0.],
                                                          [0., 0.47619048, 0.52380952, 0.28571429, 0.04761905, 0., 0.,
                                                           0.],
                                                          [0., 0., 0.42857143, 0.85714286, 0.14285714, 0., 0., 0.]]])),
            'lat': np.linspace(-67.5, 67.5, 4),
            'lon': np.linspace(-157.5, 157.5, 8),
            'time': np.array([1, 2])})
        assert_almost_equal(ds_coarse_resampled['first'].values, expected['first'].values)

        # Test that the fine dataset has been resampled (aggregated)
        # onto the grid of the coarse dataset.
        ds_fine_resampled = coregister(ds_coarse, ds_fine)
        expected = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.array([[[0.625, 0.125, 0., 0., 0., 0.],
                                                         [0.125, 0.5, 0.125, 0., 0., 0.],
                                                         [0., 0.125, 0.625, 0., 0., 0.]],

                                                        [[0.625, 0.125, 0., 0., 0., 0.],
                                                         [0.125, 0.5, 0.125, 0., 0., 0.],
                                                         [0., 0.125, 0.625, 0., 0., 0.]]])),
            'second': (['time', 'lat', 'lon'], np.array([[[0.625, 0.125, 0., 0., 0., 0.],
                                                          [0.125, 0.5, 0.125, 0., 0., 0.],
                                                          [0., 0.125, 0.625, 0., 0., 0.]],

                                                         [[0.625, 0.125, 0., 0., 0., 0.],
                                                          [0.125, 0.5, 0.125, 0., 0., 0.],
                                                          [0., 0.125, 0.625, 0., 0., 0.]]])),
            'lat': np.linspace(-60, 60, 3),
            'lon': np.linspace(-150, 150, 6),
            'time': np.array([1, 2])})

        assert_almost_equal(ds_fine_resampled['first'].values, expected['first'].values)

    def test_registered(self):
        """
        Test registered operation execution execution
        """
        reg_op = OP_REGISTRY.get_op(object_to_qualified_name(coregister))
        ds_fine = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.array([np.eye(4, 8), np.eye(4, 8)])),
            'second': (['time', 'lat', 'lon'], np.array([np.eye(4, 8), np.eye(4, 8)])),
            'lat': np.linspace(-67.5, 67.5, 4),
            'lon': np.linspace(-157.5, 157.5, 8),
            'time': np.array([1, 2])})

        ds_coarse = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.array([np.eye(3, 6), np.eye(3, 6)])),
            'second': (['time', 'lat', 'lon'], np.array([np.eye(3, 6), np.eye(3, 6)])),
            'lat': np.linspace(-60, 60, 3),
            'lon': np.linspace(-150, 150, 6),
            'time': np.array([1, 2])})

        # Test that the coarse dataset has been resampled onto the grid
        # of the finer dataset.
        ds_coarse_resampled = reg_op(ds_master=ds_fine, ds_replica=ds_coarse)
        expected = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.array([[[1., 0.28571429, 0., 0., 0., 0., 0., 0.],
                                                         [0.33333333, 0.57142857, 0.38095238, 0., 0., 0., 0., 0.],
                                                         [0., 0.47619048, 0.52380952, 0.28571429, 0.04761905, 0., 0.,
                                                          0.],
                                                         [0., 0., 0.42857143, 0.85714286, 0.14285714, 0., 0., 0.]],
                                                        [[1., 0.28571429, 0., 0., 0., 0., 0., 0.],
                                                         [0.33333333, 0.57142857, 0.38095238, 0., 0., 0., 0., 0.],
                                                         [0., 0.47619048, 0.52380952, 0.28571429, 0.04761905, 0., 0.,
                                                          0.],
                                                         [0., 0., 0.42857143, 0.85714286, 0.14285714, 0., 0., 0.]]])),
            'second': (['time', 'lat', 'lon'], np.array([[[1., 0.28571429, 0., 0., 0., 0., 0., 0.],
                                                          [0.33333333, 0.57142857, 0.38095238, 0., 0., 0., 0., 0.],
                                                          [0., 0.47619048, 0.52380952, 0.28571429, 0.04761905, 0., 0.,
                                                           0.],
                                                          [0., 0., 0.42857143, 0.85714286, 0.14285714, 0., 0., 0.]],
                                                         [[1., 0.28571429, 0., 0., 0., 0., 0., 0.],
                                                          [0.33333333, 0.57142857, 0.38095238, 0., 0., 0., 0., 0.],
                                                          [0., 0.47619048, 0.52380952, 0.28571429, 0.04761905, 0., 0.,
                                                           0.],
                                                          [0., 0., 0.42857143, 0.85714286, 0.14285714, 0., 0., 0.]]])),
            'lat': np.linspace(-67.5, 67.5, 4),
            'lon': np.linspace(-157.5, 157.5, 8),
            'time': np.array([1, 2])})
        assert_almost_equal(ds_coarse_resampled['first'].values, expected['first'].values)

        # Test that the fine dataset has been resampled (aggregated)
        # onto the grid of the coarse dataset.
        ds_fine_resampled = reg_op(ds_master=ds_coarse, ds_replica=ds_fine)
        expected = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.array([[[0.625, 0.125, 0., 0., 0., 0.],
                                                         [0.125, 0.5, 0.125, 0., 0., 0.],
                                                         [0., 0.125, 0.625, 0., 0., 0.]],

                                                        [[0.625, 0.125, 0., 0., 0., 0.],
                                                         [0.125, 0.5, 0.125, 0., 0., 0.],
                                                         [0., 0.125, 0.625, 0., 0., 0.]]])),
            'second': (['time', 'lat', 'lon'], np.array([[[0.625, 0.125, 0., 0., 0., 0.],
                                                          [0.125, 0.5, 0.125, 0., 0., 0.],
                                                          [0., 0.125, 0.625, 0., 0., 0.]],

                                                         [[0.625, 0.125, 0., 0., 0., 0.],
                                                          [0.125, 0.5, 0.125, 0., 0., 0.],
                                                          [0., 0.125, 0.625, 0., 0., 0.]]])),
            'lat': np.linspace(-60, 60, 3),
            'lon': np.linspace(-150, 150, 6),
            'time': np.array([1, 2])})

        assert_almost_equal(ds_fine_resampled['first'].values, expected['first'].values)

    def test_error(self):
        """
        Test error conditions
        """
        # Test unexpected global bounds
        ds_fine = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.array([np.eye(4, 8), np.eye(4, 8)])),
            'second': (['time', 'lat', 'lon'], np.array([np.eye(4, 8), np.eye(4, 8)])),
            'lat': np.linspace(67.5, 135, 4),
            'lon': np.linspace(-157.5, 157.5, 8),
            'time': np.array([1, 2])})

        ds_coarse = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.array([np.eye(3, 6), np.eye(3, 6)])),
            'second': (['time', 'lat', 'lon'], np.array([np.eye(3, 6), np.eye(3, 6)])),
            'lat': np.linspace(-60, 60, 3),
            'lon': np.linspace(-150, 150, 6),
            'time': np.array([1, 2])})

        with self.assertRaises(ValueError) as err:
            coregister(ds_fine, ds_coarse)
        self.assertIn('(67.5, 135.0)', str(err.exception))

        # Test non-equidistant dataset
        ds_fine = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.array([np.eye(4, 8), np.eye(4, 8)])),
            'second': (['time', 'lat', 'lon'], np.array([np.eye(4, 8), np.eye(4, 8)])),
            'lat': [-67.5, -20, 20, 67.5],
            'lon': np.linspace(-157.5, 157.5, 8),
            'time': np.array([1, 2])})

        ds_coarse = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.array([np.eye(3, 6), np.eye(3, 6)])),
            'second': (['time', 'lat', 'lon'], np.array([np.eye(3, 6), np.eye(3, 6)])),
            'lat': np.linspace(-60, 60, 3),
            'lon': np.linspace(-150, 150, 6),
            'time': np.array([1, 2])})

        with self.assertRaises(ValueError) as err:
            coregister(ds_fine, ds_coarse)
        self.assertIn('not equidistant', str(err.exception))

        # Test non-pixel registered dataset
        ds_fine = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.array([np.eye(4, 8), np.eye(4, 8)])),
            'second': (['time', 'lat', 'lon'], np.array([np.eye(4, 8), np.eye(4, 8)])),
            'lat': np.linspace(-67.5, 67.5, 4),
            'lon': np.linspace(-157.5, 157.5, 8),
            'time': np.array([1, 2])})

        ds_coarse_err = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.zeros([2, 5, 10])),
            'second': (['time', 'lat', 'lon'], np.zeros([2, 5, 10])),
            'lat': np.linspace(-90, 90, 5),
            'lon': np.linspace(-162, 162, 10),
            'time': np.array([1, 2])})

        with self.assertRaises(ValueError) as err:
            coregister(ds_fine, ds_coarse_err)
        self.assertIn('not pixel-registered', str(err.exception))

        ds_coarse_err = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.zeros([2, 5, 10])),
            'second': (['time', 'lat', 'lon'], np.zeros([2, 5, 10])),
            'lat': np.linspace(-72, 72, 5),
            'lon': np.linspace(-180, 180, 10),
            'time': np.array([1, 2])})

        with self.assertRaises(ValueError) as err:
            coregister(ds_fine, ds_coarse_err)
        self.assertIn('not pixel-registered', str(err.exception))

        # Test unexpected dimensionality
        ds_fine = xr.Dataset({
            'first': (['lat', 'longertude'], np.eye(4, 8)),
            'second': (['lat', 'longertude'], np.eye(4, 8)),
            'lat': np.linspace(-67.5, 67.5, 4),
            'longertude': np.linspace(-157.5, 157.5, 8)})

        ds_coarse = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.array([np.eye(3, 6), np.eye(3, 6)])),
            'second': (['time', 'lat', 'lon'], np.array([np.eye(3, 6), np.eye(3, 6)])),
            'lat': np.linspace(-60, 60, 3),
            'lon': np.linspace(-150, 150, 6),
            'time': np.array([1, 2])})

        with self.assertRaises(ValueError) as err:
            coregister(ds_fine, ds_coarse)
        self.assertIn('longertude', str(err.exception))

        ds_fine = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.array([np.eye(4, 8), np.eye(4, 8)])),
            'second': (['time', 'lat', 'lon'], np.array([np.eye(4, 8), np.eye(4, 8)])),
            'lat': np.linspace(-67.5, 67.5, 4),
            'lon': np.linspace(-157.5, 157.5, 8),
            'time': np.array([1, 2])})

        ds_coarse = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.array([np.eye(3, 6), np.eye(3, 6)])),
            'second': (['time', 'lon'], np.eye(2, 6)),
            'lat': np.linspace(-60, 60, 3),
            'lon': np.linspace(-150, 150, 6),
            'time': np.array([1, 2])})

        with self.assertRaises(ValueError) as err:
            coregister(ds_fine, ds_coarse)
        self.assertIn('select_var', str(err.exception))

    def test_find_intersection(self):
        """
        Test the _find_intersection method
        """
        # Test =======
        #          =========
        a = np.linspace(0.5, 9.5, 10)
        b = np.linspace(5.5, 14.5, 10)
        result = _find_intersection(a, b, (0, 15))
        self.assertEqual((5, 10), result)

        # Test   =======
        #    =========
        a = np.linspace(0.5, 9.5, 10)
        b = np.linspace(5.5, 14.5, 10)
        result = _find_intersection(b, a, (0, 15))
        self.assertEqual((5, 10), result)

        # Test   =======
        #     ==============
        a = np.linspace(5.5, 14.5, 10)
        b = np.linspace(0.5, 19.5, 20)
        result = _find_intersection(a, b, (0, 20))
        self.assertEqual((5, 15), result)

        # Test ==================
        #          ========
        a = np.linspace(0.5, 19.5, 20)
        b = np.linspace(5.5, 14.5, 10)
        result = _find_intersection(a, b, (0, 20))
        self.assertEqual((5, 15), result)

        # Test ============
        #                    ========
        a = np.linspace(0.5, 9.5, 10)
        b = np.linspace(10.5, 19.5, 10)
        with self.assertRaises(ValueError) as err:
            _find_intersection(a, b, (0, 20))
        self.assertIn('valid intersection', str(err.exception))

        # Test       ============
        #  ========
        a = np.linspace(0.5, 9.5, 10)
        b = np.linspace(10.5, 19.5, 10)
        with self.assertRaises(ValueError) as err:
            _find_intersection(b, a, (0, 20))
        self.assertIn('valid intersection', str(err.exception))

        # Test misaligned origins
        a = np.linspace(0.5, 9.5, 10)
        b = np.linspace(1, 9, 10)
        with self.assertRaises(ValueError) as err:
            _find_intersection(a, b, (0, 10))
        self.assertIn('valid intersection', str(err.exception))

        # Test differing pixel sizes
        a = np.linspace(0.5, 9.5, 10)
        b = np.linspace(5.25, 14.75, 20)
        result = _find_intersection(b, a, (0, 20))
        self.assertEqual((5, 10), result)

    def test_subset(self):
        """
        Test coregistration being run on a subset
        """
        ds_fine = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.array([np.eye(4, 8), np.eye(4, 8)])),
            'second': (['time', 'lat', 'lon'], np.array([np.eye(4, 8), np.eye(4, 8)])),
            'lat': np.linspace(-67.5, 67.5, 4),
            'lon': np.linspace(-157.5, 157.5, 8),
            'time': np.array([1, 2])})

        ds_coarse = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.array([np.eye(3, 6), np.eye(3, 6)])),
            'second': (['time', 'lat', 'lon'], np.array([np.eye(3, 6), np.eye(3, 6)])),
            'lat': np.linspace(-60, 60, 3),
            'lon': np.linspace(-150, 150, 6),
            'time': np.array([1, 2])})

        lat_slice = slice(-70, 70)
        lon_slice = slice(-40, 40)
        ds_coarse = ds_coarse.sel(lat=lat_slice, lon=lon_slice)

        # Test that the coarse dataset has been resampled onto the grid
        # of the finer dataset.
        ds_coarse_resampled = coregister(ds_fine, ds_coarse)
        assert_array_equal([-67.5, -22.5, 22.5, 67.5], ds_coarse_resampled.lat.values)
        assert_array_equal([-22.5, 22.5],
                           ds_coarse_resampled.lon.values)

        # Check if the geospatial attributes have been correctly set
        self.assertEqual(ds_coarse_resampled.lat.values[0] - 45 * 0.5,
                         ds_coarse_resampled.attrs['geospatial_lat_min'])
        self.assertEqual(ds_coarse_resampled.lat.values[-1] + 45 * 0.5,
                         ds_coarse_resampled.attrs['geospatial_lat_max'])
        self.assertEqual(ds_coarse_resampled.lon.values[0] - 45 * 0.5,
                         ds_coarse_resampled.attrs['geospatial_lon_min'])
        self.assertEqual(ds_coarse_resampled.lon.values[-1] + 45 * 0.5,
                         ds_coarse_resampled.attrs['geospatial_lon_max'])
        self.assertEqual(45.0,
                         ds_coarse_resampled.attrs['geospatial_lat_resolution'])
        self.assertEqual(45.0,
                         ds_coarse_resampled.attrs['geospatial_lon_resolution'])

    def test_recursive(self):
        """
        Test coregistration with more dimensions than lat/lon/time
        """
        slice_fine = np.eye(4, 8)
        slice_coarse = np.eye(3, 6)
        ndarr_fine = np.zeros([2, 2, 2, 4, 8])
        ndarr_coarse = np.zeros([2, 2, 2, 3, 6])
        ndarr_fine_l1 = np.zeros([2, 2, 4, 8])
        ndarr_coarse_l1 = np.zeros([2, 2, 3, 6])
        ndarr_fine_l2 = np.zeros([2, 2, 4, 8])
        ndarr_coarse_l2 = np.zeros([2, 2, 3, 6])
        ndarr_fine[:] = slice_fine
        ndarr_coarse[:] = slice_coarse
        ndarr_fine_l1[:] = slice_fine
        ndarr_coarse_l1[:] = slice_coarse
        ndarr_fine_l2[:] = slice_fine
        ndarr_coarse_l2[:] = slice_coarse

        ds_fine = xr.Dataset({
            'first': (['time', 'layer', 'layer2', 'lat', 'lon'], ndarr_fine),
            'second': (['time', 'layer', 'layer2', 'lat', 'lon'], ndarr_fine),
            'lat': np.linspace(-67.5, 67.5, 4),
            'lon': np.linspace(-157.5, 157.5, 8),
            'layer': np.array([1, 2]),
            'layer2': np.array([1, 2]),
            'time': np.array([1, 2])}).chunk(chunks={'lat': 2, 'lon': 4})

        ds_coarse = xr.Dataset({
            'first': (['time', 'layer', 'layer2', 'lat', 'lon'], ndarr_coarse),
            'second': (['time', 'layer', 'layer2', 'lat', 'lon'], ndarr_coarse),
            'lat': np.linspace(-60, 60, 3),
            'lon': np.linspace(-150, 150, 6),
            'time': np.array([1, 2]),
            'layer': np.array([1, 2]),
            'layer2': np.array([1, 2])}).chunk(chunks={'lat': 3, 'lon': 3})

        # Test that the coarse dataset has been resampled onto the grid
        # of the finer dataset.
        rm = RecordingMonitor()
        ds_coarse_resampled = coregister(ds_fine, ds_coarse, monitor=rm)

        self.assertEqual([('start', 'coregister dataset', 2),
                          ('progress', 0.0, 'coregister dataarray', 0),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 0),
                          ('progress', 0.03125, None, 2),
                          ('progress', 0.03125, None, 3),
                          ('progress', 0.03125, None, 5),
                          ('progress', 0.03125, None, 6),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 6),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 6),
                          ('progress', 0.03125, None, 8),
                          ('progress', 0.03125, None, 9),
                          ('progress', 0.03125, None, 11),
                          ('progress', 0.03125, None, 13),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 13),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 13),
                          ('progress', 0.03125, None, 14),
                          ('progress', 0.03125, None, 16),
                          ('progress', 0.03125, None, 17),
                          ('progress', 0.03125, None, 19),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 19),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 19),
                          ('progress', 0.03125, None, 20),
                          ('progress', 0.03125, None, 22),
                          ('progress', 0.03125, None, 23),
                          ('progress', 0.03125, None, 25),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 25),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 25),
                          ('progress', 0.03125, None, 27),
                          ('progress', 0.03125, None, 28),
                          ('progress', 0.03125, None, 30),
                          ('progress', 0.03125, None, 31),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 31),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 31),
                          ('progress', 0.03125, None, 33),
                          ('progress', 0.03125, None, 34),
                          ('progress', 0.03125, None, 36),
                          ('progress', 0.03125, None, 38),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 38),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 38),
                          ('progress', 0.03125, None, 39),
                          ('progress', 0.03125, None, 41),
                          ('progress', 0.03125, None, 42),
                          ('progress', 0.03125, None, 44),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 44),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 44),
                          ('progress', 0.03125, None, 45),
                          ('progress', 0.03125, None, 47),
                          ('progress', 0.03125, None, 48),
                          ('progress', 0.03125, None, 50),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 50),
                          ('progress', 0.0, 'coregister dataarray', 50),
                          ('progress', 0.0, 'coregister dataarray', 50),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 50),
                          ('progress', 0.03125, None, 52),
                          ('progress', 0.03125, None, 53),
                          ('progress', 0.03125, None, 55),
                          ('progress', 0.03125, None, 56),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 56),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 56),
                          ('progress', 0.03125, None, 58),
                          ('progress', 0.03125, None, 59),
                          ('progress', 0.03125, None, 61),
                          ('progress', 0.03125, None, 63),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 63),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 63),
                          ('progress', 0.03125, None, 64),
                          ('progress', 0.03125, None, 66),
                          ('progress', 0.03125, None, 67),
                          ('progress', 0.03125, None, 69),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 69),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 69),
                          ('progress', 0.03125, None, 70),
                          ('progress', 0.03125, None, 72),
                          ('progress', 0.03125, None, 73),
                          ('progress', 0.03125, None, 75),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 75),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 75),
                          ('progress', 0.03125, None, 77),
                          ('progress', 0.03125, None, 78),
                          ('progress', 0.03125, None, 80),
                          ('progress', 0.03125, None, 81),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 81),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 81),
                          ('progress', 0.03125, None, 83),
                          ('progress', 0.03125, None, 84),
                          ('progress', 0.03125, None, 86),
                          ('progress', 0.03125, None, 88),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 88),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 88),
                          ('progress', 0.03125, None, 89),
                          ('progress', 0.03125, None, 91),
                          ('progress', 0.03125, None, 92),
                          ('progress', 0.03125, None, 94),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 94),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 94),
                          ('progress', 0.03125, None, 95),
                          ('progress', 0.03125, None, 97),
                          ('progress', 0.03125, None, 98),
                          ('progress', 0.03125, None, 100),
                          ('progress', 0.0, 'coregister dataarray: resample slice', 100),
                          ('progress', 0.0, 'coregister dataarray', 100),
                          ('done',)], rm.records)

        slice_exp = np.array([[1., 0.28571429, 0., 0., 0., 0., 0., 0.],
                              [0.33333333, 0.57142857, 0.38095238, 0., 0., 0., 0., 0.],
                              [0., 0.47619048, 0.52380952, 0.28571429, 0.04761905, 0., 0., 0.],
                              [0., 0., 0.42857143, 0.85714286, 0.14285714, 0., 0., 0.]])
        ndarr_fine_exp = np.zeros([2, 2, 2, 4, 8])
        ndarr_fine_exp[:] = slice_exp

        expected = xr.Dataset({
            'first': (['time', 'layer', 'layer2', 'lat', 'lon'], ndarr_fine_exp),
            'second': (['time', 'layer', 'layer2', 'lat', 'lon'], ndarr_fine_exp),
            'layer': np.array([1, 2]),
            'layer2': np.array([1, 2]),
            'time': np.array([1, 2])})
        assert_almost_equal(ds_coarse_resampled['first'].values, expected['first'].values)

        # Test that the fine dataset has been resampled (aggregated)
        # onto the grid of the coarse dataset.
        ds_fine_resampled = coregister(ds_coarse, ds_fine)

        slice_exp = np.array([[0.625, 0.125, 0., 0., 0., 0.],
                              [0.125, 0.5, 0.125, 0., 0., 0.],
                              [0., 0.125, 0.625, 0., 0., 0.]])
        ndarr_coarse_exp = np.zeros([2, 2, 2, 3, 6])
        ndarr_coarse_exp[:] = slice_exp

        expected = xr.Dataset({
            'first': (['time', 'layer', 'layer2', 'lat', 'lon'], ndarr_coarse_exp),
            'second': (['time', 'layer', 'layer2', 'lat', 'lon'], ndarr_coarse_exp),
            'lat': np.linspace(-60, 60, 3),
            'lon': np.linspace(-150, 150, 6),
            'layer': np.array([1, 2]),
            'layer2': np.array([1, 2]),
            'time': np.array([1, 2])})

        assert_almost_equal(ds_fine_resampled['first'].values, expected['first'].values)

        # Test that coregistering with data arrays with less than all possible
        # dimensions works
        ds_fine = xr.Dataset({
            'first': (['time', 'layer', 'lat', 'lon'], ndarr_fine_l1),
            'second': (['time', 'layer2', 'lat', 'lon'], ndarr_fine_l2),
            'lat': np.linspace(-67.5, 67.5, 4),
            'lon': np.linspace(-157.5, 157.5, 8),
            'layer': np.array([1, 2]),
            'layer2': np.array([1, 2]),
            'time': np.array([1, 2])}).chunk(chunks={'lat': 2, 'lon': 4})

        ds_coarse = xr.Dataset({
            'first': (['time', 'layer', 'lat', 'lon'], ndarr_coarse_l1),
            'second': (['time', 'layer2', 'lat', 'lon'], ndarr_coarse_l2),
            'lat': np.linspace(-60, 60, 3),
            'lon': np.linspace(-150, 150, 6),
            'time': np.array([1, 2]),
            'layer': np.array([1, 2]),
            'layer2': np.array([1, 2])}).chunk(chunks={'lat': 3, 'lon': 3})

        ds_fine_resampled = coregister(ds_coarse, ds_fine)
        ndarr_coarse_exp = np.zeros([2, 2, 3, 6])
        ndarr_coarse_exp[:] = slice_exp

        expected = xr.Dataset({
            'first': (['time', 'layer', 'lat', 'lon'], ndarr_coarse_exp),
            'second': (['time', 'layer2', 'lat', 'lon'], ndarr_coarse_exp),
            'lat': np.linspace(-60, 60, 3),
            'lon': np.linspace(-150, 150, 6),
            'layer': np.array([1, 2]),
            'layer2': np.array([1, 2]),
            'time': np.array([1, 2])})

        assert_almost_equal(ds_fine_resampled['first'].values, expected['first'].values)

    def test_2D(self):
        """
        Test a case where a 2D lat/lon dataset is resampled or used for
        resampling
        """
        # Master dataset is 2D
        ds_fine = xr.Dataset({
            'first': (['lat', 'lon'], np.eye(4, 8)),
            'lat': np.linspace(-67.5, 67.5, 4),
            'lon': np.linspace(-157.5, 157.5, 8)}).chunk()

        ds_coarse = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.array([np.eye(3, 6), np.eye(3, 6)])),
            'second': (['time', 'lat', 'lon'], np.array([np.eye(3, 6), np.eye(3, 6)])),
            'lat': np.linspace(-60, 60, 3),
            'lon': np.linspace(-150, 150, 6),
            'time': np.array([1, 2])}).chunk(chunks={'lat': 3, 'lon': 3})

        ds_coarse_resampled = coregister(ds_fine, ds_coarse)

        slice_exp = np.array([[1., 0.28571429, 0., 0., 0., 0., 0., 0.],
                              [0.33333333, 0.57142857, 0.38095238, 0., 0., 0., 0., 0.],
                              [0., 0.47619048, 0.52380952, 0.28571429, 0.04761905, 0., 0., 0.],
                              [0., 0., 0.42857143, 0.85714286, 0.14285714, 0., 0., 0.]])
        exp_arr = np.zeros([2, 4, 8])
        exp_arr[:] = slice_exp

        expected = xr.Dataset({
            'first': (['time', 'lat', 'lon'], exp_arr),
            'second': (['time', 'lat', 'lon'], exp_arr),
            'lat': np.linspace(-67.5, 67.5, 4),
            'lon': np.linspace(-157.5, 157.5, 8),
            'time': np.array([1, 2])})

        assert_almost_equal(ds_coarse_resampled['first'].values, expected['first'].values)

        # replica dataset contains a 2D variable
        ds_coarse = xr.Dataset({
            'first': (['lat', 'lon'], np.eye(3, 6)),
            'second': (['time', 'lat', 'lon'], np.array([np.eye(3, 6), np.eye(3, 6)])),
            'lat': np.linspace(-60, 60, 3),
            'lon': np.linspace(-150, 150, 6),
            'time': np.array([1, 2])}).chunk(chunks={'lat': 3, 'lon': 3})

        ds_coarse_resampled = coregister(ds_fine, ds_coarse)

        assert_almost_equal(ds_coarse_resampled['first'].values, slice_exp)

    def test_int_array(self):
        """
        Test coregistration on integer arrays
        """
        ds_fine = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.array([np.eye(4, 8), np.eye(4, 8)], dtype='int32')),
            'second': (['time', 'lat', 'lon'], np.array([np.eye(4, 8), np.eye(4, 8,)])),
            'lat': np.linspace(-67.5, 67.5, 4),
            'lon': np.linspace(-157.5, 157.5, 8),
            'time': np.array([1, 2])}).chunk(chunks={'lat': 2, 'lon': 4})

        ds_coarse = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.array([np.eye(3, 6), np.eye(3, 6)], dtype='int32')),
            'second': (['time', 'lat', 'lon'], np.array([np.eye(3, 6), np.eye(3, 6)])),
            'lat': np.linspace(-60, 60, 3),
            'lon': np.linspace(-150, 150, 6),
            'time': np.array([1, 2])}).chunk(chunks={'lat': 3, 'lon': 3})

        # Test that the coarse dataset has been resampled onto the grid
        # of the finer dataset.
        ds_coarse_resampled = coregister(ds_fine, ds_coarse, method_us='nearest')

        expected = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.array([[[1, 1, 0, 0, 0, 0, 0, 0],
                                                         [1, 1, 0, 0, 0, 0, 0, 0],
                                                         [0, 0, 1, 0, 0, 0, 0, 0],
                                                         [0, 0, 0, 1, 0, 0, 0, 0]],
                                                        [[1, 1, 0, 0, 0, 0, 0, 0],
                                                         [1, 1, 0, 0, 0, 0, 0, 0],
                                                         [0, 0, 1, 0, 0, 0, 0, 0],
                                                         [0, 0, 0, 1, 0, 0, 0, 0]]])),
            'second': (['time', 'lat', 'lon'], np.array([[[1, 1, 0, 0, 0, 0, 0, 0],
                                                         [1, 1, 0, 0, 0, 0, 0, 0],
                                                         [0, 0, 1, 0, 0, 0, 0, 0],
                                                         [0, 0, 0, 1, 0, 0, 0, 0]],
                                                         [[1, 1, 0, 0, 0, 0, 0, 0],
                                                         [1, 1, 0, 0, 0, 0, 0, 0],
                                                         [0, 0, 1, 0, 0, 0, 0, 0],
                                                         [0, 0, 0, 1, 0, 0, 0, 0]]])),
            'lat': np.linspace(-67.5, 67.5, 4),
            'lon': np.linspace(-157.5, 157.5, 8),
            'time': np.array([1, 2])})
        assert_almost_equal(ds_coarse_resampled['first'].values, expected['first'].values)

        # Test that the fine dataset has been resampled (aggregated)
        # onto the grid of the coarse dataset.
        ds_fine_resampled = coregister(ds_coarse, ds_fine, method_ds='mode')
        expected = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.array([[[1, 0, 0, 0, 0, 0],
                                                         [0, 1, 0, 0, 0, 0],
                                                         [0, 0, 1, 0, 0, 0]],

                                                        [[1, 0, 0, 0, 0, 0],
                                                         [0, 1, 0, 0, 0, 0],
                                                         [0, 0, 1, 0, 0, 0]]])),
            'second': (['time', 'lat', 'lon'], np.array([[[1, 0, 0, 0, 0, 0],
                                                         [0, 1, 0, 0, 0, 0],
                                                         [0, 0, 1, 0, 0, 0]],

                                                        [[1, 0, 0, 0, 0, 0],
                                                         [0, 1, 0, 0, 0, 0],
                                                         [0, 0, 1, 0, 0, 0]]])),
            'lat': np.linspace(-60, 60, 3),
            'lon': np.linspace(-150, 150, 6),
            'time': np.array([1, 2])})

        assert_almost_equal(ds_fine_resampled['first'].values, expected['first'].values)

    def test_same_grid(self):
        """
        Test the case when both datasets already have the same geospatial definition
        """
        ds_fine = xr.Dataset({
            'first': (['time', 'lat', 'lon'], np.array([np.eye(4, 8), np.eye(4, 8)])),
            'second': (['time', 'lat', 'lon'], np.array([np.eye(4, 8), np.eye(4, 8)])),
            'lat': np.linspace(-67.5, 67.5, 4),
            'lon': np.linspace(-157.5, 157.5, 8),
            'time': np.array([1, 2])}).chunk(chunks={'lat': 2, 'lon': 4})

        rm = RecordingMonitor()
        ds_same = coregister(ds_fine, ds_fine, monitor=rm)
        # Make sure it returned the input as opposed to going through with
        # coregistration
        self.assertEqual([], rm.records)

        assert_almost_equal(ds_same['first'].values, ds_fine['first'].values)

        # Test that a subset is performed, but no coregistration done
        lat_slice = slice(-70, 70)
        lon_slice = slice(-40, 40)
        ds_subset = ds_fine.sel(lat=lat_slice, lon=lon_slice)

        ds_coreg = coregister(ds_subset, ds_fine, monitor=rm)
        self.assertEqual([], rm.records)
        assert_almost_equal(ds_coreg['first'].values, ds_subset['first'].values)


""" Tools for doing common subexpression elimination.
"""
from __future__ import print_function, division

import difflib

from sympy.core import Basic, Mul, Add, sympify
from sympy.core.basic import preorder_traversal
from sympy.core.function import _coeff_isneg
from sympy.core.compatibility import iterable, xrange
from sympy.utilities.iterables import numbered_symbols, \
    sift, topological_sort, ordered

from . import cse_opts

# (preprocessor, postprocessor) pairs which are commonly useful. They should
# each take a sympy expression and return a possibly transformed expression.
# When used in the function ``cse()``, the target expressions will be transformed
# by each of the preprocessor functions in order. After the common
# subexpressions are eliminated, each resulting expression will have the
# postprocessor functions transform them in *reverse* order in order to undo the
# transformation if necessary. This allows the algorithm to operate on
# a representation of the expressions that allows for more optimization
# opportunities.
# ``None`` can be used to specify no transformation for either the preprocessor or
# postprocessor.

cse_optimizations = list(cse_opts.default_optimizations)

# sometimes we want the output in a different format; non-trivial
# transformations can be put here for users
# ===============================================================

def reps_toposort(r):
    """Sort replacements `r` so (k1, v1) appears before (k2, v2)
    if k2 is in v1's free symbols. This orders items in the
    way that cse returns its results (hence, in order to use the
    replacements in a substitution option it would make sense
    to reverse the order).

    Examples
    ========
    >>> from sympy.simplify.cse_main import reps_toposort
    >>> from sympy.abc import x, y
    >>> from sympy import Eq
    >>> for l, r in reps_toposort([(x, y + 1), (y, 2)]):
    ...     print(Eq(l, r))
    ...
    y == 2
    x == y + 1

    """
    r = sympify(r)
    E = []
    for c1, (k1, v1) in enumerate(r):
        for c2, (k2, v2) in enumerate(r):
            if k1 in v2.free_symbols:
                E.append((c1, c2))
    return [r[i] for i in topological_sort((range(len(r)), E))]

def cse_separate(r, e):
    """Move expressions that are in the form (symbol, expr) out of the
    expressions and sort them into the replacements using the reps_toposort.

    Examples
    ========
    >>> from sympy.simplify.cse_main import cse_separate
    >>> from sympy.abc import x, y, z
    >>> from sympy import cos, exp, cse, Eq, symbols
    >>> x0, x1 = symbols('x:2')
    >>> eq = (x + 1 + exp((x + 1)/(y + 1)) + cos(y + 1))
    >>> cse([eq, Eq(x, z + 1), z - 2], postprocess=cse_separate) in [
    ... [[(x0, y + 1), (x, z + 1), (x1, x + 1)],
    ...  [x1 + exp(x1/x0) + cos(x0), z - 2]],
    ... [[(x1, y + 1), (x, z + 1), (x0, x + 1)],
    ...  [x0 + exp(x0/x1) + cos(x1), z - 2]]]
    ...
    True
    """
    d = sift(e, lambda w: w.is_Equality and w.lhs.is_Symbol)
    r = r + [w.args for w in d[True]]
    e = d[False]
    return [reps_toposort(r), e]

# ====end of cse postprocess idioms===========================

def preprocess_for_cse(expr, optimizations):
    """ Preprocess an expression to optimize for common subexpression
    elimination.

    Parameters
    ----------
    expr : sympy expression
        The target expression to optimize.
    optimizations : list of (callable, callable) pairs
        The (preprocessor, postprocessor) pairs.

    Returns
    -------
    expr : sympy expression
        The transformed expression.
    """
    for pre, post in optimizations:
        if pre is not None:
            expr = pre(expr)
    return expr

def postprocess_for_cse(expr, optimizations):
    """ Postprocess an expression after common subexpression elimination to
    return the expression to canonical sympy form.

    Parameters
    ----------
    expr : sympy expression
        The target expression to transform.
    optimizations : list of (callable, callable) pairs, optional
        The (preprocessor, postprocessor) pairs.  The postprocessors will be
        applied in reversed order to undo the effects of the preprocessors
        correctly.

    Returns
    -------
    expr : sympy expression
        The transformed expression.
    """
    if optimizations is None:
        optimizations = cse_optimizations
    for pre, post in reversed(optimizations):
        if post is not None:
            expr = post(expr)
    return expr

def _remove_singletons(reps, exprs):
    """
    Helper function for cse that will remove expressions that weren't
    used more than once.
    """
    u_reps = []  # the useful reps that are used more than once
    for i, ui in enumerate(reps):
        used = []  # where it was used
        ri, ei = ui

        # keep track of whether the substitution was used more
        # than once. If used is None, it was never used (yet);
        # if used is an int, that is the last place where it was
        # used (>=0 in the reps, <0 in the expressions) and if
        # it is True, it was used more than once.

        used = None

        tot = 0  # total times used so far

        # search through the reps
        for j in range(i + 1, len(reps)):
            c = reps[j][1].count(ri)
            if c:
                tot += c
                if tot > 1:
                    u_reps.append(ui)
                    used = True
                    break
                else:
                    used = j

        if used is not True:

            # then search through the expressions

            for j, rj in enumerate(exprs):
                c = rj.count(ri)
                if c:
                    # append a negative so we know that it was in the
                    # expression that used it
                    tot += c
                    if tot > 1:
                        u_reps.append(ui)
                        used = True
                        break
                    else:
                        used = j - len(exprs)

            if type(used) is int:

                # undo the change

                rep = {ri: ei}
                j = used
                if j < 0:
                    exprs[j] = exprs[j].subs(rep)
                else:
                    reps[j] = reps[j][0], reps[j][1].subs(rep)

    # reuse unused symbols so a contiguous range of symbols is returned

    if len(u_reps) != len(reps):
        for i, ri in enumerate(u_reps):
            if u_reps[i][0] != reps[i][0]:
                rep = (u_reps[i][0], reps[i][0])
                u_reps[i] = rep[1], u_reps[i][1].subs(*rep)
                for j in range(i + 1, len(u_reps)):
                    u_reps[j] = u_reps[j][0], u_reps[j][1].subs(*rep)
                for j, rj in enumerate(exprs):
                    exprs[j] = exprs[j].subs(*rep)

    reps[:] = u_reps  # change happens in-place

def cse(exprs, symbols=None, optimizations=None, postprocess=None):
    """ Perform common subexpression elimination on an expression.

    Parameters
    ==========

    exprs : list of sympy expressions, or a single sympy expression
        The expressions to reduce.
    symbols : infinite iterator yielding unique Symbols
        The symbols used to label the common subexpressions which are pulled
        out. The ``numbered_symbols`` generator is useful. The default is a
        stream of symbols of the form "x0", "x1", etc. This must be an infinite
        iterator.
    optimizations : list of (callable, callable) pairs, optional
        The (preprocessor, postprocessor) pairs. If not provided,
        ``sympy.simplify.cse.cse_optimizations`` is used.
    postprocess : a function which accepts the two return values of cse and
        returns the desired form of output from cse, e.g. if you want the
        replacements reversed the function might be the following lambda:
        lambda r, e: return reversed(r), e

    Returns
    =======

    replacements : list of (Symbol, expression) pairs
        All of the common subexpressions that were replaced. Subexpressions
        earlier in this list might show up in subexpressions later in this list.
    reduced_exprs : list of sympy expressions
        The reduced expressions with all of the replacements above.
    """
    from sympy.matrices import Matrix

    if symbols is None:
        symbols = numbered_symbols()
    else:
        # In case we get passed an iterable with an __iter__ method instead of
        # an actual iterator.
        symbols = iter(symbols)
    seen_subexp = set()
    muls = set()
    adds = set()
    to_eliminate = set()

    if optimizations is None:
        # Pull out the default here just in case there are some weird
        # manipulations of the module-level list in some other thread.
        optimizations = list(cse_optimizations)

    # Handle the case if just one expression was passed.
    if isinstance(exprs, Basic):
        exprs = [exprs]

    # Preprocess the expressions to give us better optimization opportunities.
    reduced_exprs = [preprocess_for_cse(e, optimizations) for e in exprs]

    # Find all of the repeated subexpressions.
    for expr in reduced_exprs:
        if not isinstance(expr, Basic):
            continue
        pt = preorder_traversal(expr)
        for subtree in pt:

            inv = 1/subtree if subtree.is_Pow else None

            if subtree.is_Atom or iterable(subtree) or inv and inv.is_Atom:
                # Exclude atoms, since there is no point in renaming them.
                continue

            if subtree in seen_subexp:
                if inv and _coeff_isneg(subtree.exp):
                    # save the form with positive exponent
                    subtree = inv
                to_eliminate.add(subtree)
                pt.skip()
                continue

            if inv and inv in seen_subexp:
                if _coeff_isneg(subtree.exp):
                    # save the form with positive exponent
                    subtree = inv
                to_eliminate.add(subtree)
                pt.skip()
                continue
            elif subtree.is_Mul:
                muls.add(subtree)
            elif subtree.is_Add:
                adds.add(subtree)

            seen_subexp.add(subtree)

    # process adds - any adds that weren't repeated might contain
    # subpatterns that are repeated, e.g. x+y+z and x+y have x+y in common
    adds = [set(a.args) for a in ordered(adds)]
    for i in xrange(len(adds)):
        for j in xrange(i + 1, len(adds)):
            com = adds[i].intersection(adds[j])
            if len(com) > 1:
                to_eliminate.add(Add(*com))

                # remove this set of symbols so it doesn't appear again
                adds[i] = adds[i].difference(com)
                adds[j] = adds[j].difference(com)
                for k in xrange(j + 1, len(adds)):
                    if not com.difference(adds[k]):
                        adds[k] = adds[k].difference(com)

    # process muls - any muls that weren't repeated might contain
    # subpatterns that are repeated, e.g. x*y*z and x*y have x*y in common

    # use SequenceMatcher on the nc part to find the longest common expression
    # in common between the two nc parts
    sm = difflib.SequenceMatcher()

    muls = [a.args_cnc(cset=True) for a in ordered(muls)]
    for i in xrange(len(muls)):
        if muls[i][1]:
            sm.set_seq1(muls[i][1])
        for j in xrange(i + 1, len(muls)):
            # the commutative part in common
            ccom = muls[i][0].intersection(muls[j][0])

            # the non-commutative part in common
            if muls[i][1] and muls[j][1]:
                # see if there is any chance of an nc match
                ncom = set(muls[i][1]).intersection(set(muls[j][1]))
                if len(ccom) + len(ncom) < 2:
                    continue

                # now work harder to find the match
                sm.set_seq2(muls[j][1])
                i1, _, n = sm.find_longest_match(0, len(muls[i][1]),
                                                 0, len(muls[j][1]))
                ncom = muls[i][1][i1:i1 + n]
            else:
                ncom = []

            com = list(ccom) + ncom
            if len(com) < 2:
                continue

            to_eliminate.add(Mul(*com))

            # remove ccom from all if there was no ncom; to update the nc part
            # would require finding the subexpr and then replacing it with a
            # dummy to keep bounding nc symbols from being identified as a
            # subexpr, e.g. removing B*C from A*B*C*D might allow A*D to be
            # identified as a subexpr which would not be right.
            if not ncom:
                muls[i][0] = muls[i][0].difference(ccom)
                for k in xrange(j, len(muls)):
                    if not ccom.difference(muls[k][0]):
                        muls[k][0] = muls[k][0].difference(ccom)

    # make to_eliminate canonical; we will prefer non-Muls to Muls
    # so select them first (non-Muls will have False for is_Mul and will
    # be first in the ordering.
    to_eliminate = list(ordered(to_eliminate, lambda _: _.is_Mul))

    # Substitute symbols for all of the repeated subexpressions.
    replacements = []
    reduced_exprs = list(reduced_exprs)
    hit = True
    for i, subtree in enumerate(to_eliminate):
        if hit:
            sym = next(symbols)
        hit = False
        if subtree.is_Pow and subtree.exp.is_Rational:
            update = lambda x: x.xreplace({subtree: sym, 1/subtree: 1/sym})
        else:
            update = lambda x: x.subs(subtree, sym)
        # Make the substitution in all of the target expressions.
        for j, expr in enumerate(reduced_exprs):
            old = reduced_exprs[j]
            reduced_exprs[j] = update(expr)
            hit = hit or (old != reduced_exprs[j])
        # Make the substitution in all of the subsequent substitutions.
        for j in range(i + 1, len(to_eliminate)):
            old = to_eliminate[j]
            to_eliminate[j] = update(to_eliminate[j])
            hit = hit or (old != to_eliminate[j])
        if hit:
            replacements.append((sym, subtree))

    # Postprocess the expressions to return the expressions to canonical form.
    for i, (sym, subtree) in enumerate(replacements):
        subtree = postprocess_for_cse(subtree, optimizations)
        replacements[i] = (sym, subtree)
    reduced_exprs = [postprocess_for_cse(e, optimizations)
        for e in reduced_exprs]

    # remove replacements that weren't used more than once
    _remove_singletons(replacements, reduced_exprs)

    if isinstance(exprs, Matrix):
        reduced_exprs = [Matrix(exprs.rows, exprs.cols, reduced_exprs)]
    if postprocess is None:
        return replacements, reduced_exprs
    return postprocess(replacements, reduced_exprs)


#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Date    : 2019-1-29 17:11:00
# @Author  : suke1900 (johnny824lee@gmail.com)
# @Link    : http://www.makcyun.top
# @Version : $Id$

"""
爬取 1983 至 2018 年 共 36 年春晚的节目趣事，包括这些内容：

谁导演春晚次数最多？

谁主持春晚次数最多？

哪两年的除夕刚好是同一天？

谁上春晚次数最多，堪称「钉子户」？

港台明星上春晚次数对比

歌曲、小品、相声类节目数量对比

"""

# import pandas as pd

# def get_data():
#     data = pd.read_csv('chinese_newyear2.csv',encoding='utf_8_sig')
#     data = data['category\tname\tactor\tyear'].str.split('\t',expand=True)
#     data.columns = ['category','name','actor','year']
#     # # data['actor'].str.split('、',expand=True)
#     # data2 = data['actor'].str.split('、',expand=True)
#     # # print(type(data2))
#     # data = pd.merge(data,pd.DataFrame(data2),how='left',left_index=True, right_index=True)
#     # # print(data.columns)
#     data.to_csv('chinese_newyear4.csv',encoding='utf_8_sig',index=0)

#     print(data)
#     # print(data.describe())
#     # return data

# get_data()

"""
列拆分合并
https://zhuanlan.zhihu.com/p/30529129

"""

import pandas as pd
import matplotlib.pyplot as plt
import pylab
import numpy as np
import matplotlib.dates as mdates
import matplotlib as mpl

plt.style.use('ggplot')
fig = plt.figure(figsize=(5,8))
colors = '#6D6D6D' #黑色
colorstitle = '#EE9922'

# 春晚官网提取五种颜色，
color1 = '#D50700'
color2 = '#DB2B08'
color3 = '#E14F10'
color4 = '#E77318'
color5 = '#EE9922'

fontsize_title = 15
fontsize_text = 10

# 分析导演、主持人 # # # # # # # # # # # # # # # # # #
def get_infodata():
    data = pd.read_csv('chinese_newyear3.csv',encoding='utf_8_sig')

    # 筛选导演主持人
    data = data[data['category'] == '导演']
    # data = data[data['category'] == '主持']
    data2 = data['content'].str.split('、',expand=True)

    # 统计出现次数
    data2 = data2.apply(pd.value_counts)
    data2['col_num'] = data2.sum(axis=1)
    data2.sort_values(by='col_num',ascending=False,inplace=True)
    data2 = data2['col_num'][:10][::-1]
    return data,data2

def analysis1(data,data2):
    data = data.set_index('year')
    data2.sort_values(ascending=False,inplace=True)
    lst = list(data2.index)[:10]
    lst_num = list(data2)[:10]

    colorsall = [color1,color2,color3,color4,color5,color1,color2,color3,color4,color5]
    for i,name in enumerate(lst):
        data3 = data['content'].str.contains(name,na=False).astype('int')
        data3 = pd.DataFrame(data3[data3.values == 1])
        data3['year'] = data3.index

        axs = fig.add_subplot(1, 10, 1+i)
        data3.plot(
            ax=axs,
            x='content',
            y = 'year',
            kind = 'scatter',
            subplots=True,
            sharey=True,
            color=colorsall[i],
            )
        new_ticks = np.linspace(1980,2020,41)

        plt.yticks(new_ticks)
        plt.tick_params(direction='in')  #标签朝里
        plt.tick_params(which='major',length=0)  # 不显示刻度标签长度
        plt.xticks([]) #去掉坐标标签
        plt.xlabel('%i次' %(lst_num[i]),fontsize=10)
        plt.xlim(0,2)

        plt.title(name,color=color5,fontsize=10)
        # plt.tight_layout(pad=3.4, w_pad=0.5, h_pad=1.0)
        fig.subplots_adjust(hspace=0,wspace=0) # 调整子图间距为0

    plt.suptitle('导演次数最多的导演 TOP 10',color=color1,fontsize=18)
    # 添加总标题
    # plt.tight_layout()
    plt.savefig('导演次数最多的导演 TOP 10.png',dpi=200)
    plt.show()

# 分析除夕日期
def get_date():
    fig = plt.figure(figsize=(5,8))
    ax = fig.add_subplot(111)
    data = pd.read_csv('chinese_newyear3.csv',encoding='utf_8_sig')
    data = data[data['category'] == '播出日期']

    data = data['content'].str.extract(r'.*?年(.*?)月(.*?)日.*?')
    data = data.reset_index(drop=True)
    data.columns = ['month','day']
    data['year'] = np.arange(1983,2019)

    # int 转 string
    data = data.applymap(str)
    data['year2'] = '1900'
    data['date'] = data['year2'].str.cat([data['month'],data['day']],sep='/')
    data = data.apply(pd.to_numeric,errors='ignore')
    data['date'] = pd.to_datetime(data['date'])

    ax.plot(
        data['date'],
        data['year'],
        )

    new_yticks = np.linspace(1980,2020,41)
    date_format = mpl.dates.DateFormatter("%m-%d")
    ax.xaxis.set_major_formatter(date_format)

    plt.yticks(new_yticks)
    plt.tick_params(direction='in')  #标签朝里

    content = list(zip(data['date'],data['year']))
    # print(content)

    # for x, y in content:
    #     x2 = '%s' %x.strftime('%m/%d') # 只显示月日格式
    #     # print(x,'\n',y)
    #     plt.text(x, y+0.2,x2, ha='center', color=color4)

    # plt.title('历年农历除夕日期变化',color=color4,fontsize=14)
    # plt.tight_layout()
    # plt.savefig('除夕日期变化.png',dpi=200,)
    # plt.show()

# 处理节目数据
def get_data():
    data = pd.read_csv('chinese_newyear.csv',encoding='utf_8_sig')
    data2 = data['actor'].str.split('、',expand=True) # 拆分人员名单
    data = pd.merge(data,pd.DataFrame(data2),how='left',left_index=True, right_index=True)
    return data

# 表演次数最多的演员TOP 20
def analysis2(data):
    cols = data.columns.size - 6

    cols = list(range(47))
    data = data[cols]

    # 统计演员出现次数
    data = data.apply(pd.value_counts)
    # 列求和计算出现总次数
    data['col_sum'] = data.sum(axis=1)

    data.sort_values(by='col_sum',ascending=False,inplace=True)
    # 取TOP20 [::-1]，反转顺序便于条形图大值在上
    data = data['col_sum'][:20][::-1]

    # data.plot(
    #     kind = 'barh',
    #     color = color1
    #     )
    # for y,x in enumerate(list(data.values)):
    #     plt.text(x+1,y-0.15,'%i'%x,ha='center',color=color4)

    # plt.title('表演次数最多的演员 TOP 20',color=colorstitle,fontweight='bold',fontsize=fontsize_title)

    # plt.tight_layout()
    # plt.xticks([])
    # plt.savefig('表演次数最多的演员TOP20.png',dpi=200)
    # plt.show()

    return data

# 各类节目数量对比
def analysis3(data):
    fig = plt.figure(figsize=(8,5))
    num_all = data.shape[0]
    # 歌曲节目数量
    num_song = data[data['category'].str.contains('歌|尾|开场')].shape[0]
    # 小品数量
    num_sketch = data[data['category'].str.contains('小品')].shape[0]
    # 相声数量
    num_crosstalk = data[data['category'].str.contains('相声')].shape[0]

    # 其他节目数量
    other = num_all - sum([num_song,num_sketch,num_crosstalk])
    lst = [num_song,num_sketch,num_crosstalk,other]

    sizes = [num_song,other,num_crosstalk,num_sketch]
    labels = ['歌曲','其他','相声','小品']
    colors_pie = [color1,color4,color3,color2]
    explode = [0.05,0,0,0]
    plt.pie(
        sizes,
        autopct='%.1f%%',
        labels=labels,
        colors=colors_pie,
        shadow=False,
        startangle=270,
        explode=explode,
        textprops={'fontsize':14,'color':colors}
        )
    plt.title('1983-2018 共 36 年春晚节目类型数量比较',color=colorstitle,fontsize=fontsize_title)

    plt.tight_layout()
    plt.axis('equal')
    plt.axis('off')
    plt.legend(loc='upper right')
    plt.savefig('1983-2018 共 36 年间各类节目类型数量比较.png',dpi=200)

    plt.show()

# 演员 TOP10 出演年数分布
def analysis4(data,data1):
    data = data.set_index('year')
    data1.sort_values(ascending=False,inplace=True)
    lst = list(data1.index)[:10]
    lst_num = list(data1)[:10]

    colorsall = [color1,color2,color3,color4,color5,color1,color2,color3,color4,color5]

    # for 循环绘制子图
    for i,name in enumerate(lst):
        data2 = data['actor'].str.contains(name,na=False).astype('int')

        data2 = pd.DataFrame(data2[data2.values == 1])
        data2['year'] = data2.index
        data2 = data2.drop_duplicates(subset=['year'],keep='last')

        axs = fig.add_subplot(1, 10, 1+i)
        data2.plot(
            ax=axs,
            x='actor',
            y = 'year',
            kind = 'scatter',
            subplots=True,
            sharey=True,
            color=colorsall[i],
            )

        new_ticks = np.linspace(1980,2020,41)
        plt.yticks(new_ticks) # 完整显示所有年份
        plt.tick_params(direction='in')  #标签朝里
        plt.tick_params(which='major',length=0) # 标签长度为0
        plt.xticks([]) #去掉坐标标签
        plt.xlabel('%i次' %(lst_num[i]),fontsize=10)
        plt.xlim(0,2)
        plt.title(name,color=color5,fontsize=10)

        # plt.tight_layout()
        fig.subplots_adjust(hspace=0,wspace=0) # 调整子图间距为0

    plt.suptitle('表演次数最多的演员 TOP 10',color=color1,fontsize=18)
    plt.savefig('演员表演年份.png',dpi=200)

    plt.show()

# 港台演员出演年数分布
def analysis5(data):
    data = data.set_index('year')
    lst = ['周杰伦','成龙','刘德华','王力宏','王菲','郭富城','陈奕迅','林俊杰','林志玲','黎明',]
    lst_num = []
    for i in lst:
        data2 = data[data['actor'].str.contains(i,na=False)].shape[0]
        lst_num.append(data2)

    colorsall = [color1,color2,color3,color4,color5,color1,color2,color3,color4,color5]

    for i,name in enumerate(lst):
        data2 = data['actor'].str.contains(name,na=False).astype('int')

        data2 = pd.DataFrame(data2[data2.values == 1])
        data2['year'] = data2.index
        data2 = data2.drop_duplicates(subset=['year'],keep='last')

        axs = fig.add_subplot(1, 10, 1+i)
        data2.plot(
            ax=axs,
            x='actor',
            y = 'year',
            kind = 'scatter',
            subplots=True,
            sharey=True,
            color=colorsall[i],
            )
        new_ticks = np.linspace(1980,2020,41)
        plt.yticks(new_ticks)
        plt.tick_params(direction='in')  #标签朝里
        plt.tick_params(which='major',length=0) # 标签长度为0
        plt.xticks([]) #去掉坐标标签
        plt.xlabel('%i次' %(lst_num[i]),fontsize=10)
        plt.xlim(0,2)

        plt.title(name,color=color5,fontsize=10)
        # plt.tight_layout()
        fig.subplots_adjust(hspace=0,wspace=0) # 调整子图间距为0

    plt.suptitle('十大港台演员出演次数',color=color1,fontsize=18)
    plt.savefig('港台演员表演次数对比.png',dpi=200)
    plt.show()

if __name__ == '__main__':
    # # # 导演主持除夕时间
    # data,data2 = get_infodata()
    # analysis1(data,data2)

    # # 除夕日期
    # get_date()

    # # 节目表处理
    data = get_data()

    # 表演次数最多的演员TOP 20
    # analysis2(data)

    # 各类节目数量对比
    # analysis3(data)

    # # TOP 10 演员表演年份
    # data1 = analysis2(data)
    # analysis4(data,data1)

    # # 港台演员出演年数分布
    analysis5(data)


#
# Copyright 2019 Gianluca Frison, Dimitris Kouzoupis, Robin Verschueren,
# Andrea Zanelli, Niels van Duijkeren, Jonathan Frey, Tommaso Sartor,
# Branimir Novoselnik, Rien Quirynen, Rezart Qelibari, Dang Doan,
# Jonas Koenemann, Yutao Chen, Tobias Schöls, Jonas Schlagenhauf, Moritz Diehl
#
# This file is part of acados.
#
# The 2-Clause BSD License
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.;
#

from acados_template import *
import acados_template as at
from export_ode_model import *
import numpy as np
import scipy.linalg
from ctypes import *

# create render arguments
ocp = acados_ocp_nlp()

# export model
model = export_ode_model()

# set model_name
ocp.model_name = model.name

Tf = 2.0
nx = model.x.size()[0]
nu = model.u.size()[0]
ny = nx + nu
ny_e = nx
N = 50

# set ocp_nlp_dimensions
nlp_dims     = ocp.dims
nlp_dims.nx  = nx
nlp_dims.ny  = ny
nlp_dims.ny_e = ny_e
nlp_dims.nbx = 0
nlp_dims.nbu = nu
nlp_dims.nu  = model.u.size()[0]
nlp_dims.N   = N

# set weighting matrices
nlp_cost = ocp.cost
Q = np.eye(4)
Q[0,0] = 1e0
Q[1,1] = 1e2
Q[2,2] = 1e-3
Q[3,3] = 1e-2

R = np.eye(1)
R[0,0] = 1e0

nlp_cost.W = scipy.linalg.block_diag(Q, R)

Vx = np.zeros((ny, nx))
Vx[0,0] = 1.0
Vx[1,1] = 1.0
Vx[2,2] = 1.0
Vx[3,3] = 1.0

nlp_cost.Vx = Vx

Vu = np.zeros((ny, nu))
Vu[4,0] = 1.0
nlp_cost.Vu = Vu

nlp_cost.W_e = Q

Vx_e = np.zeros((ny_e, nx))
Vx_e[0,0] = 1.0
Vx_e[1,1] = 1.0
Vx_e[2,2] = 1.0
Vx_e[3,3] = 1.0

nlp_cost.Vx_e = Vx_e

nlp_cost.yref  = np.zeros((ny, ))
nlp_cost.yref_e = np.zeros((ny_e, ))

# setting bounds
Fmax = 2.0
nlp_con = ocp.constraints
nlp_con.lbu = np.array([-Fmax])
nlp_con.ubu = np.array([+Fmax])
nlp_con.x0 = np.array([0.0, 3.14, 0.0, 0.0])
# nlp_con.x0 = np.array([0.0, 0.5, 0.0, 0.0])
nlp_con.idxbu = np.array([0])

# set constants
# ocp.constants['PI'] = 3.1415926535897932

# set QP solver
# ocp.solver_config.qp_solver = 'PARTIAL_CONDENSING_HPIPM'
ocp.solver_config.qp_solver = 'FULL_CONDENSING_QPOASES'
ocp.solver_config.hessian_approx = 'GAUSS_NEWTON'
ocp.solver_config.integrator_type = 'ERK'

# set prediction horizon
ocp.solver_config.tf = Tf
ocp.solver_config.nlp_solver_type = 'SQP'
# ocp.solver_config.nlp_solver_type = 'SQP_RTI'

# set header path
ocp.acados_include_path  = '/usr/local/include'
ocp.acados_lib_path      = '/usr/local/lib'

acados_solver = generate_solver(model, ocp, json_file = 'acados_ocp.json')

Nsim = 100

simX = np.ndarray((Nsim, nx))
simU = np.ndarray((Nsim, nu))

for i in range(Nsim):
    status = acados_solver.solve()

    # get solution
    x0 = acados_solver.get(0, "x")
    u0 = acados_solver.get(0, "u")

    for j in range(nx):
        simX[i,j] = x0[j]

    for j in range(nu):
        simU[i,j] = u0[j]

    # update initial condition
    x0 = acados_solver.get(1, "x")

    acados_solver.set(0, "lbx", x0)
    acados_solver.set(0, "ubx", x0)

    # update reference
    for j in range(N):
        acados_solver.set(j, "yref", np.array([0, 0, 0, 0, 0]))
    acados_solver.set(N, "yref", np.array([0, 0, 0, 0]))

# plot results
import matplotlib
import matplotlib.pyplot as plt
t = np.linspace(0.0, Tf/N, Nsim)
plt.subplot(2, 1, 1)
plt.step(t, simU, color='r')
plt.title('closed-loop simulation')
plt.ylabel('u')
plt.xlabel('t')
plt.grid(True)
plt.subplot(2, 1, 2)
plt.plot(t, simX[:,1])
plt.ylabel('theta')
plt.xlabel('t')
plt.grid(True)
plt.show()


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.11.1
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# ## Cleaning scraped game data
#
# Here we show how to use the data obtained with Scrapy. In order to use it for data analysis and game outcome predictions, we first need to clean the data.
#
# Let's start with importing the packages we'll use:

import pandas as pd # Dataframes
import numpy as np # number crunching, matrices and all that

# Let's now import the scraped data and perform a first simple cleaning step:
# 1. We import the `.csv` file to a pandas data frame.
# 2. There are games that were scraped multiple times because multiple of the selected top players were involved in them (a game might pop up in our data up to 10 times because of this). As these duplicates would skew the statistics, we remove them via `drop_duplicates`, using the starting time (`timestamp`) and the duration (`duration`) as unique identifiers.
# 3. We reset the index of the data frame, which can be done explicitly (see commented line) or implicitly when removing duplicates via `ignore_index=True`.

df = pd.read_csv('../data/raw/games.csv')
df.drop_duplicates(subset=['duration', 'timestamp', 'team_1', 'team_2', 'winner'], inplace=True, keep='first', ignore_index=True)
# df.reset_index(drop=True, inplace=True)
print(len(df))
df.head()

# As we can see in the print-out of the data frame head above, we now have unique games in `df`, with the columns `duration`, `server`, `summoner_name`, `team_1`, `team_2`, `timestamp` and `winner`. We will usually discard the server, player (summoner) and time information in our analysis.
#
# In order to capture the roles of the played champions, which currently is implicitly stored in their order in `team_1` and `team_2`, we create 10 new columns - 5 for the red and blue team each - and store the champions individually:

# These are the roles, in the same order as they are stored in team_1 and team_2.
roles = ['Top', 'Jng', 'Mid', 'Adc', 'Sup']
# For both teams...
for team_color, team_attr in zip(['B', 'R'], ['team_1', 'team_2']):
    # ...decompose the column in a data frame of champion names...
    team = df[team_attr].str.split(',', expand=True)
    # ...and for all 5 roles, store the role column in the corresponding column of df.
    for i, role in enumerate(roles):
        df[f"{team_color}{role}"] = team[i]
df.drop(columns=['team_1', 'team_2'], inplace=True)

# _Note on performance_: The above splitting of `team_1` and `team_2` is done for the entire data frame "at once" as we are using an internal pandas function (`pd.Series.str.split`) and then assign the full columns to the new role columns `BTop`, `BJng`... of `df`.
#
# Let's now rewrite the `winner` column to use `'Blue'` and `'Red'` instead of `'Team 1'` and `'Team 2'`, and drop the above mentioned columns of information we do not take into account.
#
# We also already can do a first step of data analysis and consider some stats:

df['winner'] = df.apply(lambda x: 'Blue' if x.winner=='Team 1' else 'Red', axis=1)
df.drop(['server', 'summoner_name', 'duration', 'timestamp'],axis=1,inplace=True)
# Some statistics:
num_games = len(df)   # Total number of games
num_blue_wins = len(df[df['winner']=='Blue'])     # No of games blue won
num_red_wins = len(df[df['winner']=='Red'])      # No of games red won
assert num_red_wins + num_blue_wins == num_games # Make sure we do not have a bad row without winner or such.
blue_winrate = num_blue_wins/num_games
red_winrate = num_red_wins/num_games
print(f"There are {num_games} games recorded, the blue team won {num_blue_wins},",
      f"the red team won {num_red_wins} of these games.",
      f"\nThis yields win rates of {blue_winrate*100:.2f}% (blue) and {red_winrate*100:.2f}% (red).")

# ### Looking at the champion stats
# Now we will prepare a second important data frame using the data above: The statistics per champion.
# To get the unique champion names, let's use `np.unique` on all role columns in `df`.

Blue = [f'B{role}' for role in roles]
Red = [f'R{role}' for role in roles]
champions = np.unique(df[Blue+Red])
# cd = pd.DataFrame(champions, columns=['Champion'])

# Now we compute the statistics per champion.
#
# In order to speed up the process by using `dict` lookups (which are very fast), we will not do the following steps in the `cd` data frame directly but make use of four separate dictionaries that capture the numbers of games/wins on the blue/red side for each champion. We also memorize the roles that the champions were played in, using a `dict` with `dict`s as values.
#
# To actually count the values we are interested in, we iterate over the data frame of games `df` _once_. For each row, we iterate over the roles and add to the counters in `blue_played` and `red_played` for the champions played on the respective side. We also memorize the role that each champion was played in. In order to count the wins on either side, we make use of python's automatic type casting and add the boolean `winner_is_blue`/`winner_is_red` to the counters in `blue_won` and `red_won`.

# +
blue_played = {champ: 0 for champ in champions}
blue_won = {champ: 0 for champ in champions}
red_played = {champ: 0 for champ in champions}
red_won = {champ: 0 for champ in champions}
roles_played = {champ: {role: 0 for role in roles} for champ in champions}

for _, row in df.iterrows():
    winner_is_blue = row.winner=='Blue'
    winner_is_red = not winner_is_blue
    for blue_role in Blue:
        champ = row[blue_role]
        blue_played[champ] += 1
        blue_won[champ] += winner_is_blue
        # Strip the "B"/"R" from blue_role to get the role
        roles_played[champ][blue_role[1:]] += 1
    for red_role in Red:
        champ = row[red_role]
        red_played[champ] += 1
        red_won[champ] += winner_is_red
        roles_played[champ][red_role[1:]] += 1
# -

# Before storing everything in a data frame, let's figure out which were the most played roles per champion. For this, we iterate over the champions and sort the roles by their occurences for each champion. The `number_of_roles_to_record` most played roles and their counters are then stored in individual lists and linked to keys, for example `"Role1"` and `"#Role1"`, in a dictionary:

# +
number_of_roles_to_record = 2 # We use 2 roles, could use up to all 5
ordered_roles_played = [[] for _ in range(number_of_roles_to_record)]
numbers_roles_played = [[] for _ in range(number_of_roles_to_record)]
for i, champ in enumerate(champions):
    # This is a list of tuples (role, #plays in the role):
    roles_for_this_champ = list(roles_played[champ].items())
    # sort by number of plays, in descending order (reverse=True)
    sorted_roles_for_this_champ = sorted(roles_for_this_champ, key=lambda x: x[1], reverse=True)

    # Now let's record the sorted tuples as order of most played roles (and their # of plays)
    for j in range(number_of_roles_to_record):
        ordered_roles_played[j].append(sorted_roles_for_this_champ[j][0]) # Record the role
        numbers_roles_played[j].append(sorted_roles_for_this_champ[j][1]) # Record the # of plays

most_played_roles = {f"Role{j+1}": ordered_roles_played[j] for j in range(number_of_roles_to_record)}
most_played_numbers = {f"#Role{j+1}": numbers_roles_played[j] for j in range(number_of_roles_to_record)}
# -

# Having all statistics sorted out, we can wrap everything up in a data frame. Because of the way we stored the most played roles above, we have a flexible pipeline that will generate the data frame for any number of most-played roles we want to store per champion.

cd = pd.DataFrame({
    'Champion': champions,
    'BluePlayed': [blue_played[champ] for champ in champions],
    'BlueWon': [blue_won[champ] for champ in champions],
    'RedPlayed': [red_played[champ] for champ in champions],
    'RedWon': [red_won[champ] for champ in champions],
    **most_played_roles,
    **most_played_numbers,
})

# We conclude the first round of data analysis by computing the total number of games played and the win rate on either side as well as in total, for each champion. For this, the column-wise operations on a data frame are very handy:

cd['TotalPlayed'] = cd['BluePlayed'] + cd['RedPlayed']
cd['Bluewinrate'] = cd['BlueWon'] / cd['BluePlayed']
cd['Redwinrate'] = cd['RedWon'] / cd['RedPlayed']
cd['Totalwinrate'] = (cd['BlueWon'] + cd['RedWon']) / cd['TotalPlayed']

# The resulting data frame looks like this:

# cd

# For other parts of the project we will want to come back to this data. Let's store it in a new `.csv` file.

cd.to_csv('../data/processed/ChampionStatsDemo.csv',index=False)


r"""
Line fronts
-----------

Using the :meth:`pygmt.Figure.plot` method you can draw a so-called
*front* which allows to plot specific symbols distributed along a line
or curve. Typical use cases are weather fronts, fault lines,
subduction zones, and more.

A front can be drawn by passing **f**\[±]\ *gap*\[/*size*] to the ``style``
parameter where *gap* defines the distance gap between the symbols and
*size* the symbol size. If *gap* is negative, it is interpreted to mean
the number of symbols along the front instead. If *gap* has a leading +
then we use the value exactly as given [Default will start and end each
line with a symbol, hence the *gap* is adjusted to fit]. If *size* is
missing it is set to 30% of the *gap*, except when *gap* is negative
and *size* is thus required. Append **+l** or **+r** to plot symbols on
the left or right side of the front [Default is centered]. Append
**+**\ *type* to specify which symbol to plot: **b**\ ox, **c**\ ircle,
**f**\ ault (default), **s**\ lip, or **t**\ riangle. Slip means left-lateral
or right-lateral strike-slip arrows (centered is not an option). The **+s**
modifier optionally accepts the angle used to draw the vector (default is
20). Alternatively, use **+S** which draws arcuate arrow heads. Append
**+o**\ *offset* to offset the first symbol from the beginning of the front
by that amount (default is 0). The chosen symbol is drawn with the same pen
as set for the line (i.e., via the ``pen`` parameter). To use an alternate
pen, append **+p**\ *pen*. To skip the outline, just use **+p** with no
argument. To make the main front line invisible, add **+i**.

"""

import numpy as np
import pygmt

# Generate a two-point line for plotting
x = np.array([1, 4])
y = np.array([20, 20])

fig = pygmt.Figure()
fig.basemap(region=[0, 10, 0, 20], projection="X15c/15c", frame='+t"Line Fronts"')

# Plot the line using different front styles
for frontstyle in [
    # line with "faults" front style, same as +f (default)
    "f1c/0.25c",
    # line with box front style
    "f1c/0.25c+b",
    # line with circle front style
    "f1c/0.25c+c",
    # line with triangle front style
    "f1c/0.3c+t",
    # line with left-lateral ("+l") slip ("+s") front style, angle is set to 45
    # and offset to 2.25 cm
    "f5c/1c+l+s45+o2.25c",
    # line with "faults" front style, symbols are plotted on the left side of
    # the front
    "f1c/0.4c+l",
    # line with box front style, symbols are plotted on the left side of the
    # front
    "f1c/0.3c+l+b",
    # line with circle front style, symbols are plotted on the right side of
    # the front
    "f1c/0.4c+r+c",
    # line with triangle front style, symbols are plotted on the left side of
    # the front
    "f1c/0.3c+l+t",
    # line with triangle front style, symbols are plotted on the right side of
    # the front, use other pen for the outline of the symbol
    "f1c/0.4c+r+t+p1.5p,dodgerblue",
    # line with triangle front style, symbols are plotted on the right side of
    # the front and offset is set to 0.3 cm, skip the outline
    "f0.5c/0.3c+r+t+o0.3c+p",
    # line with triangle front style, symbols are plotted on the right side of
    # the front and offset is set to 0.3 cm, skip the outline and make the main
    # front line invisible
    "f0.5c/0.3c+r+t+o0.3c+p+i",
]:
    y -= 1  # move the current line down
    fig.plot(x=x, y=y, pen="1.25p", style=frontstyle, color="red3")
    fig.text(
        x=x[-1],
        y=y[-1],
        text=frontstyle,
        font="Courier-Bold",
        justify="ML",
        offset="0.75c/0c",
    )

fig.show()


'''Create a random number
guessing game.'''

# import choice from the
# numpy.random module.
from numpy.random import choice

# choice function and possible nums
number = choice([1, 2, 3, 4,
    5, 6, 7, 8, 9])
# instantiate guess variable
i = 0
# if i isn't number
while i != number:
    # i is the user's guess
    i = int(input("Guess a number between 1 and 9: "))
    # break loop if i is the
    # generated number
    if i == number:
        break
# print "you win!" when while
# loop is broken.
print("You win!")


import numpy as np

name = "jeff bezos"
x = ""
for i in name.split():
    x += i.capitalize()
    x += " "

print(x)

def mypow(a, b):
    return a**b

def rect_square(h, w):
    return h*w

def sq_square(a):
    return rect_square(a, a)

def p_delim(rep=10):
    print("-"*rep)

def p_print(val):
    p_delim()
    print("value: %.2f" % val)
    p_delim(20)

def sq_circle(radius):
    return np.pi*radius**2

def add_multy(a, b):
    sum = a + b
    mult = a*b

    return sum, mult

def empty_funct():
    pass

def cels2far(deg):
    return deg * 1.8 + 32

haha = cels2far(231)


# -*- coding: utf-8 -*-
"""AI4I 2020 Predictive Maintenance.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1EPhjz5gcUOKs77FVQvbpeJNrON8IpvIb
"""

# Commented out IPython magic to ensure Python compatibility.
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
# %matplotlib inline
import seaborn as sns

"""Data Link: https://archive.ics.uci.edu/ml/datasets/AI4I+2020+Predictive+Maintenance+Dataset"""

data = pd.read_csv("ai4i2020.csv")

data

data['Machine failure'].unique()

data['Machine failure'].value_counts()

"""## It is a class imbalance problem.

# EDA
"""

data.info()

#check missing value
data.isna().sum()

data['Type'].unique()

"""# How to handle categorical column?
- One Hot Encoding

# Data Preparation
"""

def data_preparation(df):
  df = df.copy()

  #drop unnecessary columns
  df = df.drop(["UDI", "Product ID", "TWF", "HDF", "PWF", "OSF", "RNF"], axis=1)

  return df

X = data_preparation(data)

X

"""# Task
- To find "Machine Failure" if yes then 1 otherwise 0

# How we can do it?
- Traditional ML library "Scikit Learn"
- PyCaret (low code)

1. SK Learn
- Pipeline
- Individual algorithm

# Using PyCaret
"""

!pip install pycaret

import pycaret.classification as pyc

dir(pyc)

"""- setup initialization, setup()
- compare_models()
- create_model()
- tune_model()
- predict_model
-save_model()
"""

pyc.setup(
    data = X,
    target = "Machine failure",
    train_size = 0.8,
    normalize = True
)

pyc.compare_models()

best_model = pyc.create_model('lightgbm')

print(best_model)

pyc.evaluate_model(best_model)

tuned_lgbm_model = pyc.tune_model(best_model)

pyc.evaluate_model(tuned_lgbm_model)

pyc.save_model(best_model, "machine_failure")

"""## Using Sklearn's Pipeline"""

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, OneHotEncoder
# no need to scale the data for tree-based model, only for linear models
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline

from sklearn.linear_model import LogisticRegression

from sklearn.metrics import plot_confusion_matrix, classification_report

data

def preprocess_inputs(df):
  df = df.copy()

  #drop unnecessary columns
  df = df.drop(["UDI", "Product ID", "TWF", "HDF", "PWF", "OSF", "RNF"], axis=1)

  #X(features) and y(target/class/labels)
  X = df.drop('Machine failure', axis=1)
  y = df['Machine failure']

  #split
  X_train, X_test, y_train, y_test = train_test_split(X,y, train_size=0.8, shuffle=True, random_state=1)

  return X_train, X_test, y_train, y_test

X_train, X_test, y_train, y_test = preprocess_inputs(data)

X_train

y_train

print(len(X_train))
print(len(X_test))
print(len(y_train))
print(len(y_test))

"""## Benefit of Sklearn pipeline
- you can pass multiple steps and transformers in one function
"""

single_transformer = Pipeline(steps=[
                                     ("encode", OneHotEncoder(sparse=False))
])

col_transformer = ColumnTransformer(transformers=[
                                                  ("colencode", single_transformer, ['Type'])
], remainder = 'passthrough')

model = Pipeline(steps=[
                        ("coltransform", col_transformer),
                        ("scale", StandardScaler()),
                        ("classifier", LogisticRegression())
])

clf = model.fit(X_train, y_train)

print(clf)

"""# Evaluation"""

score = clf.score(X_test, y_test)
print("Model score is:", np.round(score*100), "%")

y_pred = clf.predict(X_test)
print(y_pred)

plot_confusion_matrix(clf, X_test, y_test, labels=clf.classes_)

clr = classification_report(y_test, y_pred, labels=clf.classes_)
print(clr)

import pickle

# save the model to disk
filename = 'logreg_model.pkl'
pickle.dump(clf, open(filename, 'wb'))

Ai4i Predictive Maintenance Machine Failure/ai4i_2020_predictive_maintenance.py


import csv
import numpy as np

class DataClass:

    def __init__(self):
        self.data = {
            "headers": [],
            "series": []
        }

    def read_data(self,file):
        with open(file, 'r') as dest_f:
            data_iter = csv.reader(dest_f,
                                   delimiter=",",
                                   quotechar='"')
            data = [data for data in data_iter]
            data_np = np.array(data)

            self.data["headers"] = data_np[:, 0]
            self.data["series"] = data_np[:, 1:].astype(float)
            # self.data["series_array"] = data

        return self.data

    def write_data(self, file, data, dim):
        with open(file, 'w', newline="") as dest_f:
            cw = csv.writer(dest_f,
                       delimiter=',',
                       quotechar='"')

            if dim==1:
                cw.writerow(data)
            elif dim==2:
                for row in data:
                    cw.writerow(row)


# coding=utf-8
import datetime as dt
import numpy as np
import pandas as pd

now = dt.datetime.now()  # 获得当前日期和时间
print("Now is {}".format(now))
yesterday = now - dt.timedelta(1)  # 通过timedelta函数执行日期加减
print("Yesterday is {}".format(yesterday.strftime('%Y-%m-%d')))

pd_day = pd.Timestamp('2019-01-05')  # 通过pd.Timestamp获得指定日期的时间戳
print("Weekday is {}".format(pd_day.day_name()))  # 获取对应的星期名称
print("Next day is {}".format(pd_day + pd.Timedelta('1 day')))  # 通过pd.Timedelta函数执行日期加减

this_year = pd.date_range(dt.datetime(2019, 1, 1),
                          dt.datetime(2019, 12, 31), freq='15D')  # 获取指定范围的日期序列，参数freq表示间隔时间，这里为15天
print("Selected days in 2019: \n{}\n".format(this_year))  # pandas.date_ranged的返回值是DatetimeIndex类型

df = pd.DataFrame(np.random.randint(0, 100, this_year.size),
                  index=this_year)  # 用DatetimeIndex作索引可以直接指定某个范围来选择数据
print("2019 Jan: \n{}\n".format(df['2019-01']))  #

# ### 时间日期
# http://pandas.pydata.org/pandas-docs/stable/timeseries.html
# 除了datetime, time, calendar等几个标准模块外，Pandas提供多个函数用来生成和操作时间日期序列；