# -*- coding: utf-8 -*-
# numkit.integration test cases
# Part of GromacsWrapper
# Copyright (c) Oliver Beckstein <orbeckst@gmail.com>
# Published under the Modified BSD Licence.

"""
===================================
 Test cases for numkit.observables
====================================

"""
from __future__ import division

import numkit.observables

import numpy as np
from numpy.testing import assert_equal, assert_almost_equal

import pytest

class TestQID(object):
    @staticmethod
    def test_None():
        q = numkit.observables.QID()
        assert q is None

    @staticmethod
    @pytest.mark.parametrize("identifiers", [
        (0, 1, 42), 42, ("a", "b", "a"), ("aba", ), "aba",
        ])
    def test_create(identifiers):
        q = numkit.observables.QID(identifiers)
        assert q == set(numkit.observables.asiterable(identifiers))

    @staticmethod
    def test_union():
        q1 = numkit.observables.QID([1, 2, 3])
        q2 = numkit.observables.QID([4, 2, 3, 5])

        q3 = q1.union(q2)

        assert isinstance(q3, frozenset)
        assert q3 == set([1, 2, 3, 4, 2, 3, 5])

@pytest.fixture
def Q1():
    return numkit.observables.QuantityWithError(1.0, 0.5)

@pytest.fixture
def Q1b():
    return numkit.observables.QuantityWithError(1.0, 0.5)

@pytest.fixture
def Q2():
    return numkit.observables.QuantityWithError(-1.0, 1.0)

@pytest.fixture()
def Q(Q1, Q1b, Q2):
    return {'Q1': Q1, 'Q1b': Q1b, 'Q2': Q2}

class TestQuantityWithError(object):
    def test_astuple(self, Q1):
        assert Q1.astuple() == (Q1.value, Q1.error)

    def test_astuple_other_Q(self, Q1, Q2, ref=(-1.0, 1.0)):
        self._test_astuple_other(Q1, Q2, ref)

    def test_astuple_other_float(self, Q1, other=22.3, ref=(22.3, 0)):
        self._test_astuple_other(Q1, other, ref)

    def _test_astuple_other(self, q, other, ref):
        val, err, qid = q._astuple(other)
        assert_almost_equal((val, err), ref)
        if hasattr(other, "qid"):
            assert qid == other.qid

    def test_error(self, Q1):
        assert_almost_equal(Q1.variance, Q1.error**2)

    def test_set_error(self, Q2, error=2.0):
        Q2.error = error
        assert Q2.variance == error**2

    @pytest.mark.parametrize("q1,q2,same", [
        ("Q1", "Q1", True),
        ("Q1", "Q1b", False),
        ("Q1", "Q2", False),
        ("Q2", "Q2", True),
    ])
    def test_sameness(self, Q, q1, q2, same):
        assert (Q[q1].isSame(Q[q2])) == same

    @pytest.mark.parametrize("q1,q2,same", [
        ("Q1", "Q1", True),
        ("Q1", "Q1b", True),
        ("Q1", "Q2", False),
        ("Q2", "Q2", True),
    ])
    def test_equality(self, Q, q1, q2, same):
        assert (Q[q1] == Q[q2]) == same

    @pytest.mark.parametrize("op,ref",
                             [("Q1 {} Q2".format(op),
                               "Q1.value {} Q2.value".format(op)) for op in
                              (">", ">=", "<", "<=", "==")])
    def test_comparisons(self, Q1, Q2, op, ref):
        assert eval(op) == eval(ref)

    @pytest.mark.parametrize("op,value", [
        (lambda x, y: 5*x, True),
        (lambda x, y: x*5, True),
        (lambda x, y: -x, True),
        (lambda x, y: abs(x), True),
        (lambda x, y: x + x, True),
        (lambda x, y: x - x, True),
        (lambda x, y: x * x, True),
        (lambda x, y: x * x * x, True),
        (lambda x, y: x / x, True),
        (lambda x, y: 5 + x, True),
        (lambda x, y: 5 - x, True),
        (lambda x, y: x - 5, True),
        (lambda x, y: x**x, True),
        pytest.param(lambda x, y: x**3, True,
                     marks=pytest.mark.xfail),  # need to fix
        (lambda x, y: x - 5*y, False),
        (lambda x, y: x * y, False),
        (lambda x, y: y / x, False),
        (lambda x, y: x**y, False),
    ]   )
    def test_sameness_algebra(self, Q1, Q2, op, value):
        assert Q1.isSame(op(Q1, Q2)) == value

    def test_copy(self, Q1):
        other = Q1.copy()
        assert not Q1.isSame(other)

    def test_deepcopy(self, Q1):
        other = Q1.deepcopy()
        assert Q1.isSame(other)


import unittest
import os
import numpy as npy

import skrf as rf
from skrf.io.touchstone import Touchstone

class TouchstoneTestCase(unittest.TestCase):
    '''
    TouchstoneTestCase tests the IO of Touchstone files
    '''
    def setUp(self):
        '''
        Sets up the test directory
        '''
        self.test_dir = os.path.dirname(os.path.abspath(__file__))+'/'

    def test_read_data(self):
        '''
        This test reads data from simple_touchstone.s2p and compares with known
        true values.
        '''
        filename = os.path.join(self.test_dir, 'simple_touchstone.s2p')
        touch = Touchstone(filename)
        f, s = touch.get_sparameter_arrays()
        z0 = complex(touch.resistance)
        f_true = npy.array([1.00000000e+09, 1.10000000e+09])
        s_true = npy.array([[[1.+2.j, 5.+6.j], [3.+4.j, 7.+8.j]],
                            [[9.+10.j, 13.+14.j], [11.+12.j, 15.+16.j]]])
        z0_true = 50+50j

        self.assertTrue((f == f_true).all())
        self.assertTrue((s == s_true).all())
        self.assertTrue((z0 == z0_true))

    def test_read_from_fid(self):
        '''
        This tests reading touch stone data from a file object as compared with
        a string path and name of the file.
        '''
        with open(os.path.join(self.test_dir, 'simple_touchstone.s2p')) as fid:
            touch = Touchstone(fid)
        f, s = touch.get_sparameter_arrays()
        z0 = complex(touch.resistance)
        f_true = npy.array([1.00000000e+09, 1.10000000e+09])
        s_true = npy.array([[[1.+2.j, 5.+6.j], [3.+4.j, 7.+8.j]],
                            [[9.+10.j, 13.+14.j], [11.+12.j, 15.+16.j]]])
        z0_true = 50+50j

        self.assertTrue((f == f_true).all())
        self.assertTrue((s == s_true).all())
        self.assertTrue((z0 == z0_true))

    def test_get_sparameter_data(self):
        '''
        This tests the get_sparameter_data function.

        '''
        with open(os.path.join(self.test_dir, 'simple_touchstone.s2p')) as fid:
            touch = Touchstone(fid)

        expected_keys = ["frequency", "S11R", "S11I", "S12R", "S12I",
                "S21R", "S21I", "S22R", "S22I", ]

        unexpected_keys = ['S11DB', 'S11M', ]

        # get dict data structure
        sp_ri = touch.get_sparameter_data(format="ri")

        # test data structure
        for ek in expected_keys:
            self.assertTrue(ek in sp_ri)

        for uk in unexpected_keys:
            self.assertFalse(uk in sp_ri)

        # test data contents
        expected_sp_ri = {
            'frequency': npy.array([1.0e+09, 1.1e+09]),
            'S11R': npy.array([1., 9.]),
            'S11I': npy.array([ 2., 10.]),
            'S21R': npy.array([ 3., 11.]),
            'S21I': npy.array([ 4., 12.]),
            'S12R': npy.array([ 5., 13.]),
            'S12I': npy.array([ 6., 14.]),
            'S22R': npy.array([ 7., 15.]),
            'S22I': npy.array([ 8., 16.]),
        }

        for k in sp_ri:
            self.assertTrue(k in expected_sp_ri)

            self.assertTrue( (expected_sp_ri[k] == sp_ri[k]).all(),
                    msg='Field %s does not match. Expected "%s", got "%s"'%(
                        k, str(expected_sp_ri[k]), str(sp_ri[k]))  )

suite = unittest.TestLoader().loadTestsFromTestCase(TouchstoneTestCase)
unittest.TextTestRunner(verbosity=2).run(suite)


__author__ = 'Robert Meyer'

import numpy as np
import inspect
import os # For path names being viable under Windows and Linux

from pypet import Environment, Parameter, ArrayParameter, Trajectory
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# Here we will see how we can write our own custom parameters and how we can use
# it with a trajectory.

# Now we want to do a more sophisticated simulations, we will integrate a differential equation
# with an Euler scheme

# Let's first define our job to do
def euler_scheme(traj, diff_func):
    """Simulation function for Euler integration.

    :param traj:

        Container for parameters and results

    :param diff_func:

        The differential equation we want to integrate

    """

    steps = traj.steps
    initial_conditions = traj.initial_conditions
    dimension = len(initial_conditions)

    # This array will collect the results
    result_array = np.zeros((steps,dimension))
    # Get the function parameters stored into `traj` as a dictionary
    # with the (short) names as keys :
    func_params_dict = traj.func_params.f_to_dict(short_names=True, fast_access=True)
    # Take initial conditions as first result
    result_array[0] = initial_conditions

    # Now we compute the Euler Scheme steps-1 times
    for idx in range(1,steps):
        result_array[idx] = diff_func(result_array[idx-1], **func_params_dict) * traj.dt + \
                            result_array[idx-1]
    # Note the **func_params_dict unzips the dictionary, it's the reverse of **kwargs in function
    # definitions!

    #Finally we want to keep the results
    traj.f_add_result('euler_evolution', data=result_array, comment='Our time series data!')

# Ok, now we want to make our own (derived) parameter that stores source code of python functions.
# We do NOT want a parameter that stores an executable function. This would complicate
# the problem a lot. If you have something like that in mind, you might wanna take a look
# at the marshal (http://docs.python.org/2/library/marshal) module
# or dill (https://pypi.python.org/pypi/dill) package.
# Our intention here is to define a parameter that we later on use as a derived parameter
# to simply keep track of the source code we use ('git' would be, of course, the better solution
# but this is just an illustrative example)
class FunctionParameter(Parameter):
    # We need to override the `f_set` function and simply extract the the source code if our
    # item is callable and store this instead.
    def f_set(self, data):
        if callable(data):
            data = inspect.getsource(data)
        return super(FunctionParameter, self).f_set(data)

    # For more complicate parameters you might consider implementing:
    # `f_supports` (we do not need it since we convert the data to stuff the parameter already
    #    supports, and that is strings!)
    #
    # and
    # the private functions
    #
    # `_values_of_same_type` (to tell whether data is similar, i.e. of two data items agree in their
    #   type, this is important to only allow exploration within the same dimension.
    #   For instance, a parameter that stores integers, should only explore integers etc.)
    #
    # and
    #
    # `_equal_values` (to tell if two data items are equal. This is important for merging if you
    #       want to erase duplicate parameter points. The trajectory needs to know when a
    #       parameter space point was visited before.)
    #
    # and
    #
    # `_store` (to be able to turn complex data into basic types understood by the storage service)
    #
    # and
    #
    # `_load` (to be able to recover your complex data form the basic types understood by the storage
    # service)
    #
    # But for now we will rely on the parent functions and hope for the best!

# Ok now let's follow the ideas in the final section of the cookbook and let's
# have a part in our simulation that only defines the parameters.
def add_parameters(traj):
    """Adds all necessary parameters to the `traj` container"""

    traj.f_add_parameter('steps', 10000, comment='Number of time steps to simulate')
    traj.f_add_parameter('dt', 0.01, comment='Step size')

    # Here we want to add the initial conditions as an array parameter. We will simulate
    # a 3-D differential equation, the Lorenz attractor.
    traj.f_add_parameter(ArrayParameter,'initial_conditions', np.array([0.0,0.0,0.0]),
                         comment = 'Our initial conditions, as default we will start from'
                                   ' origin!')

    # We will group all parameters of the Lorenz differential equation into the group 'func_params'
    traj.f_add_parameter('func_params.sigma', 10.0)
    traj.f_add_parameter('func_params.beta', 8.0/3.0)
    traj.f_add_parameter('func_params.rho', 28.0)

    #For the fun of it we will annotate the  group
    traj.func_params.v_annotations.info='This group contains as default the original values chosen ' \
                                   'by Edward Lorenz in 1963. Check it out on wikipedia ' \
                                   '(https://en.wikipedia.org/wiki/Lorenz_attractor)!'

# We need to define the lorenz function, we will assume that the value array is 3 dimensional,
# First dimension contains the x-component, second y-component, and third the z-component
def diff_lorenz(value_array, sigma, beta, rho):
    """The Lorenz attractor differential equation

    :param value_array: 3d array containing the x,y, and z component values.
    :param sigma: Constant attractor parameter
    :param beta: FConstant attractor parameter
    :param rho: Constant attractor parameter

    :return: 3d array of the Lorenz system evaluated at `value_array`

    """
    diff_array = np.zeros(3)
    diff_array[0] = sigma * (value_array[1]-value_array[0])
    diff_array[1] = value_array[0] * (rho - value_array[2]) - value_array[1]
    diff_array[2] = value_array[0] * value_array[1] - beta * value_array[2]

    return diff_array

# And here goes our main function
def main():

    filename = os.path.join('hdf5', 'example_05.hdf5')
    env = Environment(trajectory='Example_05_Euler_Integration',
                      filename=filename,
                      file_title='Example_05_Euler_Integration',
                      overwrite_file=True,
                      comment='Go for Euler!')

    traj = env.trajectory
    trajectory_name = traj.v_name

    # 1st a) phase parameter addition
    add_parameters(traj)

    # 1st b) phase preparation
    # We will add the differential equation (well, its source code only) as a derived parameter
    traj.f_add_derived_parameter(FunctionParameter,'diff_eq', diff_lorenz,
                                 comment='Source code of our equation!')

    # We want to explore some initial conditions
    traj.f_explore({'initial_conditions' : [
        np.array([0.01,0.01,0.01]),
        np.array([2.02,0.02,0.02]),
        np.array([42.0,4.2,0.42])
    ]})
    # 3 different conditions are enough for an illustrative example

    # 2nd phase let's run the experiment
    # We pass `euler_scheme` as our top-level simulation function and
    # the Lorenz equation 'diff_lorenz' as an additional argument
    env.run(euler_scheme, diff_lorenz)

    # We don't have a 3rd phase of post-processing here

    # 4th phase analysis.
    # I would recommend to do post-processing completely independent from the simulation,
    # but for simplicity let's do it here.

    # Let's assume that we start all over again and load the entire trajectory new.
    # Yet, there is an error within this approach, do you spot it?
    del traj
    traj = Trajectory(filename=filename)

    # We will only fully load parameters and derived parameters.
    # Results will be loaded manually later on.
    try:
        # However, this will fail because our trajectory does not know how to
        # build the FunctionParameter. You have seen this coming, right?
        traj.f_load(name=trajectory_name, load_parameters=2, load_derived_parameters=2,
                    load_results=1)
    except ImportError as e:

        print('That did\'nt work, I am sorry: %s ' % str(e))

        # Ok, let's try again but this time with adding our parameter to the imports
        traj = Trajectory(filename=filename,
                           dynamically_imported_classes=FunctionParameter)

        # Now it works:
        traj.f_load(name=trajectory_name, load_parameters=2, load_derived_parameters=2,
                    load_results=1)

    #For the fun of it, let's print the source code
    print('\n ---------- The source code of your function ---------- \n %s' % traj.diff_eq)

    # Let's get the exploration array:
    initial_conditions_exploration_array = traj.f_get('initial_conditions').f_get_range()
    # Now let's plot our simulated equations for the different initial conditions:
    # We will iterate through the run names
    for idx, run_name in enumerate(traj.f_get_run_names()):

        #Get the result of run idx from the trajectory
        euler_result = traj.results.f_get(run_name).euler_evolution
        # Now we manually need to load the result. Actually the results are not so large and we
        # could load them all at once. But for demonstration we do as if they were huge:
        traj.f_load_item(euler_result)
        euler_data = euler_result.data

        #Plot fancy 3d plot
        fig = plt.figure(idx)
        ax = fig.gca(projection='3d')
        x = euler_data[:,0]
        y = euler_data[:,1]
        z = euler_data[:,2]
        ax.plot(x, y, z, label='Initial Conditions: %s' % str(initial_conditions_exploration_array[idx]))
        plt.legend()
        plt.show()

        # Now we free the data again (because we assume its huuuuuuge):
        del euler_data
        euler_result.f_empty()

    # You have to click through the images to stop the example_05 module!

    # Finally disable logging and close all log-files
    env.disable_logging()

if __name__ == '__main__':
    main()


"""
Tests for subsetting operations
"""
from datetime import datetime
from unittest import TestCase

import numpy as np
import xarray as xr
import pandas as pd

from cate.core.op import OP_REGISTRY
from cate.core.opimpl import subset_spatial_impl
from cate.core.types import ValidationError
from cate.ops import subset
from cate.util.misc import object_to_qualified_name

def assert_dataset_equal(expected, actual):
    # this method is functionally equivalent to
    # `assert expected == actual`, but it checks each aspect
    # of equality separately for easier debugging
    assert expected.equals(actual), (expected, actual)

def get_test_subset_non_valid_lat_lon_dataset():
    temp = np.random.randn(2, 2, 3)
    precip = np.random.rand(2, 2, 3)
    lon = [[-40, 40], [-40, 40]]
    lat = [[-50, 50], [-50, 50]]
    return xr.Dataset({'temp': (['x', 'y', 'time'], temp),
                       'precip': (['x', 'y', 'time'], precip)},
                      coords={'lon': (['x', 'y'], lon),
                              'lat': (['x', 'y'], lat),
                              'time': pd.date_range('2014-09-06', periods=3)})

class TestSubsetSpatial(TestCase):
    def test_subset_non_valid_lat_lon(self):
        """
        Test whether lat and/or lon exist and if they exist whether they have dimension = 1
        :return: void
        """

        # Test whether lat lon exist
        dataset = xr.Dataset({
            'first': (['xc', 'yc', 'time'], np.ones([180, 360, 6])),
            'second': (['xc', 'yc', 'time'], np.ones([180, 360, 6])),
            'xc': np.linspace(-89.5, 89.5, 180),
            'yc': np.linspace(-179.5, 179.5, 360),
        })

        with self.assertRaises(ValidationError) as error:
            subset_spatial_impl(dataset, (-40, 40, -50, 50))
        self.assertIn('No (valid) geocoding found', str(error.exception))

        # test whether lat lon has the wrong dimension (!=1)
        dataset = get_test_subset_non_valid_lat_lon_dataset()

        with self.assertRaises(ValidationError) as error:
            subset_spatial_impl(dataset, (-40, 40, -50, 50))
        self.assertIn('Geocoding not recognised', str(error.exception))

    def test_nominal(self):
        """
        Test general 'most expected' use case functionality.
        """
        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'second': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'lat': np.linspace(-89.5, 89.5, 180),
            'lon': np.linspace(-179.5, 179.5, 360)})
        actual = subset.subset_spatial(dataset, "-20, -10, 20, 10")
        expected = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([22, 42, 6])),
            'second': (['lat', 'lon', 'time'], np.ones([22, 42, 6])),
            'lat': np.linspace(-10.5, 10.5, 22),
            'lon': np.linspace(-20.5, 20.5, 42)})
        assert_dataset_equal(expected, actual)

    def test_inverted_dims_nominal(self):
        """
        Test if the implementation is dimension order agnostic.
        """
        # Inverted lat
        dataset = xr.Dataset({
            'first': (['lon', 'lat', 'time'], np.ones([360, 180, 6])),
            'second': (['lon', 'lat', 'time'], np.ones([360, 180, 6])),
            'lat': np.linspace(89.5, -89.5, 180),
            'lon': np.linspace(-179.5, 179.5, 360)})
        actual = subset.subset_spatial(dataset, "-20, -10, 20, 10")
        expected = xr.Dataset({
            'first': (['lon', 'lat', 'time'], np.ones([42, 22, 6])),
            'second': (['lon', 'lat', 'time'], np.ones([42, 22, 6])),
            'lat': np.linspace(10.5, -10.5, 22),
            'lon': np.linspace(-20.5, 20.5, 42)})
        assert_dataset_equal(expected, actual)

    def test_generic_masked(self):
        """
        Test using a generic Polygon and masking
        """
        # Africa
        a = str('POLYGON((-10.8984375 35.60371874069731,-19.16015625 '
                '23.885837699861995,-20.56640625 17.14079039331665,-18.6328125 '
                '7.536764322084079,-10.72265625 0.7031073524364783,10.37109375 '
                '0.3515602939922709,10.37109375 -22.268764039073965,22.8515625 '
                '-42.29356419217007,37.79296875 -27.21555620902968,49.39453125 '
                '-3.5134210456400323,54.4921875 14.093957177836236,18.984375 '
                '35.88905007936091,-10.8984375 35.60371874069731))')

        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'second': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'lat': np.linspace(-89.5, 89.5, 180),
            'lon': np.linspace(-179.5, 179.5, 360)})
        actual = subset.subset_spatial(dataset, a)
        # Gulf of Guinea
        gog = actual.sel(method='nearest', **{'lon': 1.2, 'lat': -1.4})
        self.assertTrue(np.isnan(gog['first']).all())
        # Africa
        self.assertTrue(1 == actual.sel(method='nearest', **{'lon': 20.7, 'lat': 6.15}))

    def test_generic_masked_inverted(self):
        """
        Test using a generic Polygon and masking
        """
        # Africa
        a = str('POLYGON((-10.8984375 35.60371874069731,-19.16015625 '
                '23.885837699861995,-20.56640625 17.14079039331665,-18.6328125 '
                '7.536764322084079,-10.72265625 0.7031073524364783,10.37109375 '
                '0.3515602939922709,10.37109375 -22.268764039073965,22.8515625 '
                '-42.29356419217007,37.79296875 -27.21555620902968,49.39453125 '
                '-3.5134210456400323,54.4921875 14.093957177836236,18.984375 '
                '35.88905007936091,-10.8984375 35.60371874069731))')

        # Inverted lat
        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'second': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'lat': np.linspace(89.5, -89.5, 180),
            'lon': np.linspace(-179.5, 179.5, 360)})
        actual = subset.subset_spatial(dataset, a)
        # Gulf of Guinea
        gog = actual.sel(method='nearest', **{'lon': 1.2, 'lat': -1.4})
        self.assertTrue(np.isnan(gog['first']).all())
        # Africa
        self.assertTrue(1 == actual.sel(method='nearest', **{'lon': 20.7, 'lat': 6.15}))

    def test_generic_not_masked(self):
        """
        Test using a generic Polygon without masking
        """
        # Africa
        a = str('POLYGON((-10.8984375 35.60371874069731,-19.16015625 '
                '23.885837699861995,-20.56640625 17.14079039331665,-18.6328125 '
                '7.536764322084079,-10.72265625 0.7031073524364783,10.37109375 '
                '0.3515602939922709,10.37109375 -22.268764039073965,22.8515625 '
                '-42.29356419217007,37.79296875 -27.21555620902968,49.39453125 '
                '-3.5134210456400323,54.4921875 14.093957177836236,18.984375 '
                '35.88905007936091,-10.8984375 35.60371874069731))')

        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'second': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'lat': np.linspace(-89.5, 89.5, 180),
            'lon': np.linspace(-179.5, 179.5, 360)})
        actual = subset.subset_spatial(dataset, a, mask=False)
        # Gulf of Guinea
        self.assertTrue(1 == actual.sel(method='nearest', **{'lon': 1.2, 'lat': -1.4}))
        # Africa
        self.assertTrue(1 == actual.sel(method='nearest', **{'lon': 20.7, 'lat': 6.15}))

    def test_generic_not_masked_inverted(self):
        """
        Test using a generic Polygon without masking
        """
        # Africa
        a = str('POLYGON((-10.8984375 35.60371874069731,-19.16015625 '
                '23.885837699861995,-20.56640625 17.14079039331665,-18.6328125 '
                '7.536764322084079,-10.72265625 0.7031073524364783,10.37109375 '
                '0.3515602939922709,10.37109375 -22.268764039073965,22.8515625 '
                '-42.29356419217007,37.79296875 -27.21555620902968,49.39453125 '
                '-3.5134210456400323,54.4921875 14.093957177836236,18.984375 '
                '35.88905007936091,-10.8984375 35.60371874069731))')

        # Inverted lat
        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'second': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'lat': np.linspace(89.5, -89.5, 180),
            'lon': np.linspace(-179.5, 179.5, 360)})
        actual = subset.subset_spatial(dataset, a, mask=False)
        # Gulf of Guinea
        self.assertTrue(1 == actual.sel(method='nearest', **{'lon': 1.2, 'lat': -1.4}))
        # Africa
        self.assertTrue(1 == actual.sel(method='nearest', **{'lon': 20.7, 'lat': 6.15}))

    def test_registered(self):
        """
        Test if it runs as an operation registered in the op registry.
        """
        reg_op = OP_REGISTRY.get_op(object_to_qualified_name(subset.subset_spatial))
        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'second': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'lat': np.linspace(-89.5, 89.5, 180),
            'lon': np.linspace(-179.5, 179.5, 360)})
        actual = reg_op(ds=dataset, region="-20, -10, 20, 10")
        expected = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([22, 42, 6])),
            'second': (['lat', 'lon', 'time'], np.ones([22, 42, 6])),
            'lat': np.linspace(-10.5, 10.5, 22),
            'lon': np.linspace(-20.5, 20.5, 42)})
        assert_dataset_equal(expected, actual)

    def test_antimeridian_simple(self):
        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'second': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'lat': np.linspace(-89.5, 89.5, 180),
            'lon': np.linspace(-179.5, 179.5, 360)})

        # With masking
        actual = subset.subset_spatial(dataset, '170, -5, -170, 5', mask=True)
        masked = actual.sel(method='nearest', **{'lon': 0, 'lat': 0})
        self.assertTrue(np.isnan(masked['first']).all())

    def test_antimeridian_simple_inverted(self):
        # Inverted lat
        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'second': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'lat': np.linspace(89.5, -89.5, 180),
            'lon': np.linspace(-179.5, 179.5, 360)})

        # With masking
        actual = subset.subset_spatial(dataset, '170, -5, -170, 5', mask=True)
        masked = actual.sel(method='nearest', **{'lon': 0, 'lat': 0})
        self.assertTrue(np.isnan(masked['first']).all())

    def test_antimeridian_arbitrary(self):
        antimeridian_pol = str('POLYGON(('
                               '162.0703125 39.639537564366705,'
                               '-155.390625 39.774769485295465,'
                               '-155.56640625 12.726084296948184,'
                               '162.24609375 12.897489183755905,'
                               '161.89453125 26.745610382199025,'
                               '162.0703125 39.639537564366705'
                               '))')
        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'second': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'lat': np.linspace(-89.5, 89.5, 180),
            'lon': np.linspace(-179.5, 179.5, 360)})

        with self.assertRaises(Exception) as cm:
            subset.subset_spatial(dataset, antimeridian_pol)
        self.assertIn('anti-meridian', str(cm.exception))

    def test_antimeridian_arbitrary_inverted(self):
        antimeridian_pol = str('POLYGON(('
                               '162.0703125 39.639537564366705,'
                               '-155.390625 39.774769485295465,'
                               '-155.56640625 12.726084296948184,'
                               '162.24609375 12.897489183755905,'
                               '161.89453125 26.745610382199025,'
                               '162.0703125 39.639537564366705'
                               '))')
        # Inverted lat
        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'second': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'lat': np.linspace(89.5, -89.5, 180),
            'lon': np.linspace(-179.5, 179.5, 360)})

        with self.assertRaises(Exception) as cm:
            subset.subset_spatial(dataset, antimeridian_pol)
        self.assertIn('anti-meridian', str(cm.exception))

    def test_select_single_center(self):
        """
        Test subset spatial with a polygon that completely fits
        inside pixel bounds
        """
        # Masked
        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([6, 12, 3])),
            'second': (['lat', 'lon', 'time'], np.ones([6, 12, 3])),
            'lat': np.linspace(-75, 75, 6),
            'lon': np.linspace(-165, 165, 12),
            'time': [datetime(2000, x, 1) for x in range(1, 4)]})
        poly = ((32, 2), (34, 28), (58, 29), (54, 4))
        actual = subset.subset_spatial(dataset, region=poly)
        expected = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([1, 1, 3])),
            'second': (['lat', 'lon', 'time'], np.ones([1, 1, 3])),
            'lat': [15.],
            'lon': [45.],
            'time': [datetime(2000, x, 1) for x in range(1, 4)]})

        assert_dataset_equal(expected, actual)

        # Not masked
        actual = subset.subset_spatial(dataset, region=poly, mask=False)
        assert_dataset_equal(expected, actual)

    def test_select_single_vertice(self):
        """
        Test with a polygon that encloses a single pixel vertice,
        but no centers. That is, crosses four pixels
        """
        # Masked
        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([6, 12, 3])),
            'second': (['lat', 'lon', 'time'], np.ones([6, 12, 3])),
            'lat': np.linspace(-75, 75, 6),
            'lon': np.linspace(-165, 165, 12),
            'time': [datetime(2000, x, 1) for x in range(1, 4)]})
        poly = ((25, 2), (27, 40), (58, 38), (54, 4))
        actual = subset.subset_spatial(dataset, region=poly)
        expected = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([2, 2, 3])),
            'second': (['lat', 'lon', 'time'], np.ones([2, 2, 3])),
            'lat': [15., 45.],
            'lon': [15., 45.],
            'time': [datetime(2000, x, 1) for x in range(1, 4)]})

        assert_dataset_equal(expected, actual)

        # Not masked
        actual = subset.subset_spatial(dataset, region=poly, mask=False)
        assert_dataset_equal(expected, actual)

    def test_select_1d_lat(self):
        """
        Test with a polyon that runs over pixel boundaries in the
        latitude direction, resulting in selecting a single column
        """
        # Masked
        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([6, 12, 3])),
            'second': (['lat', 'lon', 'time'], np.ones([6, 12, 3])),
            'lat': np.linspace(-75, 75, 6),
            'lon': np.linspace(-165, 165, 12),
            'time': [datetime(2000, x, 1) for x in range(1, 4)]})
        poly = ((25, 2), (27, 47), (12, 50), (10, 4))
        actual = subset.subset_spatial(dataset, region=poly)
        expected = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([2, 1, 3])),
            'second': (['lat', 'lon', 'time'], np.ones([2, 1, 3])),
            'lat': [15., 45.],
            'lon': [15.],
            'time': [datetime(2000, x, 1) for x in range(1, 4)]})

        assert_dataset_equal(expected, actual)

        # Not masked
        actual = subset.subset_spatial(dataset, region=poly, mask=False)
        assert_dataset_equal(expected, actual)

    def test_select_1d_lon(self):
        """
        Test with a polygon that runs over pixel boundaries in the
        longitude direction. E.g., selects a single row
        """
        # Masked
        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([6, 12, 3])),
            'second': (['lat', 'lon', 'time'], np.ones([6, 12, 3])),
            'lat': np.linspace(-75, 75, 6),
            'lon': np.linspace(-165, 165, 12),
            'time': [datetime(2000, x, 1) for x in range(1, 4)]})
        poly = ((-28, 32), (-26, 58), (28, 52), (25, 33))
        actual = subset.subset_spatial(dataset, region=poly)
        expected = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([1, 2, 3])),
            'second': (['lat', 'lon', 'time'], np.ones([1, 2, 3])),
            'lat': [45.],
            'lon': [-15., 15.],
            'time': [datetime(2000, x, 1) for x in range(1, 4)]})

        assert_dataset_equal(expected, actual)

        # Not masked
        actual = subset.subset_spatial(dataset, region=poly, mask=False)
        assert_dataset_equal(expected, actual)

    def test_select_from_single_pixel(self):
        """
        Test subset spatial where the input dataset has only one pixel
        """
        # Masked
        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([1, 1, 3])),
            'second': (['lat', 'lon', 'time'], np.ones([1, 1, 3])),
            'lat': [45.],
            'lon': [15.],
            'time': [datetime(2000, x, 1) for x in range(1, 4)]})
        poly = ((-28, 32), (-26, 58), (28, 52), (25, 33))
        actual = subset.subset_spatial(dataset, region=poly)

        assert_dataset_equal(dataset, actual)

        # Not masked
        actual = subset.subset_spatial(dataset, region=poly, mask=False)
        assert_dataset_equal(dataset, actual)

    def test_out_of_bounds(self):
        """
        Test that an appropriate error is raised when the polygon
        wouldn't select any data from the original dataset
        """
        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([1, 1, 3])),
            'second': (['lat', 'lon', 'time'], np.ones([1, 1, 3])),
            'lat': [45.],
            'lon': [15.],
            'time': [datetime(2000, x, 1) for x in range(1, 4)]})
        poly = ((32, 32), (34, 58), (56, 56), (58, 33))

        with self.assertRaises(ValueError) as err:
            subset.subset_spatial(dataset, region=poly)
        self.assertIn('Can not select', str(err.exception))

        # Not masked
        with self.assertRaises(ValueError) as err:
            subset.subset_spatial(dataset, region=poly, mask=False)
        self.assertIn('Can not select', str(err.exception))

    def test_non_geospatial_variable(self):
        """
        Test that subsetting a dataset that contains a non-geospatial
        dataset skips such a dataset
        """
        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'second': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'third': (['time'], np.ones([6])),
            'lat': np.linspace(-89.5, 89.5, 180),
            'lon': np.linspace(-179.5, 179.5, 360),
            'time': np.array([0, 1, 2, 3, 4, 5])})
        actual = subset.subset_spatial(dataset, "-20, -10, 20, 10")
        expected = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([22, 42, 6])),
            'second': (['lat', 'lon', 'time'], np.ones([22, 42, 6])),
            'third': (['time'], np.ones([6])),
            'lat': np.linspace(-10.5, 10.5, 22),
            'lon': np.linspace(-20.5, 20.5, 42),
            'time': np.array([0, 1, 2, 3, 4, 5])})
        assert_dataset_equal(expected, actual)

        poly = "POLYGON ((6.609745636041873 45.391851854914776, 19.720614826531428 45.391851854914776, 19.720614826531428 37.06301149556095, 6.609745636041874 37.06301149556095, 6.609745636041873 45.391851854914776))"
        actual = subset.subset_spatial(dataset, poly, mask=True)
        xr.testing.assert_equal(expected.third, actual.third)

class TestSubsetTemporal(TestCase):
    def test_subset_temporal(self):
        # Test general functionality
        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'second': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'lat': np.linspace(-89.5, 89.5, 180),
            'lon': np.linspace(-179.5, 179.5, 360),
            'time': [datetime(2000, x, 1) for x in range(1, 7)]})
        actual = subset.subset_temporal(dataset, '2000-01-10, 2000-04-01')
        expected = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([180, 360, 3])),
            'second': (['lat', 'lon', 'time'], np.ones([180, 360, 3])),
            'lat': np.linspace(-89.5, 89.5, 180),
            'lon': np.linspace(-179.5, 179.5, 360),
            'time': [datetime(2000, x, 1) for x in range(2, 5)]})
        assert_dataset_equal(expected, actual)

    def test_invalid_dtype(self):
        # Test passing in a MJD dataset
        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'second': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'lat': np.linspace(-89.5, 89.5, 180),
            'lon': np.linspace(-179.5, 179.5, 360),
            'time': [2451544.5,
                     2451575.5,
                     2451604.5,
                     2451635.5,
                     2451665.5,
                     2451696.5]})
        with self.assertRaises(ValueError) as err:
            subset.subset_temporal(dataset, '2000-01-10, 2000-04-01')
        self.assertIn('type datetime', str(err.exception))

    def test_registered(self):
        """
        Test if it runs as an operation registered in the op registry.
        """
        reg_op = OP_REGISTRY.get_op(object_to_qualified_name(subset.subset_temporal))
        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'second': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'lat': np.linspace(-89.5, 89.5, 180),
            'lon': np.linspace(-179.5, 179.5, 360),
            'time': [datetime(2000, x, 1) for x in range(1, 7)]})
        actual = reg_op(ds=dataset, time_range='2000-01-10, 2000-04-01')
        expected = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([180, 360, 3])),
            'second': (['lat', 'lon', 'time'], np.ones([180, 360, 3])),
            'lat': np.linspace(-89.5, 89.5, 180),
            'lon': np.linspace(-179.5, 179.5, 360),
            'time': [datetime(2000, x, 1) for x in range(2, 5)]})
        assert_dataset_equal(expected, actual)

class TestSubsetTemporalIndex(TestCase):
    def test_subset_temporal_index(self):
        # Test general functionality
        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'second': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'lat': np.linspace(-89.5, 89.5, 180),
            'lon': np.linspace(-179.5, 179.5, 360),
            'time': ['2000-01-01',
                     '2000-02-01',
                     '2000-03-01',
                     '2000-04-01',
                     '2000-05-01',
                     '2000-06-01']})
        actual = subset.subset_temporal_index(dataset, 2, 4)
        expected = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([180, 360, 3])),
            'second': (['lat', 'lon', 'time'], np.ones([180, 360, 3])),
            'lat': np.linspace(-89.5, 89.5, 180),
            'lon': np.linspace(-179.5, 179.5, 360),
            'time': ['2000-03-01', '2000-04-01', '2000-05-01']})
        assert_dataset_equal(expected, actual)

    def test_registered(self):
        """
        Test if it runs as an operation registered in the op registry.
        """
        reg_op = OP_REGISTRY.get_op(object_to_qualified_name(subset.subset_temporal_index))
        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'second': (['lat', 'lon', 'time'], np.ones([180, 360, 6])),
            'lat': np.linspace(-89.5, 89.5, 180),
            'lon': np.linspace(-179.5, 179.5, 360),
            'time': ['2000-01-01',
                     '2000-02-01',
                     '2000-03-01',
                     '2000-04-01',
                     '2000-05-01',
                     '2000-06-01']})
        actual = reg_op(ds=dataset, time_ind_min=2, time_ind_max=4)
        expected = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.ones([180, 360, 3])),
            'second': (['lat', 'lon', 'time'], np.ones([180, 360, 3])),
            'lat': np.linspace(-89.5, 89.5, 180),
            'lon': np.linspace(-179.5, 179.5, 360),
            'time': ['2000-03-01', '2000-04-01', '2000-05-01']})
        assert_dataset_equal(expected, actual)

class TestExtractPoint(TestCase):
    @classmethod
    def setUpClass(cls):
        v1_data = np.arange(18).reshape((3, 3, 2))
        v2_data = np.arange(100, 118).reshape((3, 3, 2))
        v3_data = np.arange(36).reshape((3, 3, 2, 2))
        v4_data = np.arange(50, 59).reshape((3, 3))
        cls._ds = xr.Dataset(
            {
                'v1': (['lat', 'lon', 'd1'], v1_data),
                'v2': (['lat', 'lon', 'd2'], v2_data),
                'v3': (['lat', 'lon', 'd1', 'd2'], v3_data),
                'v4': (['lat', 'lon'], v4_data)
            },
            coords={
                'lon': np.array([12, 13, 14], np.dtype('float64')),
                'lat': np.array([22, 23, 24], np.dtype('float64')),
                'd1': [1, 2],
                'd2': [datetime(2000, 3, 1), datetime(2000, 4, 1)]
            }
        )

    def test_all_extra_dims(self):
        result = subset.extract_point(self._ds, (12.2, 23.2), indexers={'d1': 2, 'd2': '2000-03-01'})
        self.assertEqual({'lat': 23.0, 'lon': 12.0, 'v1': 7.0, 'v2': 106.0, 'v3': 14.0, 'v4': 53.0}, result)

    def test_one_extra_dims(self):
        result = subset.extract_point(self._ds, (12.2, 23.2), indexers={'d1': 2})
        self.assertEqual({'lat': 23.0, 'lon': 12.0, 'v1': 7.0, 'v4': 53.0}, result)

        result = subset.extract_point(self._ds, (12.2, 23.2), indexers={'d2': '2000-03-01'})
        self.assertEqual({'lat': 23.0, 'lon': 12.0, 'v2': 106.0, 'v4': 53.0}, result)

    def test_no_extra_dims(self):
        result = subset.extract_point(self._ds, (12.2, 23.2), indexers={})
        self.assertEqual({'lat': 23.0, 'lon': 12.0, 'v4': 53.0}, result)

        result = subset.extract_point(self._ds, (12.2, 23.2))
        self.assertEqual({'lat': 23.0, 'lon': 12.0, 'v4': 53.0}, result)

    def test_unknown_dim(self):
        result = subset.extract_point(self._ds, (12.2, 23.2), indexers={'x1': 42})
        self.assertEqual({'lat': 23.0, 'lon': 12.0, 'v4': 53.0}, result)

    def test_point_out_of_bounds(self):
        result = subset.extract_point(self._ds, (0, 0), indexers={'d1': 2, 'd2': '2000-03-01'})
        self.assertEqual({}, result)

        result = subset.extract_point(self._ds, (0, 0))
        self.assertEqual({}, result)

    def test_extra_dim_with_no_exact_match(self):
        # no exact match for 'd2', the same as if 'd2' is not given
        result = subset.extract_point(self._ds, (12.2, 23.2), indexers={'d1': 2, 'd2': '2000-03-02'})
        self.assertEqual({'lat': 23.0, 'lon': 12.0, 'v1': 7.0, 'v4': 53.0}, result)

        result = subset.extract_point(self._ds, (12.2, 23.2), indexers={'d1': 1.1, 'd2': '2000-03-01'})
        self.assertEqual({'lat': 23.0, 'lon': 12.0, 'v2': 106.0, 'v4': 53.0}, result)


# -*- coding: utf-8 -*-
"""
Created on Sun Oct 4

@author: Wang Yizhuo
"""

import random
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt

class FrozenLake:
    """Environment of frozen lake.

    Attributes:
        map_idx: 0 for 4x4 map, 1 for 10x10 map.
    """

    def __init__(self, map_idx):
        """Define and decode the map.
        Do not change the params here.

        Args:
            map_idx (int): map index
        """

        self.ACTIONS = {  # coord in the matrix
            0: (-1,  0),  # Up
            1: ( 0,  1),  # Right
            2: ( 1,  0),  # Down
            3: ( 0, -1),  # Left
        }

        self.CELL = {     # interpret the map file
            0: 'ice',
            1: 'hole',
            2: 'start',
            3: 'frisbee',
        }

        self.filepath = ['./map_4x4.txt', './map_10x10.txt']
        self.state             = None
        self.map               = np.loadtxt(self.filepath[map_idx], dtype=int)
        self.MAP_X, self.MAP_Y = self.map.shape
        self.init_states       = [(x, y) for x in range(self.MAP_X) for y in range(self.MAP_Y) if self.CELL[self.map[x, y]] == 'start'] # find init state
        self.e_fail, self.e_success, self.e_opt_policy = {}, {}, []

    def reset(self):
        """Reset the state to the start

        Returns:
            tuple: state coord
        """

        self.state = random.choice(self.init_states)
        return self.state

    def step(self, action:int):
        """Take a step given state and action, and return the next step and reward.

        Args:
            action (int): input an action

        Raises:
            RuntimeWarning: If there is no corresponding cell, the map file might contain illegal characters.

        Returns:
            tuple: next valid state given action
            float: the reward of next state
            bool: if next state is terminal
        """

        done    = False
        reward  = 0
        dx, dy  = self.ACTIONS[action]
        new_state = self.state[0] + dx, self.state[1] + dy

        reward -= 0                                                # penalty for moving each timestep, 0 by default
        if not (0 <= new_state[0] < self.MAP_X and 0 <= new_state[1] < self.MAP_Y):
            reward -= 0                                            # penalty for going out of bound, 0 by default. Keep where they are
        elif self.CELL[self.map[new_state]] in ['ice', 'start']:
            self.state = new_state
            reward -= 0                                            # penalty for moving, 0 by default
        elif self.CELL[self.map[new_state]] == 'hole':
            self.state = new_state
            reward -= 1                                            # penalty for fall into the hole, -1 by default. terminal state
            done = True
        elif self.CELL[self.map[new_state]] == 'frisbee':
            self.state = new_state
            reward += 1                                            # reward for get the frisbee, +1 by default. terminalstate
            done = True
        else:
            raise RuntimeWarning('No corresponding cell. Check your map file.')

        return self.state, reward, done

    def render(self, e_len, e, e_all, policy):
        """Store the info for each episode.
        Detect whether the episode ends with getting the frisbee, and store it in the attributes.

        Args:
            e_len (int): length of the episode
            e (int): current episode
            e_all (int): number of total episode
            policy (dict): the policy table
        """

        self._show_progress(e)                           # show progress of training on the sreen
        if self.CELL[self.map[self.state]] == 'frisbee': # store if the episode ends with getting the frisbee or not
            self.e_success[e] = e_len
        else:
            self.e_fail[e] = e_len
        if self._is_success(policy):                     # judge if the policy is optimal or not, record if it is optimal
            self.e_opt_policy.append(e)
        if e == e_all - 1:                               # judge at the last episode, show if a successful policy has been reached during the training
            if self.e_success:
                print('Frisbee firstly reach at episode:', list(self.e_success.keys())[0])
            else:
                print('More training is needed.')
            if e in self.e_success:
                self.render_policy(policy)               # show the policy

    def render_all(self, e_all, name, policy_table, Qtable):
        """Render graphs when finished the training.

        Args:
            e_all (int): number of total episode
            name (string): name of the algorithm
            policy_table (dict): policy table
            Qtable (dict): Q table
        """

        plt.figure()
        l1 = plt.scatter(self.e_fail.keys(), self.e_fail.values(), s=0.8, alpha=0.3, label='Fail')     # plot failed points in blue
        l2 = plt.scatter(self.e_success.keys(), self.e_success.values(), s=0.8, c='red', alpha=1.0)    # plot successful points in red
        l3 = plt.vlines(self.e_opt_policy, 0, max(self.e_fail.values()), colors='green', alpha=0.05)   # plot optimal policy lines in green
        plt.legend(handles=[l1,l2,l3], labels=['Fail','Success','Optimal policy'], loc='upper right')
        plt.title('Training Process of '+ name)
        plt.xlabel('#Episode')
        plt.ylabel('Step length')
        # plt.savefig(r'xxx.png', dpi=300)                # save figure for report
        self.render_heatmap(policy_table, Qtable, name)   # render the heatmap

    def render_heatmap(self, policy_table, Qtable, name):
        """Render heatmap.

        Args:
            policy_table (dict): policy table
            Qtable (dict): Q table
            name (string): name of the algorithm
        """

        num_cell = max(Qtable.keys())[0] + 1                    # number of the states
        Q = np.full((num_cell, num_cell), 0, dtype=float)       # create a square state matrix
        mask = np.full((num_cell, num_cell), False, dtype=bool) # for the use of mask of heatmap
        for x, y in Qtable.keys():
            Q[x][y] = max(Qtable[(x, y)])                       # put the values of Q table dict into the 2D array correspondingly
        for y in range(self.MAP_Y):
            for x in range(self.MAP_X):
                if self.CELL[self.map[(x, y)]] == 'hole':       # create mask, to make the cells at holes and target blank
                    mask[x][y] = True
                elif self.CELL[self.map[(x, y)]] == 'frisbee':
                    mask[x][y] = True
        plt.figure()
        sns.heatmap(Q, annot=True, cmap='RdBu_r', square=True, mask=mask, linewidths=0.3, linecolor='black', annot_kws={'size': 5}) # plot the heatmap with mask
        # sns.heatmap(Q, cmap='RdBu_r', square=True, mask=mask, linewidths=0.3, linecolor='black')   # heatmap without annotations
        plt.title('Max state-action value given state ({})'.format(name))
        # plt.savefig(r'{}.png'.format(name), dpi=300)                                               # save figure for report

    def render_policy(self, policy_table):
        """Show policy on the screen."""

        done = False
        self.state = self.reset()
        state_list, action_list = [self.state], []
        while not done:
            action = int(np.argmax(policy_table[self.state]))
            _, _, done = self.step(action)
            state_list.append(self.state)
            action_list.append(action)
        print('Policy (state): {}'.format(state_list))
        print('Policy (action): {}'.format(action_list))

    def _is_success(self, policy_table):
        """Judge whether the policy is optimal.

        Args:
            policy_table (dict)

        Returns:
            bool: True if optimal, false otherwise
        """

        state = self.reset()
        done = False
        state_list = []
        while not done:
            action = int(np.argmax(policy_table[state]))               # use argmax to extract the policy
            dx, dy = self.ACTIONS[action]
            state_list.append(state)
            new_state = state[0] + dx, state[1] + dy
            if not (0 <= new_state[0] < self.MAP_X and 0 <= new_state[1] < self.MAP_Y):
                done = True                                            # an optimal policy should not hit the wall
            elif new_state in state_list:
                done = True                                            # an optimal policy should not be recurrent
            elif self.CELL[self.map[new_state]] == 'ice':
                state = new_state
            elif self.CELL[self.map[new_state]] in ['hole', 'start']:
                done = True                                            # an optimal policy should not fall into the hole or back to the start
            elif self.CELL[self.map[new_state]] == 'frisbee':
                return True                                            # an optimal policy should lead to the target successfully
        return False

    def _show_progress(self, e):
        """Show progress on the sreen."""

        if e % 10 == 0:
            print('Episode: %d' % e, end='\r')

def render_learn_curve(envS, envQ, envM):
    """Render the learning curve of three algorithms.

    Args:
        envS (class): SARSA env
        envQ (class): Q-learning env
        envM (class): Monte Carlo env
    """

    l     = []
    num_e = len(envS.env.e_fail) + len(envS.env.e_success)
    E     = list(range(num_e))
    plt.figure()
    for env in [envS, envQ, envM]:                 # loop for all the envs
        s_all = 0                                  # accumulative step number
        s     = []                                 # accumulative step number for each episode
        for e in range(num_e):
            if e in env.env.e_fail.keys():
                s_all += env.env.e_fail[e]
                s.append(s_all)
            elif e in env.env.e_success.keys():
                s_all += env.env.e_success[e]
                s.append(s_all)
        handle, = plt.plot(E, s)                   # comma after handle to unzip
        l.append(handle)
    plt.legend(handles=l, labels=[envS.name, envQ.name, envM.name])
    plt.xlabel('#Episode')
    plt.ylabel('Accumulative steps')
    plt.title('Comparison of learning curve')

def render_success_rate(envS, envQ, envM, smooth_size=20):
    """ Render success rate for each env.

    Args:
        envS (class): SARSA env
        envQ (class): Q-learning env
        envM (class): Monte Carlo env
        smooth_size (int, optional): take an average to a batch of episodes to smooth the curve. Defaults to 20.
    """

    l = []
    num_e = len(envS.env.e_fail) + len(envS.env.e_success)
    plt.figure()
    for env in [envS, envQ, envM]:
        e_smooth, s = [], []
        for batch in range(num_e // smooth_size):
            e_smooth.append(batch * smooth_size + smooth_size // 2)     # center loc of the batch
            cnt = 0                                                     # counter for sccessful exploration
            for e in range(batch*smooth_size, (batch + 1)*smooth_size):
                if e in env.env.e_success.keys():
                    cnt += 1
            s.append(100 * cnt / smooth_size)                           # percentage of average sccess rate in a batch
        handle, = plt.plot(e_smooth, s)
        l.append(handle)
    plt.legend(handles=l, labels=[envS.name, envQ.name, envM.name])
    plt.title('Comparison of three method with a smoothness of {} episodes'.format(smooth_size))


# -*- coding: utf-8 -*-
"""
------ What is this file? ------

This test module contains some data quality tests for the istanbul_airbnb_process.csv file
The file can be found at:
    data/processed/istanbul_airbnb_processed.csv

"""
#%% --- Import Required Packages ---

import os
from pathlib import Path # To wrap around filepaths
import pytest
import numpy as np
import pandas as pd
import geopandas as gpd
import textdistance
#%% --- Set proper directory to assure integration with doit ---

abspath = os.path.abspath(__file__)
dname = os.path.dirname(abspath)
os.chdir(dname)

#%% --- Import data ---

# Dataset to test for quality
import_fp = Path("../../data/processed/istanbul_airbnb_processed.csv")
airbnb = pd.read_csv(import_fp, encoding = "utf-8-sig")

# Dataset to take as reference for lat/lon boundaries
import_fp = Path("../../data/external/istanbul_districts.shp")
istanbul_districts = gpd.read_file(import_fp)

#%% --- Data quality tests ---

class TestNullValues(object):
    def test_total_null_values(self):
        expected = 0
        actual = airbnb.isnull().sum().sum()
        error_message = "Dataset contains null values. Expected {} null values, got {}".format(expected,actual)
        assert expected == actual, error_message

class TestUniqueness(object):
    def test_total_unique_values_for_column_listing_id(self):
        expected = airbnb.shape[0]
        actual = len(airbnb.loc[:,"listing_id"].unique())
        error_message = "Column listing_id contains non-unique values. Expected {} unique values, got {}".format(expected, actual)
        assert expected == actual, error_message

class TestOutliers(object):
    expected_boundaries = istanbul_districts.geometry.total_bounds

    def test_price_outliers_min(self):
        actual = airbnb.loc[:,"price"].min()
        error_message = "Minimum price must be > 0. Got {}".format(actual)
        assert actual > 0, error_message

    def test_latitude_boundary_min(self):
        expected = self.expected_boundaries[1]
        actual = airbnb["latitude"].min()
        error_message = "At least one point is smaller than the minimum boundary latitude. Actual {} is smaller than expected {}".format(actual, expected)
        assert expected <= actual, error_message

    def test_latitude_boundary_max(self):
        expected = self.expected_boundaries[3]
        actual = airbnb["latitude"].max()
        error_message = "At least one point is bigger than the maximum boundary latitude. Actual {} is bigger than expected {}".format(actual, expected)
        assert  expected >= actual, error_message

    def test_longitude_boundary_min(self):
        expected = self.expected_boundaries[0]
        actual = airbnb["longitude"].min()
        error_message = "At least one point is smaller than the minimum boundary longitude. Actual {} is smaller than expected {}".format(actual, expected)
        assert expected <= actual, error_message

    def test_longitude_boundary_max(self):
        expected = self.expected_boundaries[2]
        actual = airbnb["longitude"].max()
        error_message = "At least one point is bigger than the maximum boundary longitude. Actual {} is bigger than expected {}".format(actual, expected)
        assert  expected >= actual, error_message

class TestDataTypes(object):
    def test_data_type_agreement_within_columns(self):
        for column_name in airbnb.columns:
            expected_dtype = type(airbnb[column_name][0])
            value_index = 0
            while value_index < len(airbnb[column_name]):
                value_type = type(airbnb[column_name][value_index])
                error_message = "Values in column \"{}\" are not all of same type. Value at index {} is type {}, expected type {}".format(column_name, value_index, value_type, expected_dtype)
                assert value_type == expected_dtype, error_message
                value_index += 1

class TestValueAgreement(object):
    def test_district_name_agreement(self):
        dataset_subset = airbnb.loc[:,["district_eng", "district_tr"]]
        for row in dataset_subset.values:
            district_eng = row[0]
            district_tr = row[1]
            actual_similarity = textdistance.jaro_winkler.normalized_similarity(district_tr, district_eng)
            similarity_threshold = 0.50
            error_message = "District name similarity is below similarity treshold. Threshold is {}, similarity is {}".format(similarity_threshold, actual_similarity)
            assert actual_similarity >= similarity_threshold, error_message


###############################################################################
# WaterTAP Copyright (c) 2021, The Regents of the University of California,
# through Lawrence Berkeley National Laboratory, Oak Ridge National
# Laboratory, National Renewable Energy Laboratory, and National Energy
# Technology Laboratory (subject to receipt of any required approvals from
# the U.S. Dept. of Energy). All rights reserved.
#
# Please see the files COPYRIGHT.md and LICENSE.md for full copyright and license
# information, respectively. These files are also available online at the URL
# "https://github.com/watertap-org/watertap/"
#
###############################################################################

import pytest

from pyomo.environ import ConcreteModel, Var, Constraint

from idaes.core import FlowsheetBlock
from idaes.core.util import get_solver
import idaes.logger as idaeslog

from watertap.util.initialization import (check_solve, check_dof,
        assert_no_degrees_of_freedom, assert_degrees_of_freedom)

__author__ = "Adam Atia"

_log = idaeslog.getLogger(__name__)

# Set up solver
solver = get_solver()

class TestCheckDOF:
    @pytest.fixture(scope="class")
    def m(self):
        m = ConcreteModel()
        m.a = Var()
        m.b = Var()
        m.abcon = Constraint(rule=m.a + m.b == 10)
        return m

    @pytest.mark.unit
    def test_expected(self, m):
        check_dof(m, fail_flag=False, expected_dof=1)
        check_dof(m, fail_flag=True, expected_dof=1)
        assert_degrees_of_freedom(m, 1)

    @pytest.mark.unit
    def test_more_expected(self, m):
        check_dof(m, fail_flag=False, expected_dof=3)
        msg = r"Unexpected degrees of freedom: Degrees of freedom on unknown = 1. Expected 3. Unfix 2 variable\(s\)"
        with pytest.raises(ValueError, match=msg):
            check_dof(m, fail_flag=True, expected_dof=3)
        with pytest.raises(ValueError, match=msg):
            assert_degrees_of_freedom(m, 3)

    @pytest.mark.unit
    def test_less_expected(self, m):
        check_dof(m, fail_flag=False, expected_dof=-1)
        msg = r"Unexpected degrees of freedom: Degrees of freedom on unknown = 1. Expected -1. Fix 2 variable\(s\)"
        with pytest.raises(ValueError, match=msg):
            check_dof(m, fail_flag=True, expected_dof=-1)
        with pytest.raises(ValueError, match=msg):
            assert_degrees_of_freedom(m, -1)

    @pytest.mark.unit
    def test_zero_expected(self, m):
        # check_dof should pass since fail_flag=False produces warning for DOF!=0
        check_dof(m, fail_flag=False)
        msg = r"Non-zero degrees of freedom: Degrees of freedom on unknown = 1. Fix 1 more variable\(s\)"
        # Verify error is raised since DOF!=0
        with pytest.raises(ValueError, match=msg):
            check_dof(m, fail_flag=True)
        with pytest.raises(ValueError, match=msg):
            assert_no_degrees_of_freedom(m)

    @pytest.mark.unit
    def test_zero(self, m):
        m.a.fix(5)
        # check should pass since DOF=0
        check_dof(m, fail_flag=True)
        check_dof(m, fail_flag=True, expected_dof=0)
        assert_no_degrees_of_freedom(m)

        m.a.unfix()

class TestCheckSolve:
    @pytest.fixture(scope="class")
    def m(self):
        m = ConcreteModel()
        m.a = Var()
        m.acon = Constraint(rule=m.a >= 10)
        m.bcon = Constraint(rule=m.a == 5)
        return m

    @pytest.mark.unit
    def test_failure(self, m):
        results = solver.solve(m)
        # check_solve should pass since fail_flag=False and only warning will be produced
        check_solve(results, logger=_log, fail_flag=False)
        # expect the solve to fail and raise error
        with pytest.raises(ValueError, match="The solver failed to converge to an optimal solution. This suggests that the "
                                             "user provided infeasible inputs or that the model is poorly scaled."):
            check_solve(results, logger=_log, fail_flag=True)

    @pytest.mark.unit
    def test_failure_checkpoint(self, m):
        results = solver.solve(m)
        # check_solve should pass since fail_flag=False and only warning will be produced
        check_solve(results, checkpoint='test', logger=_log, fail_flag=False)
        # expect the solve to fail and raise error
        with pytest.raises(ValueError, match="test failed. The solver failed to converge to an optimal solution. "
                "This suggests that the user provided infeasible inputs or that the model is poorly scaled."):
            check_solve(results, checkpoint='test', logger=_log, fail_flag=True)

    @pytest.mark.unit
    def test_success(self, m):
        m.acon.deactivate()

        results = solver.solve(m)
        # both check_solve's should pass
        check_solve(results, logger=_log, fail_flag=False)
        check_solve(results, logger=_log, fail_flag=True)

        m.acon.activate()


"""
Checking Imported Data III - DataFrame Labels
"""

import numpy as np
import pandas as pd
wine_reviews = pd.read_csv('../../winemag-data-130k.csv')

# Access the labels on the rows of data.

print(wine_reviews.index) # RangeIndex(start=0, stop=129971, step=1)

# Access the labels on the columns of data.

print(wine_reviews.columns) # Index(['country', 'description', 'designation', 'points', 'price', 'province', 'region_1', 'region_2', 'taster_name', 'taster_twitter_handle', 'title', 'variety', 'winery'], dtype='object')

# Return the labels for the rows and columns in wine_reviews in one command.

print(wine_reviews.axes)
"""
[RangeIndex(start=0, stop=129971, step=1), Index(['country', 'description', 'designation', 'points', 'price', 'province', 'region_1', 'region_2', 'taster_name', 'taster_twitter_handle', 'title', 'variety', 'winery'], dtype='object')
"""


# import altair as alt
# import numpy as np
# import pandas as pd
# import pytest
from vega_datasets import data
from simpler_eda.corr_map import corr_map

df = data.cars()

def test_input_type():
    """
    Tests for exceptions handling.

    Returns
    -------
    None
        All test should pass and no asserts error should be displayed.
    """

    try:
        corr_map(
            ["a", "b", "c"],
            ["Horsepower", "Displacement", "Cylinders", "Acceleration"],
        )

    except Exception as err:
        assert str(err) == "The input data is not a panda dataframe"

    try:
        corr_map(df, "ab")

    except Exception as err:
        assert str(err) == "The input for feature should be a list"

    try:
        corr_map(df, ["Horsepower"])

    except Exception as err:
        assert str(err) == "There should be at least 2 features in the list"

    try:
        corr_map(df, ["Horsepower", [123, 456]])

    except Exception as err:
        assert (
            str(err)
            == "All the entries in the feature list should be a string"
        )

    # added a check to ensure the input list of features are all numeric
    try:
        corr_map(df, ["Horsepower", "Name"])

    except Exception as err:
        assert str(err) == "All features in the list should be numeric"

    try:
        corr_map(
            df,
            ["Horsepower", "Displacement", "Cylinders", "Acceleration"],
            corr_method="kvb",
        )

    except Exception as err:
        assert (
            str(err)
            == """The correlation method should be 'pearson', 'kendall',
            or 'spearman' """
        )

    try:
        corr_map(
            df,
            ["Horsepower", "Displacement", "Cylinders", "Acceleration"],
            color_scheme=["blueyellow"],
        )

    except Exception as err:
        assert str(err) == "The color scheme should be given as a string"

    try:
        corr_map(
            df,
            ["Horsepower", "Displacement", "Cylinders", "Acceleration"],
            plot_width="200",
        )

    except Exception as err:
        assert str(err) == "The plot_width should be given as an integer"

    try:
        corr_map(
            df,
            ["Horsepower", "Displacement", "Cylinders", "Acceleration"],
            plot_height="200",
        )

    except Exception as err:
        assert str(err) == "The plot_height should be given as an integer"

    try:
        corr_map(
            df,
            ["Horsepower", "Displacement", "Cylinders", "Acceleration"],
            title=200,
        )

    except Exception as err:
        assert str(err) == "The title should be given as a string"

out = corr_map(
    df,
    ["Horsepower", "Displacement", "Cylinders", "Acceleration"],
    corr_method="pearson",
    color_scheme="blueorange",
    plot_width=450,
    plot_height=450,
    title="Correlation Map",
)

def test_corr_map():
    """
    Tests the corr_map function to make sure the outputs are correct.

    Returns
    --------
    None
        All test should pass and no asserts should be displayed.s
    """
    assert (
        str(type(out)) == "<class 'altair.vegalite.v4.api.Chart'>"
    ), "The function should retrun an altair plot"
    assert (
        out.encoding.x.shorthand == "level_0"
    ), "The level_0 should be mapped to the x axis"
    assert (
        out.encoding.y.shorthand == "level_1"
    ), "The level 1 should be mapped to the y axis"
    assert (
        out.mark.type == "rect"
    ), "the plot type (mark) should be a rect plot (heatmap)"
    assert out.encoding.color.scale.scheme == "blueorange", (
        "the color scheme should blueorange (default) or the inputted color"
        " scheme"
    )
    assert (
        out.title == "Correlation Map"
    ), "The title should be Correlation Map or the given value"


import warnings
warnings.simplefilter("ignore")

import numpy as np
import mlrose_hiive as mlrose
import numpy as np
from matplotlib import pyplot as plt
import pandas as pd
import os

def plot_ga_param_search(runs_ga, problem_name, outdir, limit_iter1=3000, limit_iter2=5000):

    pop = [100, 200, 300]
    mut = [0.1, 0.2, 0.3]

    combinations = [(p, m) for p in pop for m in mut]
    color = {100: 'r', 200:'b', 300:'g'}
    marker = {0.1: 'o', 0.2:'+', 0.3:'x'}

    fig, ax = plt.subplots(1,2, figsize=(15,6))
    #fig.figure(figsize=(10,10))
    for p,m in combinations:
        runs = runs_ga[(runs_ga['Population Size'] == p) & (runs_ga['Mutation Rate'] == m) & (runs_ga['Iteration'] <= limit_iter1)]

        ax[0].plot(runs['Iteration'], runs['Fitness'], c=color[p], marker=marker[m], ls='--',label='Population Size='+str(p)+' Mutation Rate='+str(m))
        #ax[0].legend()
        ax[0].set_xlabel('Iteration')
        ax[0].set_ylabel('Fitness')
        ax[0].set_title('GA-'+problem_name+'-Fitness')
        ax[0].grid()
        ax[0].legend()

    for p,m in combinations:
        runs = runs_ga[(runs_ga['Population Size'] == p) & (runs_ga['Mutation Rate'] == m)  & (runs_ga['Iteration'] <= limit_iter2)]

        ax[1].plot(runs['Iteration'], runs['Time'],  c=color[p], marker=marker[m], ls='--', label='Population Size='+str(p)+' Mutation Rate='+str(m))
        ax[1].set_xlabel('Iteration')
        ax[1].set_ylabel('Time')
        ax[1].set_title('GA-'+problem_name+'-Time')
        ax[1].grid()
        #ax[1].legend()
    #plt.show()
    plt.savefig(os.path.join(outdir, problem_name+'_GA_param_search.png'))
    plt.close()

def plot_sa_param_search(runs, problem_name, outdir, limit_iter1=10000, limit_iter2=5000):

    temp = [100, 250, 500]
    decay = ["exponential", "geometric"]

    combinations = [(t, d) for t in temp for d in decay]
    color = {100: 'r', 250:'b', 500:'g'}
    marker = {"exponential": 'o', "geometric":'+'}

    fig, ax = plt.subplots(1,2, figsize=(15,6))
    #fig.figure(figsize=(10,10))
    for t,d in combinations:
        run = runs[(runs['Temperature'] == t) & (runs['schedule_type'] == d) & (runs['Iteration'] <= limit_iter1)]
        #print(run[['Temperature', 'schedule_type']].head())
        ax[0].plot(run['Iteration'], run['Fitness'], c=color[t], marker=marker[d], ls='--',
                   label='Temperature='+str(t)+' Decay='+str(d))
        #ax[0].legend()
        ax[0].set_xlabel('Iteration')
        ax[0].set_ylabel('Fitness')
        ax[0].set_title('SA-'+problem_name+'-Fitness')
        ax[0].grid()
        ax[0].legend()

    for t,d in combinations:
        run = runs[(runs['Temperature'] == t) & (runs['schedule_type'] == d)  & (runs['Iteration'] <= limit_iter2)]

        ax[1].plot(run['Iteration'], run['Time'],  c=color[t], marker=marker[d], ls='--',
                   label='Temperature='+str(t)+' Decay='+str(d))
        ax[1].set_xlabel('Iteration')
        ax[1].set_ylabel('Time')
        ax[1].set_title('SA-'+problem_name+'-Time')
        ax[1].grid()
        #ax[1].legend()
    ##plt.show()
    plt.savefig(os.path.join(outdir, problem_name+'_SA_param_search.png'))
    plt.close()

def plot_rhc_param_search(rhc_curves_runs, rhc_runs_stats, problem_name, outdir, limit_iter1=np.inf, limit_iter2=np.inf):
    restarts = np.arange(1,11)

    #color = {0: 'r', 10:'b', 50:'g'}
    markers = {0:'o', 1:'^', 2:'s', 3:'+', 4:'*', 5:'x', 6:'p', 7:'h', 8:'<', 9:'>', 10:''}
    #markers = {i:'$'+str(i)+'$' for i in restarts}
    #print(markers)

    fig, ax = plt.subplots(1,2, figsize=(15,6))
    for t in restarts:
        run = rhc_curves_runs[(rhc_curves_runs['current_restart'] == t) & (rhc_curves_runs['Iteration'] <=limit_iter1)]
        run = run.iloc[::300, :]
        ax[0].plot(run['Iteration']-run.iloc[0]['Iteration'], run['Fitness'], ls='--', marker=markers[t],
                   label='Restart='+str(t))
        #ax[0].plot(run['Iteration'], run['Fitness'], ls='--',  marker=markers[t],
        #           label='Restart='+str(t))
        ax[0].set_xlabel('Iteration')
        ax[0].set_ylabel('Fitness')
        ax[0].set_title('RHC-'+problem_name+'-Fitness')
        ax[0].grid()
        ax[0].legend()

    for t in restarts:
        run = rhc_runs_stats[(rhc_runs_stats['current_restart'] == t) & (rhc_runs_stats['Iteration'] <= limit_iter2)]

        #ax[1].plot(run['Iteration']-run.iloc[0]['Iteration'], run['Time']-run.iloc[0]['Time'],  ls='--',
        #           label='Restart='+str(t))
        ax[1].plot(run['Iteration'], run['Time']-run.iloc[0]['Time'],  ls='--',  marker=markers[t],
               label='Restart='+str(t))

        ax[1].set_xlabel('Iteration')
        ax[1].set_ylabel('Time')
        ax[1].set_title('RHC-'+problem_name+'-Time')
        ax[1].grid()
        ax[1].legend()
    #plt.show()
    plt.savefig(os.path.join(outdir, problem_name+'_RHC_param_search.png'))
    plt.close()

def plot_mimic_param_search(runs, problem_name, outdir, limit_iter1=1000, limit_iter2=2000):
    keep_pct = [0.1, 0.2, 0.3]
    population_sizes=[100, 200, 300]
    combinations = [(t, p) for t in keep_pct for p in population_sizes]

    color = {0.1: 'r', 0.2:'b', 0.3:'g'}
    marker = {100: '^', 200:'x', 300:'+'}

    fig, ax = plt.subplots(1,2, figsize=(15,6))
    #fig.figure(figsize=(10,10))
    for t, p in combinations:
        run = runs[(runs['Keep Percent'] == t) & (runs['Population Size'] == p) &  (runs['Iteration'] <= limit_iter1)]
        ax[0].plot(run['Iteration'], run['Fitness'],
                   c=color[t], marker=marker[p],
                   ls='--',
                   label='Keep Pct='+str(t)+' Population Size='+str(p))
        #ax[0].legend()
        ax[0].set_xlabel('Iteration')
        ax[0].set_ylabel('Fitness')
        ax[0].set_title('MIMIC-'+problem_name+'-Fitness')
        ax[0].grid()
        ax[0].legend()

    for t, p in combinations:
        run = runs[(runs['Keep Percent'] == t) & (runs['Population Size'] == p) &  (runs['Iteration'] <= limit_iter2)]

        ax[1].plot(run['Iteration'], run['Time'],
                   c=color[t], marker=marker[p],
                   ls='--',
                   label='Keep Pct='+str(t)+' Population Size='+str(p))
        ax[1].set_xlabel('Iteration')
        ax[1].set_ylabel('Time')
        ax[1].set_title('MIMIC-'+problem_name+'-Time')
        ax[1].grid()
        ax[1].legend()
    #plt.show()
    plt.savefig(os.path.join(outdir, problem_name+'_MIMIC_param_search.png'))
    plt.close()

def plot_best_models(ga_best, sa_best, rhc_best, mimic_best, ax, problem_name):
    ax[0].plot(ga_best['Iteration'], ga_best['Fitness'], 'ro--', label='GA best')
    ax[0].plot(sa_best['Iteration'], sa_best['Fitness'], 'b^--', label='SA best')
    ax[0].plot(rhc_best['Iteration'], rhc_best['Fitness'], 'yd--', label='RHC best')
    ax[0].plot(mimic_best['Iteration'], mimic_best['Fitness'], 'cs--', label='MIMIC best')
    ax[0].legend()
    ax[0].grid() #ax[0].set_xscale('log')
    ax[0].set_xlabel('Iteration')
    ax[0].set_ylabel('Fitesss')
    ax[0].set_title(problem_name+' Fitness Comparison')

    ax[1].plot(ga_best['Iteration'], ga_best['Time'], 'ro--', label='GA best')
    ax[1].plot(sa_best['Iteration'], sa_best['Time'], 'b^--', label='SA best')
    ax[1].plot(rhc_best['Iteration'], rhc_best['Time'] - rhc_best.iloc[0]['Time'], 'yd--', label='RHC best')
    ax[1].plot(mimic_best['Iteration'], mimic_best['Time'], 'cs--', label='MIMIC best')
    ax[1].legend()
    ax[1].grid()
    ax[1].set_xlabel('Iteration')
    ax[1].set_ylabel('Time')
    ax[1].set_title(problem_name+' Time Comparison')

def run_param_search(problem, output_dir):
    rhc = mlrose.RHCRunner(problem=problem, experiment_name="RHC",
        output_directory=output_dir,
        seed=42,
        iteration_list=2 ** np.arange(15),
        max_attempts=1000,
        restart_list=[10])
    rhc_run_stats, rhc_run_curves = rhc.run()
    sa = mlrose.SARunner(problem=problem,
        experiment_name="SA",
        output_directory=output_dir,
        seed=42,
        iteration_list=2 ** np.arange(20),
        max_attempts=1000,
        temperature_list=[100, 250, 500],
        decay_list=[mlrose.ExpDecay, mlrose.GeomDecay])
    sa_run_stats, sa_run_curves = sa.run()
    ga = mlrose.GARunner(problem=problem,
        experiment_name="GA",
        output_directory=output_dir,
        seed=42,
        iteration_list=2 ** np.arange(13),
        max_attempts=1000,
        population_sizes=[100, 200, 300],
        mutation_rates=[0.1, 0.2, 0.3])
    ga_run_stats, ga_run_curves = ga.run()
    mimic = mlrose.MIMICRunner(problem=problem,
        experiment_name="MIMIC",
        output_directory=output_dir,
        seed=42,
        iteration_list=2 ** np.arange(13),
        population_sizes=[100, 200, 300],
        max_attempts=500,
        keep_percent_list=[0.1, 0.2, 0.3],
        use_fast_mimic=True)
    mimic_run_stats, mimic_run_curves = mimic.run()