import logging
import os
import random
import string
import time
import unittest

import neurolib.utils.paths as paths
import neurolib.utils.pypetUtils as pu
import numpy as np
import pytest
import xarray as xr
from neurolib.models.aln import ALNModel
from neurolib.models.fhn import FHNModel
from neurolib.models.multimodel import MultiModel
from neurolib.models.multimodel.builder.fitzhugh_nagumo import FitzHughNagumoNetwork
from neurolib.optimize.exploration import BoxSearch
from neurolib.utils.loadData import Dataset
from neurolib.utils.parameterSpace import ParameterSpace


def randomString(stringLength=10):
    """Generate a random string of fixed length"""
    letters = string.ascii_lowercase
    return "".join(random.choice(letters) for i in range(stringLength))


class TestBoxSearch(unittest.TestCase):
    """
    Basic tests.
    """

    def test_assertions(self):
        parameters = ParameterSpace(
            {"mue_ext_mean": np.linspace(0, 3, 2), "mui_ext_mean": np.linspace(0, 3, 2)}, kind="sequence"
        )
        with pytest.raises(AssertionError):
            _ = BoxSearch(model=None, parameterSpace=parameters)

        with pytest.raises(AssertionError):
            _ = BoxSearch(model=None, parameterSpace=None)

        with pytest.raises(AssertionError):
            _ = BoxSearch(model=None, parameterSpace=parameters, evalFunction=None)

    def test_fillin_default_parameters_for_sequential(self):
        in_dict = {"a": [None, None, 1, 2], "b": [4, 5, None, None]}
        SHOULD_BE = {"a": [0, 0, 1, 2], "b": [4, 5, 12, 12]}
        model_params = {"a": 0, "b": 12}

        parameters = ParameterSpace({"mue_ext_mean": [1.0, 2.0]})
        search = BoxSearch(model=ALNModel(), parameterSpace=parameters)
        out_dict = search._fillin_default_parameters_for_sequential(in_dict, model_params)
        self.assertDictEqual(out_dict, SHOULD_BE)


class TestExplorationSingleNode(unittest.TestCase):
    """
    ALN single node exploration.
    """

    def test_single_node(self):
        start = time.time()

        model = ALNModel()
        parameters = ParameterSpace({"mue_ext_mean": np.linspace(0, 3, 2), "mui_ext_mean": np.linspace(0, 3, 2)})
        search = BoxSearch(model, parameters, filename="test_single_nodes.hdf")
        search.run()
        search.loadResults()
        dataarray = search.xr()
        self.assertTrue(isinstance(dataarray, xr.DataArray))
        self.assertFalse(dataarray.attrs)

        for i in search.dfResults.index:
            search.dfResults.loc[i, "max_r"] = np.max(
                search.results[i]["rates_exc"][:, -int(1000 / model.params["dt"]) :]
            )

        end = time.time()
        logging.info("\t > Done in {:.2f} s".format(end - start))


class TestExplorationBrainNetwork(unittest.TestCase):
    """
    FHN brain network simulation with BOLD simulation.
    """

    def test_fhn_brain_network_exploration(self):
        ds = Dataset("hcp")
        model = FHNModel(Cmat=ds.Cmat, Dmat=ds.Dmat)
        model.params.duration = 10 * 1000  # ms
        model.params.dt = 0.2
        model.params.bold = True
        parameters = ParameterSpace(
            {
                "x_ext": [np.ones((model.params["N"],)) * a for a in np.linspace(0, 2, 2)],
                "K_gl": np.linspace(0, 2, 2),
                "coupling": ["additive", "diffusive"],
            },
            kind="grid",
        )
        search = BoxSearch(model=model, parameterSpace=parameters, filename="test_fhn_brain_network_exploration.hdf")

        search.run(chunkwise=True, bold=True)

        pu.getTrajectorynamesInFile(os.path.join(paths.HDF_DIR, "test_fhn_brain_network_exploration.hdf"))
        search.loadDfResults()
        search.getRun(0, pypetShortNames=True)
        search.getRun(0, pypetShortNames=False)
        search.loadResults()
        # firing rate xr
        dataarray = search.xr()
        self.assertTrue(isinstance(dataarray, xr.DataArray))
        self.assertFalse(dataarray.attrs)
        # bold xr
        dataarray = search.xr(bold=True)
        self.assertTrue(isinstance(dataarray, xr.DataArray))
        self.assertFalse(dataarray.attrs)
        search.info()


class TestExplorationBrainNetworkPostprocessing(unittest.TestCase):
    """
    ALN brain network simulation with custom evaluation function.
    """

    @classmethod
    def setUpClass(cls):
        # def test_brain_network_postprocessing(self):
        ds = Dataset("hcp")
        model = ALNModel(Cmat=ds.Cmat, Dmat=ds.Dmat)
        # Resting state fits
        model.params["mue_ext_mean"] = 1.57
        model.params["mui_ext_mean"] = 1.6
        model.params["sigma_ou"] = 0.09
        model.params["b"] = 5.0
        model.params["signalV"] = 2
        model.params["dt"] = 0.2
        model.params["duration"] = 0.2 * 60 * 1000

        # multi stage evaluation function
        def evaluateSimulation(traj):
            model = search.getModelFromTraj(traj)
            model.randomICs()
            model.params["dt"] = 0.2
            model.params["duration"] = 4 * 1000.0
            model.run(bold=True)

            result_dict = {"outputs": model.outputs}

            search.saveToPypet(result_dict, traj)

        # define and run exploration
        parameters = ParameterSpace({"mue_ext_mean": np.linspace(0, 3, 2), "mui_ext_mean": np.linspace(0, 3, 2)})
        search = BoxSearch(
            evalFunction=evaluateSimulation,
            model=model,
            parameterSpace=parameters,
            filename=f"test_brain_postprocessing_{randomString(20)}.hdf",
        )
        search.run()
        cls.model = model
        cls.search = search
        cls.ds = ds

    def test_getRun(self):
        self.search.getRun(0)

    def test_loadResults(self):
        self.search.loadResults()

    def test_loadResults_all_False(self):
        self.search.loadResults(all=False)


class TestCustomParameterExploration(unittest.TestCase):
    """Exploration with custom function"""

    def test_circle_exploration(self):
        def explore_me(traj):
            pars = search.getParametersFromTraj(traj)
            # let's calculate the distance to a circle
            computation_result = abs((pars["x"] ** 2 + pars["y"] ** 2) - 1)
            result_dict = {"scalar_result": computation_result, "list_result": [1, 2, 3, 4], "array_result": np.ones(3)}
            search.saveToPypet(result_dict, traj)

        parameters = ParameterSpace({"x": np.linspace(-2, 2, 2), "y": np.linspace(-2, 2, 2)})
        search = BoxSearch(evalFunction=explore_me, parameterSpace=parameters, filename="test_circle_exploration.hdf")
        search.run()
        search.loadResults(pypetShortNames=False)

        # call the result dataframe
        search.dfResults

        # test integrity of dataframe
        for i in search.dfResults.index:
            self.assertEqual(search.dfResults.loc[i, "scalar_result"], search.results[i]["scalar_result"])
            self.assertListEqual(search.dfResults.loc[i, "list_result"], search.results[i]["list_result"])
            np.testing.assert_array_equal(search.dfResults.loc[i, "array_result"], search.results[i]["array_result"])


class TestExplorationMultiModel(unittest.TestCase):
    """
    MultiModel exploration test - uses FHN network.
    """

    def test_multimodel_explore(self):
        start = time.time()

        DELAY = 13.0
        fhn_net = FitzHughNagumoNetwork(np.random.rand(2, 2), np.array([[0.0, DELAY], [DELAY, 0.0]]))
        model = MultiModel(fhn_net)
        parameters = ParameterSpace({"*input*sigma": [0.0, 0.05], "*epsilon*": [0.5, 0.6]}, allow_star_notation=True)
        search = BoxSearch(model, parameters, filename="test_multimodel.hdf")
        search.run()
        search.loadResults()
        dataarray = search.xr()
        self.assertTrue(isinstance(dataarray, xr.DataArray))
        self.assertTrue(isinstance(dataarray.attrs, dict))
        self.assertListEqual(
            list(dataarray.attrs.keys()),
            [k.replace("*", "_").replace(".", "_").replace("|", "_") for k in parameters.dict().keys()],
        )

        end = time.time()
        logging.info("\t > Done in {:.2f} s".format(end - start))


class TestExplorationMultiModelSequential(unittest.TestCase):
    """
    MultiModel exploration test with sequential exploration - uses FHN network.
    """

    def test_multimodel_explore(self):
        start = time.time()

        DELAY = 13.0
        fhn_net = FitzHughNagumoNetwork(np.random.rand(2, 2), np.array([[0.0, DELAY], [DELAY, 0.0]]))
        model = MultiModel(fhn_net)
        parameters = ParameterSpace(
            {"*input*sigma": [0.0, 0.05], "*epsilon*": [0.5, 0.6, 0.7]}, allow_star_notation=True, kind="sequence"
        )
        search = BoxSearch(model, parameters, filename="test_multimodel.hdf")
        search.run()
        search.loadResults()
        dataarray = search.xr()
        self.assertTrue(isinstance(dataarray, xr.DataArray))
        self.assertTrue("run_no" in dataarray.dims)
        self.assertEqual(len(dataarray["run_no"]), 5)
        self.assertTrue(isinstance(dataarray.attrs, dict))
        self.assertListEqual(
            list(dataarray.attrs.keys()),
            [k.replace("*", "_").replace(".", "_").replace("|", "_") for k in parameters.dict().keys()],
        )

        end = time.time()
        logging.info("\t > Done in {:.2f} s".format(end - start))


class TestExplorationSingleNodeSequential(unittest.TestCase):
    """
    ALN single node test with sequential exploration.
    """

    def test_single_node(self):
        start = time.time()

        model = ALNModel()
        parameters = ParameterSpace({"mue_ext_mean": [0.0, 1.5, 3.0], "mui_ext_mean": [1.5, 3.0]}, kind="sequence")
        search = BoxSearch(model, parameters, filename="test_single_nodes.hdf")
        search.run()
        search.loadResults()
        dataarray = search.xr()
        self.assertTrue(isinstance(dataarray, xr.DataArray))
        self.assertTrue("run_no" in dataarray.dims)
        self.assertEqual(len(dataarray["run_no"]), 5)
        self.assertFalse(dataarray.attrs)

        for i in search.dfResults.index:
            search.dfResults.loc[i, "max_r"] = np.max(
                search.results[i]["rates_exc"][:, -int(1000 / model.params["dt"]) :]
            )

        end = time.time()
        logging.info("\t > Done in {:.2f} s".format(end - start))


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

[((289, 4, 289, 19), 'unittest.main', 'unittest.main', ({}, {}), '()', False, 'import unittest\n'), ((51, 21, 51, 65), 'neurolib.utils.parameterSpace.ParameterSpace', 'ParameterSpace', ({(51, 36, 51, 64): "{'mue_ext_mean': [1.0, 2.0]}"}, {}), "({'mue_ext_mean': [1.0, 2.0]})", False, 'from neurolib.utils.parameterSpace import ParameterSpace\n'), ((63, 16, 63, 27), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((65, 16, 65, 26), 'neurolib.models.aln.ALNModel', 'ALNModel', ({}, {}), '()', False, 'from neurolib.models.aln import ALNModel\n'), ((67, 17, 67, 79), 'neurolib.optimize.exploration.BoxSearch', 'BoxSearch', (), '', False, 'from neurolib.optimize.exploration import BoxSearch\n'), ((79, 14, 79, 25), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((89, 13, 89, 27), 'neurolib.utils.loadData.Dataset', 'Dataset', ({(89, 21, 89, 26): '"""hcp"""'}, {}), "('hcp')", False, 'from neurolib.utils.loadData import Dataset\n'), ((90, 16, 90, 52), 'neurolib.models.fhn.FHNModel', 'FHNModel', (), '', False, 'from neurolib.models.fhn import FHNModel\n'), ((102, 17, 102, 117), 'neurolib.optimize.exploration.BoxSearch', 'BoxSearch', (), '', False, 'from neurolib.optimize.exploration import BoxSearch\n'), ((130, 13, 130, 27), 'neurolib.utils.loadData.Dataset', 'Dataset', ({(130, 21, 130, 26): '"""hcp"""'}, {}), "('hcp')", False, 'from neurolib.utils.loadData import Dataset\n'), ((131, 16, 131, 52), 'neurolib.models.aln.ALNModel', 'ALNModel', (), '', False, 'from neurolib.models.aln import ALNModel\n'), ((188, 17, 188, 118), 'neurolib.optimize.exploration.BoxSearch', 'BoxSearch', (), '', False, 'from neurolib.optimize.exploration import BoxSearch\n'), ((208, 16, 208, 27), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((212, 16, 212, 35), 'neurolib.models.multimodel.MultiModel', 'MultiModel', ({(212, 27, 212, 34): 'fhn_net'}, {}), '(fhn_net)', False, 'from neurolib.models.multimodel import MultiModel\n'), ((213, 21, 213, 117), 'neurolib.utils.parameterSpace.ParameterSpace', 'ParameterSpace', (), '', False, 'from neurolib.utils.parameterSpace import ParameterSpace\n'), ((214, 17, 214, 77), 'neurolib.optimize.exploration.BoxSearch', 'BoxSearch', (), '', False, 'from neurolib.optimize.exploration import BoxSearch\n'), ((225, 14, 225, 25), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((235, 16, 235, 27), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((239, 16, 239, 35), 'neurolib.models.multimodel.MultiModel', 'MultiModel', ({(239, 27, 239, 34): 'fhn_net'}, {}), '(fhn_net)', False, 'from neurolib.models.multimodel import MultiModel\n'), ((240, 21, 242, 9), 'neurolib.utils.parameterSpace.ParameterSpace', 'ParameterSpace', (), '', False, 'from neurolib.utils.parameterSpace import ParameterSpace\n'), ((243, 17, 243, 77), 'neurolib.optimize.exploration.BoxSearch', 'BoxSearch', (), '', False, 'from neurolib.optimize.exploration import BoxSearch\n'), ((256, 14, 256, 25), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((266, 16, 266, 27), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((268, 16, 268, 26), 'neurolib.models.aln.ALNModel', 'ALNModel', ({}, {}), '()', False, 'from neurolib.models.aln import ALNModel\n'), ((269, 21, 269, 115), 'neurolib.utils.parameterSpace.ParameterSpace', 'ParameterSpace', (), '', False, 'from neurolib.utils.parameterSpace import ParameterSpace\n'), ((270, 17, 270, 79), 'neurolib.optimize.exploration.BoxSearch', 'BoxSearch', (), '', False, 'from neurolib.optimize.exploration import BoxSearch\n'), ((284, 14, 284, 25), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((25, 19, 25, 41), 'random.choice', 'random.choice', ({(25, 33, 25, 40): 'letters'}, {}), '(letters)', False, 'import random\n'), ((37, 13, 37, 42), 'pytest.raises', 'pytest.raises', ({(37, 27, 37, 41): 'AssertionError'}, {}), '(AssertionError)', False, 'import pytest\n'), ((38, 16, 38, 64), 'neurolib.optimize.exploration.BoxSearch', 'BoxSearch', (), '', False, 'from neurolib.optimize.exploration import BoxSearch\n'), ((40, 13, 40, 42), 'pytest.raises', 'pytest.raises', ({(40, 27, 40, 41): 'AssertionError'}, {}), '(AssertionError)', False, 'import pytest\n'), ((41, 16, 41, 58), 'neurolib.optimize.exploration.BoxSearch', 'BoxSearch', (), '', False, 'from neurolib.optimize.exploration import BoxSearch\n'), ((43, 13, 43, 42), 'pytest.raises', 'pytest.raises', ({(43, 27, 43, 41): 'AssertionError'}, {}), '(AssertionError)', False, 'import pytest\n'), ((44, 16, 44, 83), 'neurolib.optimize.exploration.BoxSearch', 'BoxSearch', (), '', False, 'from neurolib.optimize.exploration import BoxSearch\n'), ((106, 36, 106, 105), 'os.path.join', 'os.path.join', ({(106, 49, 106, 62): 'paths.HDF_DIR', (106, 64, 106, 104): '"""test_fhn_brain_network_exploration.hdf"""'}, {}), "(paths.HDF_DIR, 'test_fhn_brain_network_exploration.hdf')", False, 'import os\n'), ((199, 12, 199, 117), 'numpy.testing.assert_array_equal', 'np.testing.assert_array_equal', ({(199, 42, 199, 81): "search.dfResults.loc[i, 'array_result']", (199, 83, 199, 116): "search.results[i]['array_result']"}, {}), "(search.dfResults.loc[i, 'array_result'],\n    search.results[i]['array_result'])", True, 'import numpy as np\n'), ((211, 40, 211, 60), 'numpy.random.rand', 'np.random.rand', ({(211, 55, 211, 56): '2', (211, 58, 211, 59): '2'}, {}), '(2, 2)', True, 'import numpy as np\n'), ((211, 62, 211, 100), 'numpy.array', 'np.array', ({(211, 71, 211, 99): '[[0.0, DELAY], [DELAY, 0.0]]'}, {}), '([[0.0, DELAY], [DELAY, 0.0]])', True, 'import numpy as np\n'), ((238, 40, 238, 60), 'numpy.random.rand', 'np.random.rand', ({(238, 55, 238, 56): '2', (238, 58, 238, 59): '2'}, {}), '(2, 2)', True, 'import numpy as np\n'), ((238, 62, 238, 100), 'numpy.array', 'np.array', ({(238, 71, 238, 99): '[[0.0, DELAY], [DELAY, 0.0]]'}, {}), '([[0.0, DELAY], [DELAY, 0.0]])', True, 'import numpy as np\n'), ((35, 29, 35, 49), 'numpy.linspace', 'np.linspace', ({(35, 41, 35, 42): '0', (35, 44, 35, 45): '3', (35, 47, 35, 48): '2'}, {}), '(0, 3, 2)', True, 'import numpy as np\n'), ((35, 67, 35, 87), 'numpy.linspace', 'np.linspace', ({(35, 79, 35, 80): '0', (35, 82, 35, 83): '3', (35, 85, 35, 86): '2'}, {}), '(0, 3, 2)', True, 'import numpy as np\n'), ((52, 33, 52, 43), 'neurolib.models.aln.ALNModel', 'ALNModel', ({}, {}), '()', False, 'from neurolib.models.aln import ALNModel\n'), ((66, 53, 66, 73), 'numpy.linspace', 'np.linspace', ({(66, 65, 66, 66): '0', (66, 68, 66, 69): '3', (66, 71, 66, 72): '2'}, {}), '(0, 3, 2)', True, 'import numpy as np\n'), ((66, 91, 66, 111), 'numpy.linspace', 'np.linspace', ({(66, 103, 66, 104): '0', (66, 106, 66, 107): '3', (66, 109, 66, 110): '2'}, {}), '(0, 3, 2)', True, 'import numpy as np\n'), ((97, 24, 97, 44), 'numpy.linspace', 'np.linspace', ({(97, 36, 97, 37): '0', (97, 39, 97, 40): '2', (97, 42, 97, 43): '2'}, {}), '(0, 2, 2)', True, 'import numpy as np\n'), ((154, 53, 154, 73), 'numpy.linspace', 'np.linspace', ({(154, 65, 154, 66): '0', (154, 68, 154, 69): '3', (154, 71, 154, 72): '2'}, {}), '(0, 3, 2)', True, 'import numpy as np\n'), ((154, 91, 154, 111), 'numpy.linspace', 'np.linspace', ({(154, 103, 154, 104): '0', (154, 106, 154, 107): '3', (154, 109, 154, 110): '2'}, {}), '(0, 3, 2)', True, 'import numpy as np\n'), ((184, 109, 184, 119), 'numpy.ones', 'np.ones', ({(184, 117, 184, 118): '(3)'}, {}), '(3)', True, 'import numpy as np\n'), ((187, 42, 187, 63), 'numpy.linspace', 'np.linspace', ({(187, 54, 187, 56): '-2', (187, 58, 187, 59): '2', (187, 61, 187, 62): '2'}, {}), '(-2, 2, 2)', True, 'import numpy as np\n'), ((187, 70, 187, 91), 'numpy.linspace', 'np.linspace', ({(187, 82, 187, 84): '-2', (187, 86, 187, 87): '2', (187, 89, 187, 90): '2'}, {}), '(-2, 2, 2)', True, 'import numpy as np\n'), ((96, 26, 96, 55), 'numpy.ones', 'np.ones', ({(96, 34, 96, 54): "(model.params['N'],)"}, {}), "((model.params['N'],))", True, 'import numpy as np\n'), ((96, 69, 96, 89), 'numpy.linspace', 'np.linspace', ({(96, 81, 96, 82): '0', (96, 84, 96, 85): '2', (96, 87, 96, 88): '2'}, {}), '(0, 2, 2)', True, 'import numpy as np\n')]