"""Batch Simulation."""

# Authors: Abdul Samad Siddiqui <abdulsamadsid1@gmail.com>
#          Nick Tolley <nicholas_tolley@brown.edu>
#          Ryan Thorpe <ryan_thorpe@brown.edu>
#          Mainak Jas <mjas@mgh.harvard.edu>

import numpy as np
import os
from joblib import Parallel, delayed, parallel_config

from .network import Network
from .externals.mne import _validate_type, _check_option
from .dipole import simulate_dipole
from .network_models import jones_2009_model


class BatchSimulate(object):
    def __init__(self, set_params, net=jones_2009_model(), tstop=170,
                 dt=0.025, n_trials=1, record_vsec=False,
                 record_isec=False, postproc=False, save_outputs=False,
                 save_folder='./sim_results', batch_size=100,
                 overwrite=True, summary_func=None,
                 save_dpl=True, save_spiking=False,
                 save_lfp=False, save_voltages=False,
                 save_currents=False, save_calcium=False,
                 clear_cache=False):
        """Initialize the BatchSimulate class.

        Parameters
        ----------
        set_params : func
            User-defined function that sets parameters in network drives.

                ``set_params(net, params) -> None``

            where ``net`` is a Network object and ``params`` is a dictionary
            of the parameters that will be set inside the function.
        net : Network object, optional
            The network model to use for simulations. Must be an instance of
            jones_2009_model, law_2021_model, or calcium_model.
            Default is jones_2009_model().
        tstop : float, optional
            The stop time for the simulation. Default is 170 ms.
        dt : float, optional
            The time step for the simulation. Default is 0.025 ms.
        n_trials : int, optional
            The number of trials for the simulation. Default is 1.
        record_vsec : 'all' | 'soma' | False
            Option to record voltages from all sections ('all'), or just
            the soma ('soma'). Default: False.
        record_isec : 'all' | 'soma' | False
            Option to record voltages from all sections ('all'), or just
            the soma ('soma'). Default: False.
        postproc : bool
            If True, smoothing (``dipole_smooth_win``) and scaling
            (``dipole_scalefctr``) values are read from the parameter file, and
            applied to the dipole objects before returning.
            Note that this setting
            only affects the dipole waveforms, and not somatic voltages,
            possible extracellular recordings etc.
            The preferred way is to use the
            :meth:`~hnn_core.dipole.Dipole.smooth` and
            :meth:`~hnn_core.dipole.Dipole.scale` methods instead.
            Default: False.
        save_outputs : bool, optional
            Whether to save the simulation outputs to files. Default is False.
        save_folder : str, optional
            The path to save the simulation outputs.
            Default is './sim_results'.
        batch_size : int, optional
            The maximum number of simulations saved in a single file.
            Default is 100.
        overwrite : bool, optional
            Whether to overwrite existing files and create file paths
            if they do not exist. Default is True.
        summary_func : callable, optional
            A function to calculate summary statistics from the simulation
            results. Default is None.
        save_dpl : bool
            If True, save dipole results. Note, `save_outputs` must be True.
            Default: True.
        save_spiking : bool
            If True, save spiking results. Note, `save_outputs` must be True.
            Default: False.
        save_lfp : bool
            If True, save local field potential (lfp) results.
            Note, `save_outputs` must be True.
            Default: False.
        save_voltages : bool
            If True, save voltages results. Note, `save_outputs` must be True.
            Default: False.
        save_currents : bool
            If True, save currents results. Note, `save_outputs` must be True.
            Default: False.
        save_calcium : bool
            If True, save calcium concentrations.
            Note, `save_outputs` must be True.
            Default: False.
        clear_cache : bool, optional
            Whether to clear the results cache after saving each batch.
            Default is False.
        Notes
        -----
        When `save_output=True`, the saved files will appear as
        `sim_run_{start_idx}-{end_idx}.npz` in the specified `save_folder`
        directory. The `start_idx` and `end_idx` indicate the range of
        simulation indices contained in each file. Each file will contain
        a maximum of `batch_size` simulations, split evenly among the
        available files. If `overwrite=True`, existing files with the same name
        will be overwritten.
        """

        _validate_type(net, Network, 'net', 'Network')
        _validate_type(tstop, types='numeric', item_name='tstop')
        _validate_type(dt, types='numeric', item_name='dt')
        _validate_type(n_trials, types='int', item_name='n_trials')
        _check_option('record_vsec', record_vsec, ['all', 'soma', False])
        _check_option('record_isec', record_isec, ['all', 'soma', False])
        _validate_type(save_folder, types='path-like', item_name='save_folder')
        _validate_type(batch_size, types='int', item_name='batch_size')
        _validate_type(save_outputs, types=(bool,), item_name='save_outputs')
        _validate_type(save_dpl, types=(bool,), item_name='save_dpl')
        _validate_type(save_spiking, types=(bool,), item_name='save_spiking')
        _validate_type(save_lfp, types=(bool,), item_name='save_lfp')
        _validate_type(save_voltages, types=(bool,), item_name='save_voltages')
        _validate_type(save_currents, types=(bool,), item_name='save_currents')
        _validate_type(save_calcium, types=(bool,), item_name='save_calcium')
        _validate_type(clear_cache, types=(bool,), item_name='clear_cache')

        if set_params is not None and not callable(set_params):
            raise TypeError("set_params must be a callable function")

        if summary_func is not None and not callable(summary_func):
            raise TypeError("summary_func must be a callable function")

        self.net = net
        self.set_params = set_params
        self.tstop = tstop
        self.dt = dt
        self.n_trials = n_trials
        self.record_vsec = record_vsec
        self.record_isec = record_isec
        self.postproc = postproc
        self.save_outputs = save_outputs
        self.save_folder = save_folder
        self.batch_size = batch_size
        self.overwrite = overwrite
        self.summary_func = summary_func
        self.save_dpl = save_dpl
        self.save_spiking = save_spiking
        self.save_lfp = save_lfp
        self.save_voltages = save_voltages
        self.save_currents = save_currents
        self.save_calcium = save_calcium
        self.clear_cache = clear_cache

    def run(self, param_grid, return_output=True,
            combinations=True, n_jobs=1, backend='loky',
            verbose=50):
        """Run batch simulations.

        Parameters
        ----------
        param_grid : dict
            Dictionary with parameter names and ranges.
        return_output : bool, optional
            Whether to return the simulation outputs. Default is True.
        combinations : bool, optional
            Whether to generate the Cartesian product of the parameter ranges.
            If False, generate combinations based on corresponding indices.
            Default is True.
        n_jobs : int, optional
            Number of parallel jobs. Default is 1.
        backend : str or joblib.parallel.ParallelBackendBase instance, optional
            The parallel backend to use. Can be one of `loky`, `threading`,
            `multiprocessing`, or `dask`. Default is `loky`
        verbose : int, optional
            The verbosity level for parallel execution. Default is 50.

        Returns
        -------
        results : dict
            Dictionary containing 'summary_statistics' and optionally
            'simulated_data'.
            'simulated_data' may include keys: 'dpl', 'lfp', 'spikes',
            'voltages', 'param_values', 'net', 'times'.

        Notes
        -----
        Return content depends on summary_func, return_output, and
        clear_cache settings.

        """
        _validate_type(param_grid, types=(dict,), item_name='param_grid')
        _validate_type(n_jobs, types='int', item_name='n_jobs')
        _check_option('backend', backend, ['loky', 'threading',
                                           'multiprocessing', 'dask'])
        _validate_type(verbose, types='int', item_name='verbose')

        param_combinations = self._generate_param_combinations(
            param_grid, combinations)
        total_sims = len(param_combinations)
        num_sims_per_batch = max(total_sims // self.batch_size, 1)
        batch_size = min(self.batch_size, total_sims)

        results = []
        simulated_data = []
        for i in range(batch_size):
            start_idx = i * num_sims_per_batch
            end_idx = start_idx + num_sims_per_batch
            if i == batch_size - 1:
                end_idx = len(param_combinations)
            batch_results = self.simulate_batch(
                param_combinations[start_idx:end_idx],
                n_jobs=n_jobs,
                backend=backend,
                verbose=verbose)

            if self.save_outputs:
                self._save(batch_results, start_idx, end_idx)

            if self.summary_func is not None:
                summary_statistics = self.summary_func(batch_results)
                results.append(summary_statistics)
                if not self.clear_cache:
                    simulated_data.append(batch_results)

            elif return_output and not self.clear_cache:
                simulated_data.append(batch_results)

            if self.clear_cache:
                del batch_results

        if return_output:
            if self.clear_cache:
                return {"summary_statistics": results}
            else:
                return {"summary_statistics": results,
                        "simulated_data": simulated_data}

    def simulate_batch(self, param_combinations, n_jobs=1,
                       backend='loky', verbose=50):
        """Simulate a batch of parameter sets in parallel.

        Parameters
        ----------
        param_combinations : list
            List of parameter combinations.
        n_jobs : int, optional
            Number of parallel jobs. Default is 1.
        backend : str or joblib.parallel.ParallelBackendBase instance, optional
            The parallel backend to use. Can be one of `loky`, `threading`,
            `multiprocessing`, or `dask`. Default is `loky`
        verbose : int, optional
            The verbosity level for parallel execution. Default is 50.

        Returns
        -------
        res: list
            List of dictionaries containing simulation results.
            Each dictionary contains the following keys:
            - `net`: The network model used for the simulation.
            - `dpl`: The simulated dipole.
            - `param_values`: The parameter values used for the simulation.
        """
        _validate_type(param_combinations, types=(list,),
                       item_name='param_combinations')
        _validate_type(n_jobs, types='int', item_name='n_jobs')
        _check_option('backend', backend, ['loky', 'threading',
                                           'multiprocessing', 'dask'])
        _validate_type(verbose, types='int', item_name='verbose')

        with parallel_config(backend=backend):
            res = Parallel(n_jobs=n_jobs, verbose=verbose)(
                delayed(self._run_single_sim)(
                    params) for params in param_combinations)
        return res

    def _run_single_sim(self, param_values):
        """Run a single simulation.

        Parameters
        ----------
        param_values : dict
            Dictionary of parameter values.

        Returns
        -------
        dict
            Dictionary containing the simulation results.
            The dictionary contains the following keys:
            - `net`: The network model used for the simulation.
            - `dpl`: The simulated dipole.
            - `param_values`: The parameter values used for the simulation.
        """

        net = self.net.copy()
        self.set_params(param_values, net)

        results = {'net': net, 'param_values': param_values}

        if self.save_dpl:
            dpl = simulate_dipole(net,
                                  tstop=self.tstop,
                                  dt=self.dt,
                                  n_trials=self.n_trials,
                                  record_vsec=self.record_vsec,
                                  record_isec=self.record_isec,
                                  postproc=self.postproc)
            results['dpl'] = dpl

        if self.save_spiking:
            results['spiking'] = {
                'spike_times': net.cell_response.spike_times,
                'spike_types': net.cell_response.spike_types,
                'spike_gids': net.cell_response.spike_gids
            }

        if self.save_lfp:
            results['lfp'] = net.rec_arrays

        if self.save_voltages:
            results['voltages'] = net.cell_response.vsec

        if self.save_currents:
            results['currents'] = net.cell_response.isec

        if self.save_calcium:
            results['calcium'] = net.cell_response.ca

        return results

    def _generate_param_combinations(self, param_grid, combinations=True):
        """Generate combinations of parameters from the grid.

        Parameters
        ----------
        param_grid : dict
            Dictionary with parameter names and ranges.
        combinations : bool, optional
            Whether to generate the Cartesian product of the parameter ranges.
            If False, generate combinations based on corresponding indices.
            Default is True.

        Returns
        -------
        param_combinations: list
            List of parameter combinations.
        """
        from itertools import product

        keys, values = zip(*param_grid.items())
        if combinations:
            param_combinations = [dict(zip(keys, combination))
                                  for combination in product(*values)]
        else:
            param_combinations = [dict(zip(keys, combination))
                                  for combination in zip(*values)]
        return param_combinations

    def _save(self, results, start_idx, end_idx):
        """Save simulation results to files.

        Parameters
        ----------
        results : list
            The results to save.
        start_idx : int
            The index of the first simulation in this batch.
        end_idx : int
            The index of the last simulation in this batch.
        """
        _validate_type(results, types=(list,), item_name='results')
        _validate_type(start_idx, types='int', item_name='start_idx')
        _validate_type(end_idx, types='int', item_name='end_idx')

        if not os.path.exists(self.save_folder):
            os.makedirs(self.save_folder)

        save_data = {
            'param_values': [result['param_values'] for result in results]
        }

        attributes_to_save = ['dpl', 'spiking', 'lfp',
                              'voltages', 'currents', 'calcium']
        for attr in attributes_to_save:
            if getattr(self, f'save_{attr}') and attr in results[0]:
                save_data[attr] = [result[attr] for result in results]

        file_name = os.path.join(self.save_folder,
                                 f'sim_run_{start_idx}-{end_idx}.npz')
        if os.path.exists(file_name) and not self.overwrite:
            raise FileExistsError(
                f"File {file_name} already exists and "
                "overwrite is set to False.")

        np.savez(file_name, **save_data)

    def load_results(self, file_path, return_data=None):
        """Load simulation results from a file.

        Parameters
        ----------
        file_path : str
            The path to the file containing the simulation results.
        return_data : list of str, optional
            List of data types to return. If None,
            defaults to the types specified during initialization.

        Returns
        -------
        results : dict
            Dictionary containing the loaded results and parameter values.
        """
        if return_data is None:
            return_data = []
            if self.save_dpl:
                return_data.append('dpl')
            if self.save_spiking:
                return_data.append('spiking')
            if self.save_lfp:
                return_data.append('lfp')
            if self.save_voltages:
                return_data.append('voltages')
            if self.save_currents:
                return_data.append('currents')
            if self.save_calcium:
                return_data.append('calcium')

        data = np.load(file_path, allow_pickle=True)
        results = {key: data[key].tolist() for key in data.files
                   if key in return_data or key == 'param_values'}
        return results

    def load_all_results(self):
        """Load all simulation results from the save folder.

        Returns
        -------
        all_results : list
            List of dictionaries containing all loaded simulation results.
        """
        all_results = []
        for file_name in os.listdir(self.save_folder):
            if file_name.startswith('sim_run_') and file_name.endswith('.npz'):
                file_path = os.path.join(self.save_folder, file_name)
                results = self.load_results(file_path)
                all_results.append(results)
        return all_results
