import os
import glob
import pandas as pd
from bmtk.utils.reports import CompartmentReport
from bmtk.simulator.pointnet.io_tools import io
from bmtk.simulator.pointnet.nest_utils import nest_version
from bmtk.simulator.pointnet.modules.sim_module import SimulatorMod

import nest


try:
    MPI_RANK = nest.Rank()
    N_HOSTS = nest.NumProcesses()

except Exception as e:
    MPI_RANK = 0
    N_HOSTS = 1


def create_multimeter_nest2(tstart, tstop, variable_name, label):
    return nest.Create(
        'multimeter',
        params={
            'start': tstart,
            'stop': tstop,
            'to_file': True,
            'to_memory': False,
            'withtime': True,
            'record_from': variable_name,
            'label': label
        }
    )


def create_multimeter_nest3(tstart, tstop, variable_name, label):
    return nest.Create(
        'multimeter',
        params={
            'start': tstart,
            'stop': tstop,
            'record_to': 'ascii',
            'record_from': variable_name,
            'label': label
        }
    )


def read_dat_nest2(dat_file, variable_name):
    return pd.read_csv(dat_file, index_col=False, names=['nest_id', 'time']+variable_name, sep='\t')


def read_dat_nest3(dat_file, variable_name):
    report_df = pd.read_csv(dat_file, index_col=False, sep='\t', comment='#')
    report_df = report_df.rename(columns={'sender': 'nest_id', 'time_ms': 'time'})
    return report_df



if nest_version[0] >= 3:
    create_multimeter = create_multimeter_nest3
    read_dat = read_dat_nest3
else:
    create_multimeter = create_multimeter_nest2
    read_dat = read_dat_nest2


class MultimeterMod(SimulatorMod):
    def __init__(self, tmp_dir, file_name, variable_name, cells, tstart=None, tstop=None, interval=None, to_h5=True,
                 delete_dat=True, **opt_params):
        """For recording neuron properties using a NEST multimeter object

        :param tmp_dir: ouput directory
        :param file_name: Name of (SONATA hdf5) file that will be saved to
        :param variable_name: A list of the variable(s) being recorded. Must be valid according to the cells
        :param cells: A node-set or list of gids to record from
        :param tstart: Start time of the recording (if None will default to sim.tstart)
        :param tstop: Stop time of recording (if None will default to sim.tstop)
        :param interval: Recording time step (if None will default to sim.dt)
        :param to_h5: True to save to sonata .h5 format (default: True)
        :param delete_dat: True to delete the .dat files created by NEST (default True)
        :param opt_params:
        """
        self._output_dir = tmp_dir
        self._file_name = file_name if os.path.isabs(file_name) else os.path.join(self._output_dir, file_name)
        self._variable_name = variable_name
        self._node_set = cells
        self._tstart = tstart
        self._tstop = tstop
        self._interval = interval
        self._to_h5 = to_h5
        self._delete_dat = delete_dat

        self._gids = None  # global ids will be the NEST ids assigned to each cell
        self._multimeter = None
        self._population = None

        self._min_delay = 1.0  # Required for calculating steps recorded

        self.__output_label = os.path.join(self._output_dir, '__bmtk_nest_{}'.format(os.path.basename(self._file_name)))
        self._var_recorder = None  # CellVarRecorder(self._file_name, self._output_dir, self._variable_name, buffer_data=False)

    def initialize(self, sim):
        node_set = sim.net.get_node_set(self._node_set)

        self._gids = list(set(node_set.gids()))
        self._gids.sort()
        self._population = node_set.population_names()[0]
        self._tstart = self._tstart or sim.tstart
        self._tstop = self._tstop or sim.tstop
        self._interval = self._interval or sim.dt
        self._multimeter = create_multimeter(self._tstart, self._tstop, self._variable_name, self.__output_label)

        nest.SetStatus(self._multimeter, 'interval', self._interval)
        nest.Connect(self._multimeter, self._gids)

    def finalize(self, sim):
        io.barrier()  # Makes sure all nodes finish, but not sure if actually required by nest

        # min_delay needs to be fetched after simulation otherwise the value will be off. There also seems to be some
        # MPI barrier inside GetKernelStatus
        self._min_delay = nest.GetKernelStatus('min_delay')
        if self._to_h5 and MPI_RANK == 0:
            # Initialize hdf5 file including preallocated data block of recorded variables
            #   Unfortantely with NEST the final time-step recorded can't be calculated in advanced, and even with the
            # same min/max_delay can be different. We need to read the output-file to get n_steps
            def get_var_recorder(node_recording_df):
                if self._var_recorder is None:
                    self._var_recorder = CompartmentReport(self._file_name, mode='w', variable=self._variable_name[0],
                                                           default_population=self._population,
                                                           tstart=node_recording_df['time'].min(),
                                                           tstop=node_recording_df['time'].max(),
                                                           dt=self._interval,
                                                           n_steps=len(node_recording_df),
                                                           mpi_size=1)
                    if self._to_h5 and MPI_RANK == 0:
                        for gid in self._gids:
                            pop_id = gid_map.get_pool_id(gid)
                            self._var_recorder.add_cell(pop_id.node_id, element_ids=[0], element_pos=[0.0],
                                                        population=pop_id.population)

                    self._var_recorder.initialize()

                return self._var_recorder

            gid_map = sim.net.gid_map
            for nest_file in glob.glob('{}*'.format(self.__output_label)):
                # report_df = pd.read_csv(nest_file, index_col=False, names=['nest_id', 'time']+self._variable_name,
                #                         sep='\t', comment='#')
                report_df = read_dat(nest_file, self._variable_name)
                # print(report_df)
                # exit()

                for grp_id, grp_df in report_df.groupby(by='nest_id'):
                    pop_id = gid_map.get_pool_id(grp_id)
                    vr = get_var_recorder(grp_df)
                    for var_name in self._variable_name:
                        vr.record_cell_block(node_id=pop_id.node_id, vals=grp_df[var_name], beg_step=0,
                                             end_step=vr[pop_id.population].n_steps(), population=pop_id.population)

                if self._delete_dat:
                    # remove csv file created by nest
                    os.remove(nest_file)

            self._var_recorder.close()

        io.barrier()
