import json
import unittest

from connectome.analysis import (AnalysisBuilder, ConnectivityEstimator, RelativeReciprocityEstimator,
                                 RelativeCycleAnalysis, InOutDegreeCorrelation)
from connectome.model import ER, Network
from connectome.noise import OutDegreePreservingNoise
from connectome.builder import InvalidModelType
from connectome.builder import PipelineBuilder, concatenate_model_noise_analysis


config = """
{
    "models": {
        "joint_parameters": {
            "nr_neurons": 300,
            "inh_ratio": 0.33333333333333333333333333,
            "p_exc": {
                "type": "uniform",
                "args": [0.15, 0.1]
            },
            "p_inh": {
                "type": "uniform",
                "args": [0.54, 0.02]
            }
        },
        "model_specific_parameters": {
            "ER": {
            },
            "FF": {
                "nr_layers":  {
                    "type": "truncnorm",
                    "args": [],
                    "kwargs": {"a": -2, "b": 2, "loc": 1500, "scale": 150}
                }
            },
            "LL": {
            }
        }
    },
    "noise": {
        "type": "NoNoise",
        "parameters": {
        }
    },
    "analysis": {
        "NrNeurons": {
        },
        "ConnectivityEstimator": {
        }
    }
}
"""


noisyconfig = """
{
    "models": {
        "joint_parameters": {
            "nr_neurons": 300,
            "inh_ratio": 0.33333333333333333333333333,
            "p_exc": {
                "type": "uniform",
                "args": [0.15, 0.1]
            },
            "p_inh": {
                "type": "uniform",
                "args": [0.54, 0.02]
            }
        },
        "model_specific_parameters": {
            "ER": {
            },
            "FF": {
                "nr_layers":  {
                    "type": "truncnorm",
                    "args": [],
                    "kwargs": {"a": -2, "b": 2, "loc": 1500, "scale": 150}
                }
            },
            "LL": {
            }
        }
    },
    "noise": {
    "type": "RemoveAddNoiseAndSubsample",
    "parameters": {
      "fraction_remove_and_add": {
          "type": "beta",
          "args": [
            2,
            10
          ]
        },
      "subsampling_fraction": {
          "type": "uniform",
          "args": [
            0.05,
            0.05
          ]
        }
    }
  },
    "analysis": {
        "NrNeurons": {
        },
        "ConnectivityEstimator": {
        }
    }
}
"""

class TestModelBuilder(unittest.TestCase):
    def setUp(self):
        self.modelbuilder = PipelineBuilder(json.loads(config))

    def test_build_ER_check_fixed_parameter(self):
        er = self.modelbuilder.build_pipeline("ER")
        er_prior = self.modelbuilder.build_prior("ER")
        parameter = er_prior.rvs()
        analysis_result = er(parameter)
        self.assertEqual(200, analysis_result['nr_exc'])
        self.assertEqual(100, analysis_result['nr_inh'])

    def test_build_ER_check_sampled_parameter(self):
        er = self.modelbuilder.build_pipeline("ER")
        er_prior = self.modelbuilder.build_prior("ER")
        parameter = er_prior.rvs()
        analysis_result = er(parameter)
        self.assertTrue(.15 <= analysis_result['p_ee'] <=.25)
        self.assertTrue(.5 <= analysis_result['p_ii'] <=.6)

    def test_defined_models(self):
        self.assertEqual(["ER", "FF", "LL"],sorted(self.modelbuilder.defined_models()))


class TestModelBuilderInvalid(unittest.TestCase):
    def test_invalid_model(self):
        dct = json.loads(config)
        dct['models']['model_specific_parameters']['INVALIDMODEL'] = {}
        modelbuilder = PipelineBuilder(dct)
        with self.assertRaises(InvalidModelType):
            modelbuilder.build_pipeline("INVALIDMODEL")

    def test_invalid_model_not_Defined(self):
        dct = json.loads(config)
        modelbuilder = PipelineBuilder(dct)
        with self.assertRaises(InvalidModelType):
            modelbuilder.build_pipeline("NOTDEFINED")


class TestConcatenateModelNoiseAnalysis(unittest.TestCase):
    def make_analysis(self):
        analysis_builder = AnalysisBuilder()
        analysis_builder.add(ConnectivityEstimator)
        analysis_builder.add(RelativeReciprocityEstimator)
        analysis_builder.add(RelativeCycleAnalysis, length=5)
        analysis_builder.add(InOutDegreeCorrelation)
        return analysis_builder.get_result()

    def setUp(self):
        self.model = concatenate_model_noise_analysis(ER(inh_ratio=.1), OutDegreePreservingNoise(), self.make_analysis())

    def test_result_keys(self):
        result = self.model(network_nr_neurons=200, network_p_exc=.1, network_p_inh=.2,
                            noise_fraction_draws_over_nr_synapses=.2)
        self.assertEqual({"relative_reciprocity_ee",
                          "relative_reciprocity_ii",
                          "relative_reciprocity_ei",
                          "relative_reciprocity_ie",
                          "p_ee",
                          "p_ii",
                          "p_ei",
                          "p_ie",
                          "in_out_degree_correlation_exc",
                          "relative_cycles_5"}, set(result.keys()))

    def test_different_results(self):
        result1 = self.model(network_nr_neurons=200, network_p_exc=.1, network_p_inh=.2,
                                noise_fraction_draws_over_nr_synapses=.2)
        result2 = self.model()
        self.assertNotEqual(result1, result2)

    def test_p_ee_in_range(self):
        result = self.model(network_nr_neurons=200, network_p_exc=.1, network_p_inh=.2,
                                noise_fraction_draws_over_nr_synapses=.2)
        self.assertLess(abs(result['p_ee'] - .1)/.1, .05)


class TestNoisyModel(unittest.TestCase):
    def test_noisy_model(self):
        builder = PipelineBuilder(json.loads(noisyconfig))
        model = builder.build_model_with_noise("ER")
        prior = builder.build_prior("ER")
        par = prior.rvs()
        res = model(par)
        self.assertTrue(isinstance(res, Network))



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