# Copyright 2020. Allen Institute. All rights reserved
#
# Redistribution and use in source and binary forms, with or without modification, are permitted provided that the
# following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following
# disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following
# disclaimer in the documentation and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote
# products derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
# INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
import csv
import pandas as pd
import numpy as np
import h5py
import six
from collections import defaultdict
import warnings

from .spike_trains_api import SpikeTrainsReadOnlyAPI
from .core import SortOrder, csv_headers, col_population, col_timestamps, col_node_ids, pop_na


GRP_spikes_root = 'spikes'
DATASET_timestamps = 'timestamps'
DATASET_node_ids = 'node_ids'


sorting_attrs = {
    'time': SortOrder.by_time,
    'by_time': SortOrder.by_time,
    'id': SortOrder.by_id,
    'by_id': SortOrder.by_id,
    'none': SortOrder.none,
    'unknown': SortOrder.unknown
}


def load_sonata_file(path, version=None, **kwargs):
    """Loads a Sonata file reader, making sure it matches the correct version.

    :param path:
    :param version:
    :param kwargs:
    :return:
    """
    try:
        with h5py.File(path, 'r') as h5:
            spikes_root = h5[GRP_spikes_root]
            for name, h5_obj in spikes_root.items():
                if isinstance(h5_obj, h5py.Group):
                    # In case there exists a population subgroup
                    return SonataSTReader(path, **kwargs)
    except Exception:
        pass

    try:
        with h5py.File(path, 'r') as h5:
            spikes_root = h5[GRP_spikes_root]
            if 'gids' in spikes_root and 'timestamps' in spikes_root:
                return SonataOldReader(path, **kwargs)
    except Exception:
        pass

    try:
        with h5py.File(path, 'r') as h5:
            if '/spikes' in h5:
                return EmptySonataReader(path, **kwargs)
    except Exception:
        pass

    raise Exception('Could not open file {}, does not contain SONATA spike-trains'.format(path))


def to_list(v):
    if v is not None and np.isscalar(v):
        return [v]
    else:
        return v


class SonataSTReader(SpikeTrainsReadOnlyAPI):
    def __init__(self, path, **kwargs):
        self._path = path
        self._h5_handle = h5py.File(self._path, 'r')
        self._DATASET_node_ids = 'node_ids'
        self._n_spikes = None
        # TODO: Create a function for looking up population and can return errors if more than one
        self._default_pop = None
        # self._node_list = None

        self._indexed = False
        self._index_nids = {}

        if GRP_spikes_root not in self._h5_handle:
            raise Exception('Could not find /{} root'.format(GRP_spikes_root))
        else:
            self._spikes_root = self._h5_handle[GRP_spikes_root]

        if 'population' in kwargs:
            pop_filter = to_list(kwargs['population'])
        elif 'populations' in kwargs:
            pop_filter = to_list(kwargs['populations'])
        else:
            pop_filter = None

        # get a map of 'pop_name' -> pop_group
        self._population_map = {}
        for name, h5_obj in self._h5_handle[GRP_spikes_root].items():
            if isinstance(h5_obj, h5py.Group):
                if pop_filter is not None and name not in pop_filter:
                    continue

                if 'node_ids' not in h5_obj or 'timestamps' not in h5_obj:
                    warnings.warn('population {} in {} is missing spikes, skipping.'.format(name, path))
                    continue

                self._population_map[name] = h5_obj

        if not self._population_map:
            # In old version of the sonata standard there was no 'population' subgroup. For backwards compatability
            # use a default dictionary
            # TODO: Remove so we only have to support latest version of SONATA
            self._population_map = defaultdict(lambda: self._h5_handle[GRP_spikes_root])
            self._population_map[pop_na] = self._h5_handle[GRP_spikes_root]
            self._DATASET_node_ids = 'gids'

        self._default_pop = kwargs.get('default_population', list(self._population_map.keys())[0])

        self._population_sorting_map = {}
        for pop_name, pop_grp in self._population_map.items():
            if 'sorting' in pop_grp[self._DATASET_node_ids].attrs.keys():
                # Found a few existing sonata files put the 'sorting' attribute in the node_ids dataset, remove later
                attr_str = pop_grp[self._DATASET_node_ids].attrs['sorting']
                sort_order = sorting_attrs.get(attr_str, SortOrder.unknown)

            elif 'sorting' in pop_grp.attrs.keys():
                attr_str = pop_grp.attrs['sorting']
                sort_order = sorting_attrs.get(attr_str, SortOrder.unknown)

            else:
                sort_order = SortOrder.unknown

            self._population_sorting_map[pop_name] = sort_order

        # TODO: Add option to skip building indices
        self._build_node_index()

        # units are not instrinsic to a csv file, but allow users to pass it in if they know
        self._units_maps = {}
        for pop_name, pop_grp in self._population_map.items():
            if 'units' in pop_grp['timestamps'].attrs:
                pop_units = pop_grp['timestamps'].attrs['units']
            elif 'units' in kwargs:
                pop_units = kwargs['units']
            else:
                pop_units = 'ms'

            self._units_maps[pop_name] = pop_units

    def _build_node_index(self):
        self._indexed = False
        for pop_name, pop_grp in self._population_map.items():
            sort_order = self._population_sorting_map[pop_name]
            nodes_indices = {}
            # loop on h5 is slow, so convert it to np before the loop.
            node_ids_ds = np.array(pop_grp[self._DATASET_node_ids])
            if sort_order == SortOrder.by_id:
                indx_beg = 0
                last_id = node_ids_ds[0]
                for indx, cur_id in enumerate(node_ids_ds):
                    if cur_id != last_id:
                        # nodes_indices[last_id] = np.arange(indx_beg, indx)
                        nodes_indices[last_id] = slice(indx_beg, indx)
                        last_id = cur_id
                        indx_beg = indx
                # nodes_indices[last_id] = np.arange(indx_beg, indx + 1)
                nodes_indices[last_id] = slice(indx_beg, indx + 1)  # capture the last node_id
            else:
                nodes_indices = {int(node_id): [] for node_id in np.unique(node_ids_ds)}
                for indx, node_id in enumerate(node_ids_ds):
                    nodes_indices[node_id].append(indx)

            self._index_nids[pop_name] = nodes_indices

        self._indexed = True

    @property
    def populations(self):
        return list(self._population_map.keys())

    def units(self, population=None):
        population = population if population is not None else self._default_pop
        return self._units_maps[population]

    def set_units(self, u, population=None):
        self._units_maps[population] = u

    def sort_order(self, population=None):
        return self._population_sorting_map[population]

    def node_ids(self, population=None):
        population = population if population is not None else self._default_pop
        pop_grp = self._population_map[population]
        return np.unique(pop_grp[self._DATASET_node_ids][()])

    def n_spikes(self, population=None):
        population = population if population is not None else self._default_pop
        return len(self._population_map[population][DATASET_timestamps])

    def time_range(self, populations=None):
        if populations is None:
            populations = [self._default_pop]

        if isinstance(populations, six.string_types) or np.isscalar(populations):
            populations = [populations]

        min_time = np.inf
        max_time = -np.inf
        for pop_name, pop_grp in self._population_map.items():
            if pop_name in populations:
                # TODO: Check if sorted by time
                for ts in pop_grp[DATASET_timestamps]:
                    if ts < min_time:
                        min_time = ts

                    if ts > max_time:
                        max_time = ts

        return min_time, max_time

    def to_dataframe(self, populations=None, sort_order=SortOrder.none, with_population_col=True, **kwargs):
        populations = populations if populations is not None else self.populations
        if isinstance(populations, six.string_types) or np.isscalar(populations):
            populations = [populations]

        ret_df = None
        for pop_name, pop_grp in self._population_map.items():
            if pop_name in populations:
                pop_df = pd.DataFrame({
                    col_timestamps: pop_grp[DATASET_timestamps],
                    # col_population: pop_name,
                    col_node_ids: pop_grp[self._DATASET_node_ids]
                })

                if with_population_col:
                    pop_df['population'] = pop_name

                if sort_order == SortOrder.by_id:
                    pop_df = pop_df.sort_values('node_ids')
                elif sort_order == SortOrder.by_time:
                    pop_df = pop_df.sort_values('timestamps')

                ret_df = pop_df if ret_df is None else pd.concat((ret_df, pop_df))

        if sort_order == SortOrder.by_time:
            ret_df.sort_values(by=col_timestamps, inplace=True)
        elif sort_order == SortOrder.by_id:
            ret_df.sort_values(by=col_node_ids, inplace=True)

        return ret_df

    def get_times(self, node_id, population=None, time_window=None, **kwargs):
        if population is None:
            if not isinstance(self._default_pop, six.string_types) and len(self._default_pop) > 1:
                raise Exception('Error: Multiple populations, must select one.')

            population = self._default_pop

        elif population not in self._population_map:
            return []

        spikes_index = self._index_nids[population].get(node_id, None)
        if spikes_index is None:
            return []
        spike_times = self._population_map[population][DATASET_timestamps][spikes_index]

        if time_window is not None:
            spike_times = spike_times[(time_window[0] <= spike_times) & (spike_times <= time_window[1])]

        return spike_times

    def spikes(self, node_ids=None, populations=None, time_window=None, sort_order=SortOrder.none, **kwargs):
        populations = populations or self.populations
        if np.isscalar(populations):
            populations = [populations]

        if sort_order == SortOrder.by_id:
            for pop_name in populations:
                if pop_name not in self.populations:
                    continue

                timestamps_ds = self._population_map[pop_name][DATASET_timestamps]
                index_map = self._index_nids[pop_name]
                node_ids = list(index_map.keys())
                node_ids.sort()
                for node_id in node_ids:
                    st_indices = index_map[node_id]
                    for st in timestamps_ds[st_indices]:  # st_indices:
                        yield st, pop_name, node_id

        elif sort_order == SortOrder.by_time:
            # TODO: Reimplement using a heap
            index_ranges = []
            for pop_name in populations:
                if pop_name not in self.populations:
                    continue

                pop_grp = self._population_map[pop_name]
                if self._population_sorting_map[pop_name] == SortOrder.by_time:
                    ts_ds = pop_grp[DATASET_timestamps]
                    index_ranges.append([pop_name, 0, len(ts_ds), np.arange(len(ts_ds)), pop_grp, ts_ds[0]])
                else:
                    ts_ds = pop_grp[DATASET_timestamps]
                    ts_indices = np.argsort(ts_ds[()])
                    index_ranges.append([pop_name, 0, len(ts_ds), ts_indices, pop_grp, ts_ds[ts_indices[0]]])

            while index_ranges:
                selected_r = index_ranges[0]
                for i, r in enumerate(index_ranges[1:]):
                    if r[5] < selected_r[5]:
                        selected_r = r

                ds_index = selected_r[1]
                timestamp = selected_r[5]  # pop_grp[DATASET_timestamps][ds_index]
                node_id = selected_r[4][self._DATASET_node_ids][ds_index]
                pop_name = selected_r[0]
                ds_index += 1
                if ds_index >= selected_r[2]:
                    index_ranges.remove(selected_r)
                else:
                    selected_r[1] = ds_index
                    ts_index = selected_r[3][ds_index]
                    next_ts = self._population_map[pop_name][DATASET_timestamps][selected_r[3][ds_index]]
                    selected_r[5] = next_ts  # pop_grp[DATASET_timestamps][selected_r[3][ds_index]]

                yield timestamp, pop_name, node_id

        else:
            for pop_name in populations:
                if pop_name not in self.populations:
                    continue

                pop_grp = self._population_map[pop_name]
                for i in range(len(pop_grp[DATASET_timestamps])):
                    yield pop_grp[DATASET_timestamps][i], pop_name, pop_grp[self._DATASET_node_ids][i]

    def __len__(self):
        if self._n_spikes is None:
            self._n_spikes = 0
            for _, pop_grp in self._population_map.items():
                self._n_spikes += len(pop_grp[self._DATASET_node_ids])

        return self._n_spikes


class SonataOldReader(SonataSTReader):
    """Older version of SONATA
    """

    def node_ids(self, population=None):
        return super(SonataOldReader, self).node_ids(population=None)

    def n_spikes(self, population=None):
        return super(SonataOldReader, self).n_spikes(population=self._default_pop)

    def time_range(self, populations=None):
        return super(SonataOldReader, self).time_range(populations=None)

    def to_dataframe(self, node_ids=None, populations=None, time_window=None, sort_order=SortOrder.none, **kwargs):
        return super(SonataOldReader, self).to_dataframe(node_ids=node_ids, populations=None,
                                                         time_window=time_window, sort_order=sort_order, **kwargs)

    def get_times(self, node_id, population=None, time_window=None, **kwargs):
        return super(SonataOldReader, self).get_times(node_id=node_id, population=None,
                                                      time_window=time_window, **kwargs)

    def spikes(self, node_ids=None, populations=None, time_window=None, sort_order=SortOrder.none, **kwargs):
        return super(SonataOldReader, self).spikes(node_ids=node_ids, populations=None,
                                                   time_window=time_window, sort_order=sort_order, **kwargs)


class EmptySonataReader(SpikeTrainsReadOnlyAPI):
    """A Hack that is needed for when a simulation produces a file with no spikes, since there won't/can't be
    <population_name> subgroup and/or gids/timestamps datasets.

    """
    def __init__(self, path, **kwargs):
        pass

    @property
    def populations(self):
        return []

    def node_ids(self, population=None):
        return []

    def n_spikes(self, population=None):
        return 0

    def time_range(self, populations=None):
        return None

    def to_dataframe(self, node_ids=None, populations=None, time_window=None, sort_order=SortOrder.none, **kwargs):
        return pd.DataFrame(columns=csv_headers)

    def get_times(self, node_id, population=None, time_window=None, **kwargs):
        return []

    def spikes(self, node_ids=None, populations=None, time_window=None, sort_order=SortOrder.none, **kwargs):
        return []


class CSVSTReader(SpikeTrainsReadOnlyAPI):
    def __init__(self, path, sep=' ', default_population=None, **kwargs):
        self._n_spikes = None
        self._populations = None

        try:
            # check to see if file contains headers
            with open(path, 'r') as csvfile:
                sniffer = csv.Sniffer()
                has_headers = sniffer.has_header(csvfile.read(1024))
        except Exception:
            has_headers = True

        self._spikes_df = pd.read_csv(path, sep=sep, header=0 if has_headers else None)

        if not has_headers:
            self._spikes_df.columns = csv_headers[0::2]

        self._defaul_population = default_population if default_population is not None \
            else self._spikes_df[col_population][0]

        if col_population not in self._spikes_df.columns:
            pop_name = kwargs.get(col_population, self._defaul_population)
            self._spikes_df[col_population] = pop_name

        # TODO: Check all the necessary columns exits
        self._spikes_df = self._spikes_df[csv_headers]

    @property
    def populations(self):
        if self._populations is None:
            self._populations = self._spikes_df['population'].unique()

        return self._populations

    def to_dataframe(self, populations=None, sort_order=SortOrder.none, with_population_col=True, **kwargs):
        selected = self._spikes_df.copy()

        mask = True
        if populations is not None:
            if isinstance(populations, six.string_types) or np.isscalar(populations):
                mask &= selected[col_population] == populations
            else:
                mask &= selected[col_population].isin(populations)

        if isinstance(mask, pd.Series):
            selected = selected[mask]

        if sort_order == SortOrder.by_time:
            selected.sort_values(by=col_timestamps, inplace=True)
        elif sort_order == SortOrder.by_id:
            selected.sort_values(by=col_node_ids, inplace=True)

        if not with_population_col:
            selected = selected.drop(col_population, axis=1)

        selected.index = pd.RangeIndex(len(selected.index))
        return selected

    def get_times(self, node_id, population=None, time_window=None, **kwargs):
        selected = self._spikes_df.copy()
        mask = (selected[col_node_ids] == node_id)

        if population is not None:
            mask &= (selected[col_population] == population)

        if time_window is not None:
            mask &= (selected[col_timestamps] >= time_window[0]) & (selected[col_timestamps] <= time_window[1])

        return np.array(self._spikes_df[mask][col_timestamps])

    def node_ids(self, population=None):
        population = population if population is not None else self._defaul_population
        return np.unique(self._spikes_df[self._spikes_df[col_population] == population][col_node_ids])

    def n_spikes(self, population=None):
        population = population if population is not None else self._defaul_population
        return len(self.to_dataframe(populations=population))

    # def time_range(self, populations=None):
    #     selected = self._spikes_df.copy()
    #     if populations is not None:
    #         if isinstance(populations, six.string_types) or np.isscalar(populations):
    #             mask = selected[col_population] == populations
    #         else:
    #             mask = selected[col_population].isin(populations)
    #
    #         selected = selected[mask]
    #
    #     return selected[col_timestamps].agg([np.min, np.max]).values

    def spikes(self, node_ids=None, populations=None, time_window=None, sort_order=SortOrder.none, **kwargs):
        selected = self._spikes_df.copy()

        mask = True
        if populations is not None:
            if isinstance(populations, six.string_types) or np.isscalar(populations):
                mask &= selected[col_population] == populations
            else:
                mask &= selected[col_population].isin(populations)

        if node_ids is not None:
            node_ids = [node_ids] if np.isscalar(node_ids) else node_ids
            mask &= selected[col_node_ids].isin(node_ids)

        if time_window is not None:
            mask &= (selected[col_timestamps] >= time_window[0]) & (selected[col_timestamps] <= time_window[1])

        if isinstance(mask, pd.Series):
            selected = selected[mask]

        if sort_order == SortOrder.by_time:
            selected.sort_values(by=col_timestamps, inplace=True)
        elif sort_order == SortOrder.by_id:
            selected.sort_values(by=col_node_ids, inplace=True)

        indicies = selected.index.values
        for indx in indicies:
            yield tuple(self._spikes_df.iloc[indx])

    def __len__(self):
        if self._n_spikes is None:
            self._n_spikes = len(self._spikes_df)

        return self._n_spikes


class NWBSTReader(SpikeTrainsReadOnlyAPI):
    def __init__(self, path, **kwargs):
        self._path = path
        self._h5_file = h5py.File(self._path, 'r')
        self._n_spikes = None
        self._spikes_df = None

        # TODO: Check for other versions
        self._population = kwargs.get('population', pop_na)
        if 'trial' in kwargs.keys():
            self._trial = kwargs['trial']
        elif len(self._h5_file['/processing']) == 1:
            self._trial = list(self._h5_file['/processing'].keys())[0]
        else:
            raise Exception('Please specify a trial')

        self._trial_grp = self._h5_file['processing'][self._trial]['spike_train']

    @property
    def populations(self):
        return [self._population]

    def node_ids(self, population=None):
        # if populations is None:
        #     populations = [self._population]
        # elif isinstance(populations, six.string_types):
        #     populations = [populations]

        if self._population != population:
            return []
        return [(self._population, np.uint64(node_id)) for node_id in self._trial_grp.keys()]

    def n_spikes(self, population=None):
        if population != self._population:
            return 0

        return self.__len__()

    # def time_range(self, populations=None):
    #     data_df = self.to_dataframe()
    #     return data_df[col_timestamps].agg([np.min, np.max]).values

    def get_times(self, node_id, population=None, time_window=None, **kwargs):
        try:
            spiketimes = self._trial_grp[str(node_id)]['data'][()]

            if time_window is not None:
                spiketimes = spiketimes[(time_window[0] <= spiketimes) & (spiketimes <= time_window[1])]

            return spiketimes
        except KeyError:
            return []

    def to_dataframe(self, node_ids=None, populations=None, time_window=None, sort_order=SortOrder.none, **kwargs):
        if self._spikes_df is None:
            self._spikes_df = pd.DataFrame({
                col_timestamps: pd.Series(dtype=float),
                col_population: pd.Series(dtype=np.string_),
                col_node_ids: pd.Series(dtype=np.uint64)
            })
            for node_id, node_grp in self._trial_grp.items():
                timestamps = node_grp['data'][()]
                node_df = pd.DataFrame({
                    col_timestamps: timestamps,
                    col_population: self._population,
                    col_node_ids: np.uint64(node_id)
                })
                self._spikes_df = self._spikes_df.append(node_df, ignore_index=True)

        selected = self._spikes_df.copy()
        mask = True
        if populations is not None:
            if isinstance(populations, six.string_types) or np.isscalar(populations):
                mask &= selected[col_population] == populations
            else:
                mask &= selected[col_population].isin(populations)

        if node_ids is not None:
            node_ids = [node_ids] if np.isscalar(node_ids) else node_ids
            mask &= selected[col_node_ids].isin(node_ids)

        if time_window is not None:
            mask &= (selected[col_timestamps] >= time_window[0]) & (selected[col_timestamps] <= time_window[1])

        if isinstance(mask, pd.Series):
            selected = selected[mask]

        if sort_order == SortOrder.by_time:
            selected.sort_values(by=col_timestamps, inplace=True)
        elif sort_order == SortOrder.by_id:
            selected.sort_values(by=col_node_ids, inplace=True)

        return selected

    def spikes(self, node_ids=None, populations=None, time_window=None, sort_order=SortOrder.none, **kwargs):
        if populations is not None:
            if np.isscalar(populations) and populations != self._population:
                raise StopIteration
            elif self._population not in populations:
                raise StopIteration

        if sort_order == SortOrder.by_time:
            spikes_df = self.to_dataframe()
            spikes_df = spikes_df.sort_values(col_timestamps)
            for indx in spikes_df.index:
                r = spikes_df.loc[indx]
                yield (r[col_timestamps], r[col_population], r[col_node_ids])
        else:
            node_ids = np.array(list(self._trial_grp.keys()), dtype=np.uint64)
            if sort_order == SortOrder.by_id:
                node_ids.sort()

            for node_id in node_ids:
                timestamps = self._trial_grp[str(node_id)]['data']
                for ts in timestamps:
                    yield (ts, self._population, node_id)

    def __len__(self):
        if self._n_spikes is None:
            self._n_spikes = 0
            for node_id in self._trial_grp.keys():
                self._n_spikes += len(self._trial_grp[node_id]['data'])

        return self._n_spikes
