import os
import glob
import nixio as nix
import numpy as np
import matplotlib.pyplot as plt

from util import firing_rate

from IPython import embed

def read_baseline(block):
    spikes = []
    if "baseline" not in block.name:
        print("Block %s does not appear to be a baseline block!" % block.name )
        return spikes
    spikes = block.data_arrays[0][:]
    return spikes


def sort_blocks(nix_file):
    block_map = {}
    contrasts = []
    deltafs = []
    conditions = []
    for b in nix_file.blocks:
        if "baseline" not in b.name.lower(): 
            name_parts = b.name.split("_")
            cntrst = float(name_parts[1])
            if cntrst not in contrasts:
                contrasts.append(cntrst)
            cndtn = name_parts[3]
            if cndtn not in conditions:
                conditions.append(cndtn)
            dltf = float(name_parts[5])
            if dltf not in deltafs:
                deltafs.append(dltf)
            block_map[(cntrst, dltf, cndtn)] = b
        else:
            block_map["baseline"] = b
    return block_map, contrasts, deltafs, conditions


def get_firing_rate(block_map, df, contrast, condition, kernel_width=0.0005):
    block = block_map[(contrast, df, condition)]
    print((contrast, df, condition))
    response_map = {}
    for da in block.data_arrays:
        if "spike_times" in da.type and "response" in da.name:
            resp_id = int(da.name.split("_")[-1])
            response_map[resp_id] = da
    duration = float(block.metadata["stimulus parameter"]["duration"])
    dt = float(block.metadata["stimulus parameter"]["dt"])
    time = np.arange(0.0, duration, dt)
    rates = np.zeros((len(response_map.keys()), len(time)))
    for i, k in enumerate(response_map.keys()):
        rates[i,:] = firing_rate(response_map[k][:], duration, kernel_width, dt)
    return time, rates


def create_response_plot(block_map, all_dfs, all_contrasts, all_conditions, current_df):
    conditions = ["no-other", "self", "other"]
    condition_labels = ["alone", "self", "other"]
    
    fig = plt.figure(figsize=(4.5, 5.5))
    
    
    embed()


def process_cell(filename, dfs=[], contrasts=[], conditions=[]):
    nf = nix.File.open(filename, nix.FileMode.ReadOnly)
    block_map, all_dfs, all_contrasts, all_conditions  = sort_blocks(nf)
    if "baseline" in block_map.keys():
        baseline_spikes = read_baseline(block_map["baseline"])
    else:
        print("ERROR: no baseline data for file %s!" % filename)


    """
    if len(dfs) == 0:
        dfs = all_dfs
    if len(contrasts) == 0:
        contrasts = all_contrasts
    if len(conditions) == 0:
        conditions = all_conditions
    
    for df in dfs:
        for condition in conditions:
            for contrast in contrasts:
                time, rates = get_firing_rate(block_map, df, contrast, condition, kernel_width=0.0025)
    """         
    
    nf.close()


def main():
    nix_files = sorted(glob.glob("cell*.nix"))
    for nix_file in nix_files:
        process_cell(nix_file, dfs=[20], contrasts=[20], conditions=["self"])


if __name__ == "__main__":
    main()