import os
from pathlib import Path
import numpy as np
from spikeextractors import SortingExtractor
from spikeextractors.extractors.bindatrecordingextractor import BinDatRecordingExtractor
from spikeextractors.extraction_tools import save_to_probe_file, load_probe_file, check_valid_unit_id

try:
    import hybridizer.io as sbio
    import hybridizer.probes as sbprb
    HAVE_SBEX = True
except ImportError:
    HAVE_SBEX = False


class SHYBRIDRecordingExtractor(BinDatRecordingExtractor):
    extractor_name = 'SHYBRIDRecording'
    installed = HAVE_SBEX
    is_writable = True
    mode = 'file'
    installation_mesg = "To use the SHYBRID extractors, install SHYBRID: \n\n pip install shybrid\n\n"

    def __init__(self, file_path):
        # load params file related to the given shybrid recording
        assert HAVE_SBEX, self.installation_mesg
        params = sbio.get_params(file_path)['data']

        # create a shybrid probe object
        probe = sbprb.Probe(params['probe'])
        nb_channels = probe.total_nb_channels

        # translate the byte ordering
        # TODO still ambiguous, shybrid should assume time_axis=1, since spike interface makes an assumption on the byte ordering
        byte_order = params['order']
        if byte_order == 'C':
            time_axis = 1
        elif byte_order == 'F':
            time_axis = 0

        # piggyback on binary data recording extractor
        BinDatRecordingExtractor.__init__(self,
                                          file_path,
                                          params['fs'],
                                          nb_channels,
                                          params['dtype'],
                                          time_axis=time_axis)
        self._kwargs = {'file_path': str(Path(file_path).absolute())}
        self = load_probe_file(self, params['probe'])


    @staticmethod
    def write_recording(recording, save_path, initial_sorting_fn, dtype='float32'):
        """ Convert and save the recording extractor to SHYBRID format

        parameters
        ----------
        recording: RecordingExtractor
            The recording extractor to be converted and saved
        save_path: str
            Full path to desired target folder
        initial_sorting_fn: str
            Full path to the initial sorting csv file (can also be generated
            using write_sorting static method from the SHYBRIDSortingExtractor)
        dtype: dtype
            Type of the saved data. Default float32.
        """
        assert HAVE_SBEX, SHYBRIDRecordingExtractor.installation_mesg
        RECORDING_NAME = 'recording.bin'
        PROBE_NAME = 'probe.prb'
        PARAMETERS_NAME = 'recording.yml'

        # location information has to be present in order for shybrid to
        # be able to operate on the recording
        if 'location' not in recording.get_shared_channel_property_names():
            raise GeometryNotLoadedError("Channel locations were not found")

        # write recording
        recording_fn = os.path.join(save_path, RECORDING_NAME)
        BinDatRecordingExtractor.write_recording(recording, recording_fn,
                                                 time_axis=0, dtype=dtype)

        # write probe file
        probe_fn = os.path.join(save_path, PROBE_NAME)
        save_to_probe_file(recording, probe_fn)

        # create parameters file
        parameters = params_template.format(initial_sorting_fn=initial_sorting_fn,
                                            data_type=dtype,
                                            sampling_frequency=str(recording.get_sampling_frequency()),
                                            byte_ordering='F',
                                            probe_fn=probe_fn)

        # write parameters file
        parameters_fn = os.path.join(save_path, PARAMETERS_NAME)
        with open(parameters_fn, 'w') as fp:
            fp.write(parameters)


class SHYBRIDSortingExtractor(SortingExtractor):
    extractor_name = 'SHYBRIDSortingExtractor'
    installed = HAVE_SBEX
    is_writable = True
    installation_mesg = "To use the SHYBRID extractors, install SHYBRID: \n\n pip install shybrid\n\n"

    def __init__(self, file_path, delimiter=','):
        assert HAVE_SBEX, self.installation_mesg
        SortingExtractor.__init__(self)

        if os.path.isfile(file_path):
            self._spike_clusters = sbio.SpikeClusters()
            self._spike_clusters.fromCSV(file_path, None, delimiter=delimiter)
        else:
            raise FileNotFoundError('the ground truth file "{}" could not be found'.format(file_path))
        self._kwargs = {'file_path': str(Path(file_path).absolute()), 'delimiter': delimiter}

    def get_unit_ids(self):
        return self._spike_clusters.keys()

    @check_valid_unit_id
    def get_unit_spike_train(self, unit_id, start_frame=None, end_frame=None):
        start_frame, end_frame = self._cast_start_end_frame(start_frame, end_frame)
        train = self._spike_clusters[unit_id].get_actual_spike_train().spikes

        if start_frame is None:
            start_frame = 0
        if end_frame is None:
            end_frame = np.Inf

        idxs = np.where((start_frame <= train) & (train < end_frame))
        return train[idxs]

    @staticmethod
    def write_sorting(sorting, save_path):
        """ Convert and save the sorting extractor to SHYBRID CSV format

        parameters
        ----------
        sorting : SortingExtractor
            The sorting extractor to be converted and saved
        save_path : str
            Full path to the desired target folder
        """
        assert HAVE_SBEX, SHYBRIDSortingExtractor.installation_mesg
        dump = np.empty((0, 2))

        for unit_id in sorting.get_unit_ids():
            spikes = sorting.get_unit_spike_train(unit_id)[:,np.newaxis]
            expanded_id = (np.ones(spikes.size) * unit_id)[:,np.newaxis]
            tmp_concat = np.concatenate((expanded_id, spikes), axis=1)

            dump = np.concatenate((dump, tmp_concat), axis=0)

        sorting_fn = os.path.join(save_path, 'initial_sorting.csv')
        np.savetxt(sorting_fn, dump, delimiter=',', fmt='%i')


class GeometryNotLoadedError(Exception):
    """ Raised when the recording extractor has no associated channel locations
    """
    pass

params_template = \
"""clusters:
  csv: {initial_sorting_fn}
data:
  dtype: {data_type}
  fs: {sampling_frequency}
  order: {byte_ordering}
  probe: {probe_fn}
"""
