import numpy as np
from ..utils import extract_barcodes_from_times, get_first_last_times, sync_timestamp_offline, convert_port_number
from ..data import load_neuropixel_data, load_ks_output, get_channel_bank_name
import os
import glob

def sync_neuropixel_ecube(raw_timestamp,on_times_np,off_times_np,on_times_ecube,off_times_ecube,inter_barcode_interval=20,bar_duration=0.02,verbose=False):
    '''
    This function is specfic to synchronization between neuropixels and ecube.
    
    Args:
        raw_timestamp (nt) : raw timestamp that is not synchronized
        on_times_np (ndarray) : timings when sync line rises to 1 in neruopixel streams
        off_times_np (ndarray): timings when sync line returns to 0 in neruopixel streams
        on_times_ecube (ndarray): timings when sync line rises to 1 in ecube streams
        off_times_ecube (ndarray): timings when sync line returns to 0 in ecube streams
        bar_duration (float): duration of each bar that is sent to each stream
        verbose (bool): print barcode times and barcodes for each stream
    
    Returns:
        tuple: tuple containing:
            | **sync_timestamps (nt):** synchronized timestamps
            | **scaling (float):** scaling factor between streams
    '''
    
    while bar_duration >0.001:
        # Get each barcode timing and label
        barcode_ontimes_np,barcode_np = extract_barcodes_from_times(on_times_np,off_times_np,inter_barcode_interval=inter_barcode_interval,bar_duration=bar_duration)
        barcode_ontimes_ecube,barcode_ecube = extract_barcodes_from_times(on_times_ecube,off_times_ecube,inter_barcode_interval=inter_barcode_interval,bar_duration=bar_duration)

        # Check if barcodes are consistent across streams
        n_barcode = min(len(barcode_ecube),len(barcode_np))
        num_different_barcodes = 0
        for idx in range(n_barcode):
            barcode = barcode_ecube[idx]
            if barcode != barcode_np[idx]:
                num_different_barcodes += 1
        
        # If barcodes are the same across streams, break this loop
        if num_different_barcodes == 0:
            break
        bar_duration -= 0.0001
            
    if verbose:
        print(f'bar duration: {bar_duration}\n')
        print(f'neuropixel barcode times: {barcode_ontimes_np}\n')
        print(f'neuropixel barcode: {barcode_np}\n')
        print(f'ecube barcode times: {barcode_ontimes_ecube}\n')
        print(f'ecube barcode: {barcode_ecube}\n')
        
    # Get the first and last barcode timing in the recording at each stream
    first_last_times_np, first_last_times_ecube = get_first_last_times(barcode_ontimes_np,barcode_ontimes_ecube,barcode_np, barcode_ecube) 

    return sync_timestamp_offline(raw_timestamp, first_last_times_np, first_last_times_ecube)

def classify_ks_unit(spike_times, spike_label):
    '''
    Classify unit activity identified by kilosort into each single unit
    
    Args:
        spike_times (nspikes): spike times generated by kilosort
        spike_label (nspikes): cluster labels of each spike generated by kilosort
        
    Returns:
        spike_times_unit (dict): spike times for each unit
    '''
    
    spike_times_unit = {}
    
    for unit_label in np.unique(spike_label):
        spike_times_unit[f'{unit_label}'] = spike_times[spike_label==unit_label.astype(int)]
    
    return spike_times_unit

def parse_ksdata_entries(kilosort_dir, concat_data_dir):
    '''
    Parse concatenated data processed by kilosort into the task entries and save relevant data (spike_indices, spike_clusters, and ks_label)
    
    Args:
        kilosort_dir (str): kilosort directory (ex. '/data/preprocessed/kilosort')
        concat_data_dir (str): data directory that contains concatenated data and kilosort_output (ex. '2023-06-30_Neuropixel_ks_affi_bottom_port1')
        
    Returns
        None
    '''
    
    # Load kilosort data
    kilosort_output = load_ks_output(kilosort_dir, concat_data_dir, flag='spike')
    spike_indices = kilosort_output['spike_indices']
    spike_clusters = kilosort_output['spike_clusters']
    ks_label = kilosort_output['ks_label']
    
    # Load datasize of each entry and filename
    data_path = os.path.join(kilosort_dir, concat_data_dir)
    datasize_entry = np.load(os.path.join(data_path, 'datasize_entry.npy'))
    task_paths = np.load(os.path.join(data_path, 'task_path.npy'))

    nentry = datasize_entry.shape[0]

    datasize_entry = np.cumsum(datasize_entry)
    for idx in range(nentry):
        # Get spike indices included in a corresponding entry
        if idx == 0:
            ientry_idx = spike_indices < datasize_entry[idx]
        else:
            ientry_idx = (spike_indices >= datasize_entry[idx-1])&(spike_indices < datasize_entry[idx])

        # Subtract previous entries' datasize so that the first index of each entry can become 0.
        spike_indices_entry = spike_indices[ientry_idx]
        if idx > 0:
            spike_indices_entry -= datasize_entry[idx-1]
        
        # Make new directory to save parsed data
        taskid = task_paths[idx].split('_')[-1][2:]
        task_save_dir = f'{concat_data_dir}_{taskid}'
        task_save_path = os.path.join(kilosort_dir, task_save_dir)
        if not os.path.exists(task_save_path):
            os.makedirs(task_save_path)
        
        # save data
        np.save(os.path.join(task_save_path, 'spike_indices_entry'), spike_indices_entry)
        np.save(os.path.join(task_save_path, 'spike_clusters_entry'), spike_clusters[ientry_idx])
        np.save(os.path.join(task_save_path, 'ks_label'), ks_label.astype('str'))
        
def concat_neuropixel_within_day(np_datadir, kilosort_dir, subject, date, ch_config_dir='/data/channel_config_np', port_number = 1, bad_taskid = None, max_memory_gb = 0.1):
    '''
    Concatenate continuous.dat files of different sessions within the same day, subject, and electrode configuration
    Then save the concatenated continuous data into the kilosort directory with the datasize and filenames of original task entries
    
    Args:
        np_datadir (str): neuropixel data directory where the data files are located (ex. '/data/raw/neuropixels')
        kilosort_dir (str): kilosort directory (ex. '/data/preprocessed/kilosort')
        subject (str): subject name
        date (str): date (str): date of recording (ex. '2023-04-14')
        ch_config_dir (str, optional): directory that contains the channel configuration file
        port_number (int, optional): port number which a probe connected to. natural number from 1 to 4.
        bad_taskid (list): a list of bad task id that should not be included in concatenating data
        max_memory_gb (float): max memory used to load binary data at one time
        
    Returns:
        savedir_names (list): a list of directory names where concatenated data was saved
    '''
    # Find path for task entries within day and subject
    data_path = np.array(glob.glob(os.path.join(np_datadir,f'{date}_Neuropixel_{subject}_te*'),recursive=True))
    data_path = np.sort(data_path) # sort in the order of task_id
    
    print(data_path)

    if data_path.size > 0:
        # Chack which bank for each channel was used
        bank_name_list = []
        for idx, path in enumerate(data_path):
            np_recorddir = os.path.split(path)[1]
            _, metadata = load_neuropixel_data(np_datadir,np_recorddir,'ap',port_number=port_number)
            bank_name_list.append(get_channel_bank_name(metadata['ch_bank'], ch_config_dir=ch_config_dir))
        bank_name_list = np.array(bank_name_list)
        unique_bank = np.unique(bank_name_list)
        nch = metadata['num_channels'] # assumed nch is the same across sessions
        
        dtype = 'int16'
        savedir_names = []
        for ibank in unique_bank:
            # data path that has the same bank information
            bank_data_path = np.array(data_path[bank_name_list == ibank])
            
            if bad_taskid:
                bad_taskid_label = np.zeros(len(bank_data_path)) == 0
                for badid in bad_taskid:
                    for idx, path in enumerate(bank_data_path):
                        bad_taskid_label[idx] *= not badid in path

                bank_data_path = bank_data_path[bad_taskid_label]

            # Make a new directory to save concatenated data by subject, bank, port number
            savedir_name = f'{date}_Neuropixel_ks_{subject}_{ibank}_port{port_number}'
            savedir_path = os.path.join(kilosort_dir, savedir_name)
            if not os.path.exists(savedir_path):
                os.makedirs(savedir_path)
            savedir_names.append(savedir_name)
            
            datasize_entry = []
            file_path = []
            for idx, path in enumerate(bank_data_path):
                print(f'Processing {path} ({ibank})')
                    
                # Find data path for continuous.dat recorded by the same probe
                probe_dir = convert_port_number(port_number)
                continuous_data_path = glob.glob(os.path.join(path,f'**/*{probe_dir}/continuous.dat'),recursive=True)[0]

                # Load continuous data as numpy.memmap
                np_recorddir = os.path.split(path)[1]
                data = np.memmap(continuous_data_path, dtype=dtype)

                # Store file path and datasize in original data
                file_path.append(np_recorddir)
                datasize_entry.append(int(data.shape[0]/nch))

                # Save data (concatentate data)
                chunksize = int(max_memory_gb * 1e9 / np.dtype(dtype).itemsize / nch)
                nchunk =int(np.ceil(data.shape[0]/chunksize))
                
                for ichunk in range(nchunk):
                    if ichunk != nchunk-1:
                        data_reshape = data[ichunk*chunksize*nch:(ichunk+1)*chunksize*nch].reshape(-1,nch)
                    else:
                        data_reshape = data[ichunk*chunksize*nch:].reshape(-1,nch)
                        
                    if (idx == 0) & (ichunk == 0):
                        save_filename = 'continuous.dat'
                        f = open(os.path.join(savedir_path,save_filename), 'wb') # save
                        data_reshape.tofile(f)
                        f.close()
                    else:
                        f = open(os.path.join(savedir_path,save_filename), 'a') # append
                        data_reshape.tofile(f)
                        f.close()

            # Save datasize and filename of each entry for preprocessing after spike sorting
            np.save(os.path.join(savedir_path, 'datasize_entry'), np.array(datasize_entry))
            np.save(os.path.join(savedir_path, 'task_path'), np.array(file_path))

    else:
        print('No data to concatenate')

    return savedir_names