import glob
import json
import os

from snudda.neurons import NeuronMorphologyExtended
from snudda.utils.snudda_path import snudda_parse_path


class NeuronPrototype:
    """ Helper class, returns a neuron prototype based on parameter_id, morph_id and modulation_id """

    def __init__(self,
                 neuron_path,
                 neuron_name,
                 morphology_path=None,
                 parameter_path=None,
                 mechanism_path=None,
                 modulation_path=None,
                 reaction_diffusion_path=None,
                 snudda_data=None,
                 meta_path=None,
                 virtual_neuron=False,
                 load_morphology=True,
                 verbose=False):

        self.verbose = verbose
        self.snudda_data = snudda_data

        if neuron_path:
            self.neuron_path = snudda_parse_path(neuron_path, self.snudda_data)
        elif parameter_path:
            self.neuron_path = os.path.dirname(snudda_parse_path(parameter_path, self.snudda_data))
        elif morphology_path:
            # morphology path usually points to "morphology" directory, but it can also point to a specific morphology
            # which should be in neuron_path then (old format), but it might be in morphology folder...

            if os.path.isdir(snudda_parse_path(morphology_path, self.snudda_data)):
                self.neuron_path = os.path.dirname(snudda_parse_path(morphology_path, self.snudda_data))
            elif os.path.isfile(snudda_parse_path(morphology_path, self.snudda_data)):
                morph_path = os.path.dirname(snudda_parse_path(morphology_path, self.snudda_data))
                if "morphology" in morph_path:
                    self.neuron_path = os.path.dirname(morph_path)
                else:
                    self.neuron_path = morph_path
        else:
            self.neuron_path = None

        if morphology_path:
            self.morphology_path = snudda_parse_path(morphology_path, self.snudda_data)
        elif self.neuron_path:
            self.morphology_path = snudda_parse_path(os.path.join(self.neuron_path, "morphology"), self.snudda_data)
        else:
            self.morphology_path = None

        if mechanism_path:
            self.mechanism_path = snudda_parse_path(mechanism_path, self.snudda_data)
        elif self.neuron_path:
            self.mechanism_path = snudda_parse_path(os.path.join(self.neuron_path, "mechanisms.json"), self.snudda_data)
        else:
            self.mechanism_path = None

        if parameter_path:
            self.parameter_path = snudda_parse_path(parameter_path, self.snudda_data)
        elif self.neuron_path:
            self.parameter_path = snudda_parse_path(os.path.join(self.neuron_path, "parameters.json"), self.snudda_data)
        else:
            self.parameter_path = None

        if meta_path:
            self.meta_path = snudda_parse_path(meta_path, self.snudda_data)
        elif self.neuron_path:
            self.meta_path = snudda_parse_path(os.path.join(self.neuron_path, "meta.json"), self.snudda_data)
        else:
            self.meta_path = None

        if modulation_path:
            self.modulation_path = snudda_parse_path(modulation_path, self.snudda_data)
        elif self.neuron_path:
            self.modulation_path = snudda_parse_path(os.path.join(self.neuron_path, "modulation.json"), self.snudda_data)

            if not os.path.exists(self.modulation_path):
                self.modulation_path = None
        else:
            self.modulation_path = None

        if isinstance(reaction_diffusion_path, dict):
            self.reaction_diffusion_path = reaction_diffusion_path

        elif reaction_diffusion_path:
            self.reaction_diffusion_path = snudda_parse_path(reaction_diffusion_path, self.snudda_data)

        elif self.neuron_path:
            self.reaction_diffusion_path = snudda_parse_path(os.path.join(self.neuron_path, "reaction_diffusion.json"),
                                                             self.snudda_data)

            if not os.path.exists(self.reaction_diffusion_path):
                self.reaction_diffusion_path = None
        else:
            self.reaction_diffusion_path = None

        self.neuron_name = neuron_name
        self.parameter_info = None
        self.meta_info = None
        self.modulation_info = None
        self.virtual_neuron = virtual_neuron
        self.load_morphology = load_morphology

        self.morphology_cache = dict()
        self.morphology_lookup = dict()

        self.load_info()

    def load_info(self):
        """ Reads information about the neuron prototype. """

        self.morphology_cache = dict()
        self.morphology_lookup = dict()

        par_path = self.parameter_path


        if self.meta_path and os.path.exists(self.meta_path):
            try:
                with open(self.meta_path, "r") as fm:
                    self.meta_info = json.load(fm)
            except Exception as e:
                print(f"!! Attempting to load {self.meta_path}")
                import traceback
                print(traceback.format_exc())
                raise e

        if par_path is None or not os.path.exists(par_path):
            if self.verbose:
                print(f"Missing parameters.json : {par_path}")

            self.parameter_info = None
            return

        with open(par_path, "r") as f:
            self.parameter_info = json.load(f)

        if self.meta_info and isinstance(self.parameter_info, dict):
            # Check that all parameter_keys are in meta, otherwise remove them
            for par_key in list(self.parameter_info.keys()):
                if par_key not in self.meta_info:
                    del self.parameter_info[par_key]
                    if self.verbose:
                        print(f"Parameter key {par_key} missing in {self.meta_path}")

        # We now expect a dictionary of parameter sets. If it is a list, we convert it to a dictionary
        if type(self.parameter_info) == list:
            self.parameter_info = {"default": self.parameter_info}

        mod_path = self.modulation_path

        if mod_path is not None and os.path.exists(mod_path):
            with open(mod_path, "r") as f:
                self.modulation_info = json.load(f)
        else:
            self.modulation_info = None

    def get_num_morphologies(self, parameter_id=None, parameter_key=None):

        if parameter_key is None:
            parameter_key = self.get_parameter_key(parameter_id=parameter_id)
        elif parameter_id is not None:
            assert parameter_key == self.get_parameter_key(parameter_id=parameter_id), \
                f"Internal error, parameter_id {parameter_id} does not match parameter_key {parameter_key} specified."

        if self.meta_info and parameter_key:
            morph_key_list = self.meta_info[parameter_key].keys()
            return len(morph_key_list)
        else:
            return 1

    def get_parameter_key(self, parameter_id):

        if parameter_id is None:
            return None

        # assert parameter_id is not None, (f"get_parameter_key: parameter_id is None, for {self.neuron_name}"
        #                                   f"\nNeuron path: {self.neuron_path}")

        if self.parameter_info:
            par_key_list = sorted(list(self.parameter_info.keys()))
            par_key = par_key_list[parameter_id % len(par_key_list)]
        elif self.meta_info:
            if self.verbose:
                print(f"Missing parameter file for {self.parameter_path}, using meta.json for parameter key list.")
            par_key_list = sorted(list(self.meta_info.keys()))
            par_key = par_key_list[parameter_id % len(par_key_list)]
        else:
            par_key = None

        return par_key

    def get_morph_key(self,  morphology_id, parameter_id=None, parameter_key=None):

        if parameter_id is None:
            par_key = parameter_key
        else:
            par_key = self.get_parameter_key(parameter_id=parameter_id)
            assert parameter_key is None or par_key == parameter_key, \
                (f"parameter_id = {parameter_id} gives parameter_key {par_key}, "
                 f"different from parameter_key {parameter_key} provided")

        if self.meta_info:
            assert par_key in self.meta_info, f"Parameter key {par_key} missing in {self.meta_path}"
            morph_key_list = sorted(list(self.meta_info[par_key].keys()))
            assert len(morph_key_list) > 0, f"No morphologies available for parameter key {par_key} in {self.meta_path}"
            morph_key = morph_key_list[morphology_id % len(morph_key_list)]
        else:
            morph_key = None

        return morph_key

    def get_modulation_key(self, modulation_id,
                           parameter_id=None, morphology_id=None,
                           parameter_key=None, morphology_key=None):

        if modulation_id is None:
            return None

        if self.modulation_info:
            if type(self.modulation_info) == list:
                # Old format without keys
                return None

            if parameter_key is None:
                parameter_key = self.get_parameter_key(parameter_id)

            if morphology_key is None:
                morphology_key = self.get_morph_key(parameter_id=None, parameter_key=parameter_key,
                                                    morphology_id=morphology_id)

            modulation_key_list = self.meta_info[parameter_key][morphology_key]["neuromodulation"]

            if len(modulation_key_list) == 0:
                print("Modulation key list empty")
                print(f"Path: {self.neuron_path}, param: {parameter_key}, morph: {morphology_key}")
                import pdb
                pdb.set_trace()

            # modulation_key_list = list(self.modulation_info.keys())
            modulation_key = modulation_key_list[modulation_id % len(modulation_key_list)]
        else:
            modulation_key = None

        return modulation_key

    def get_input_parameters(self, parameter_id=None, morphology_id=None, parameter_key=None, morphology_key=None):

        if parameter_key is not None and morphology_key is not None:
            if self.meta_info and "input" in self.meta_info[parameter_key][morphology_key]:
                input_info = self.meta_info[parameter_key][morphology_key]["input"]
            else:
                input_info = dict()

        else:
            # Fallback if the user gave ID instead of Keys (will remove this at some point)
            par_key = self.get_parameter_key(parameter_id=parameter_id)
            morph_key = self.get_morph_key(parameter_id=parameter_id, morphology_id=morphology_id)

            if self.meta_info and "input" in self.meta_info[par_key][morph_key]:
                input_info = self.meta_info[par_key][morph_key]["input"]
            else:
                input_info = dict()

        return input_info

    def get_morphology(self, parameter_key=None, morphology_key=None, parameter_id=None, morphology_id=None):

        """
        Returns morphology for a given parameter_id, morphology_id combination.
        (Each parameter set has a set of morphologies that it is valid for)
        """

        # TODO: In the future remove parameter_id and morphology_id and only use keys
        if parameter_id is not None:
            par_key = self.get_parameter_key(parameter_id=parameter_id)
            if parameter_key:
                assert par_key == parameter_key, \
                    f"Mismatch. Parameter ID {parameter_id} has parameter_key {par_key}, " \
                    f"not {parameter_key} (Is snudda_data set correctly?)"
        elif parameter_key:
            par_key = parameter_key
        else:
            par_key = None

        if self.meta_info and par_key is not None:

            if morphology_id is not None:
                morph_key = self.get_morph_key(parameter_id=parameter_id, morphology_id=morphology_id,
                                               parameter_key=par_key)
                if morphology_key:
                    if morph_key != morphology_key:
                        raise KeyError(f"Mismatch: Expected morphology_key {morph_key}, got {morphology_key}")

            elif morphology_key is not None:
                morph_key = morphology_key
                if morphology_key not in self.meta_info[par_key]:
                    raise KeyError(f"Missing morphology_key {morphology_key}")
            else:
                raise KeyError(f"morphology_id or morphology_key must be specified if meta_info is used")

            morph_path = os.path.join(self.morphology_path, self.meta_info[par_key][morph_key]["morphology"])
            assert os.path.isfile(morph_path), f"Morphology file {morph_path} is missing (listed in {self.meta_path})"

        elif self.morphology_path and os.path.isfile(self.morphology_path):
            # Fallback if morphology file is specified in the path
            morph_path = self.morphology_path
            morph_key = morphology_key
        else:
            # No morphology
            file_list = glob.glob(os.path.join(self.neuron_path, "*swc"))
            morph_key = morphology_key

            if len(file_list) > 0:
                morph_path = file_list[0]

                assert len(file_list) == 1, f"More than one morphology available in {self.neuron_path}"

            else:
                print(f"No morphology in {self.neuron_path}")
                morph_path = None

        if morph_path is None:
            print(f"morph_path is None for {self.neuron_name} path: {self.neuron_path}. "
                  f"Is SNUDDA_DATA set correctly? ({os.environ['SNUDDA_DATA'] if 'SNUDDA_DATA' in os.environ else None})")
            import pdb
            pdb.set_trace()

        assert morph_path is not None, \
            (f"morph_path is None for {self.neuron_name} path: {self.neuron_path} (neuron_path should be a directory, not a SWC file). "
             f"Is SNUDDA_DATA set correctly? ({os.environ['SNUDDA_DATA']})")

        return morph_path, morph_key

    def get_parameters(self, parameter_key=None, parameter_id=None):
        """
        Returns neuron parameter set

        Args:
            parameter_key (str) : parameter key for the parameter set
            parameter_id (int) : parameter ID, which of the parameter sets to use

        Returns:
            One parameter set as a list of dictionaries
        """

        if self.parameter_info:
            if parameter_key is not None:
                assert parameter_key in self.parameter_info, \
                    f"Parameter key {parameter_key} not in parameter_info {self.parameter_info}. ERROR, no key: {parameter_key}"
                par_set = self.parameter_info[parameter_key]
            else:
                assert parameter_id is not None, "If parameter_key is not set, then parameter_id must be set"
                par_key_list = list(self.parameter_info.keys())
                par_key = par_key_list[parameter_id % len(par_key_list)]
                par_set = self.parameter_info[par_key]
        else:
            par_set = []

        return par_set

    def get_modulation_parameters(self, modulation_key=None, modulation_id=None):

        """
        Returns modulation parameter set, give either ID or key

        Args:
            modulation_key (str) : modulation key
            modulation_id (int) : modulation ID, which of the modulation parameter sets to use

        Returns:
            One parameter set as a list of dictionaries
        """

        if self.modulation_info:
            if modulation_key:
                modulation_set = self.modulation_info[modulation_key]
            elif modulation_id is not None:
                possible_keys = [x for x in self.modulation_info.keys()]
                modulation_set = self.modulation_info[possible_keys[modulation_id % len(self.modulation_info)]]
            else:
                assert False, f"You need to specify modulation_key or modulation_id to look up a modulation set"
        else:
            modulation_set = []

        return modulation_set

    def instantiate(self):
        """
        Instantiates all morphologies at once, instead of on demand.
        """

        n_par = len(self.parameter_info) if self.parameter_info is not None else 1

        for par_id in range(0, n_par):
            if self.verbose:
                print(f"Instantiates par_id = {par_id}")

            par_data = self.get_parameters(parameter_id=par_id)
            n_morph = self.get_num_morphologies(parameter_id=par_id)

            for morph_id in range(0, n_morph):
                morph_path, morph_key = self.get_morphology(parameter_id=par_id, morphology_id=morph_id)
                morph_tag = os.path.basename(morph_path)

                if morph_tag not in self.morphology_cache:
                    if self.verbose:
                        print(f"morph_tag = {morph_tag}")

                    self.morphology_cache[morph_tag] = NeuronMorphologyExtended(swc_filename=morph_path,
                                                                                param_data=self.parameter_path,
                                                                                mech_filename=self.mechanism_path,
                                                                                reaction_diffusion=self.reaction_diffusion_path,
                                                                                neuron_path=self.neuron_path,
                                                                                snudda_data=self.snudda_data,
                                                                                name=self.neuron_name,
                                                                                morphology_key=morph_key,
                                                                                load_morphology=self.load_morphology,
                                                                                virtual_neuron=self.virtual_neuron)

    def apply(self, function_name, arguments):
        """
        Applies function to all morphology prototypes
        """

        res = []

        for m in self.morphology_cache.values():
            function = getattr(m, function_name)
            res.append(function(*arguments))

        return res

    def all_have_axon(self):

        all_axon = True

        for m in self.morphology_cache.values():
            if 2 not in m.morphology_data["neuron"].sections \
              or len(m.morphology_data["neuron"].sections[2]) == 0:
                all_axon = False

        return all_axon

    def all_have_dend(self):

        all_dend = True

        for m in self.morphology_cache.values():
            if 3 not in m.morphology_data["neuron"].sections \
              or len(m.morphology_data["neuron"].sections[3]) == 0:
                all_dend = False

        return all_dend

    def clone(self, parameter_id=None, morphology_id=None, modulation_id=None,
              parameter_key=None, morphology_key=None, modulation_key=None,
              position=None, rotation=None, get_cache_original=False,
              morphology_path=None):
        """
        Creates a clone of the neuron prototype, with given position and rotation.

        Use either keys or id (keys have precedence if both given)

        Args:
            parameter_id (int) : parameter set to use
            morphology_id (int) : morphology set to use, a parameter set can have multiple morphology variations
            modulation_id (int) : neuro-modulation parameter set to use

            parameter_key (str) : parameter key for parameter set in parameters.json
            morphology_key (str) : morphology key for morphology in meta.json
            modulation_key (str) : modulation key for neuromodulation in modulation.json

            position (float,float,float) : position of neuron clone
            rotation : rotation (3x3 numpy matrix)
            get_cache_original (bool) : return the original rather than a clone

            morphology_path : If morphology path is given it can override morphology key
        """

        if morphology_path is not None:
            morph_path = morphology_path
        else:
            morph_path, morph_key = self.get_morphology(parameter_key=parameter_key, morphology_key=morphology_key,
                                                        parameter_id=parameter_id, morphology_id=morphology_id)

            if morphology_key is None:
                morphology_key = morph_key
            else:
                assert morphology_key == morph_key, \
                    f"Internal mismatch requested morphology_key {morphology_key}, got {morph_key}"

        morph_tag = os.path.basename(morph_path)

        if morph_tag not in self.morphology_cache:

            # TODO: hoc file will depend on both morphology_id and parameter_id, we ignore it for now
            self.morphology_cache[morph_tag] = NeuronMorphologyExtended(name=self.neuron_name,
                                                                        position=None,  # This is set further down when using clone
                                                                        rotation=None,
                                                                        swc_filename=morph_path,
                                                                        snudda_data=self.snudda_data,
                                                                        param_data=self.parameter_path,
                                                                        mech_filename=self.mechanism_path,
                                                                        reaction_diffusion=self.reaction_diffusion_path,
                                                                        neuron_path=self.neuron_path,
                                                                        parameter_key=parameter_key,
                                                                        morphology_key=morphology_key,
                                                                        modulation_key=modulation_key,
                                                                        load_morphology=self.load_morphology,
                                                                        virtual_neuron=self.virtual_neuron,
                                                                        verbose=self.verbose)


        if get_cache_original:
            assert position is None and rotation is None and modulation_id is None, \
                "If get_cache_original is passed position, rotation and modulation_id must be None"
            morph = self.morphology_cache[morph_tag]
        else:
            # Add the keys also for parameter, morphology and modulation
            if parameter_key is None:
                parameter_key = self.get_parameter_key(parameter_id=parameter_id)

            if morphology_key is None:
                morphology_key = self.get_morph_key(parameter_key=parameter_key, morphology_id=morphology_id)

            if modulation_key is None:
                modulation_key = self.get_modulation_key(parameter_key=parameter_key,
                                                         morphology_key=morphology_key,
                                                         modulation_id=modulation_id)

            morph = self.morphology_cache[morph_tag].clone(position=position,
                                                           rotation=rotation,
                                                           parameter_key=parameter_key,
                                                           morphology_key=morphology_key,
                                                           modulation_key=modulation_key)

            if self.meta_info is not None and "axon_density" in self.meta_info.get(parameter_key, {}).get(morphology_key, {}):
                axon_density_type, axon_density, axon_density_bounds = self.meta_info[parameter_key][morphology_key]["axon_density"]

                if axon_density_type == "r":
                    morph.set_axon_voxel_radial_density(density=axon_density,
                                                        max_axon_radius=axon_density_bounds)
                elif axon_density_type == "xyz":
                    morph.set_axon_voxel_xyz_density(density=axon_density,
                                                     axon_density_bounds_xyz=axon_density_bounds)
                else:
                    raise ValueError(f"Unknown axon_density type {axon_density_type}, {self.neuron_path =}")

        return morph
