#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Jul  1 11:50:02 EDT 2019

@author: Bryan Medina
"""
###### Imports ########
from nwb_plots_functions import *
from scipy.interpolate import LSQUnivariateSpline

import h5py as h5
import matplotlib.pyplot as plt
import numpy as np
import os
import pickle
import sys
########################

###### UPDATE PATH #################################
DIRECTORY = '/Users/bjm/Documents/CMU/Research/data'
VAR_DIREC = '/Users/bjm/Documents/CMU/Research/data/plots/variations/'
SUMMARY_PLOTS_DIRECTORY = '/Users/bjm/Documents/CMU/Research/data/plots/'
MOUSE_ID = '405751'
####################################################

# Get file from directory
spikes_nwb_file = os.path.join(DIRECTORY, 'mouse' + MOUSE_ID + '.spikes.nwb')
nwb = h5.File(spikes_nwb_file, 'r')

probe_names = nwb['processing']

# save all curves for all regions
mid = {}
top = {}
bot = {}

for probe_name in probe_names:
    # Calculate median neuron, and also 90th and 10th percentile neuron
    median_n = []
    top_ten  = []
    bot_ten  = []
    
    probe_filename = MOUSE_ID + "_" + probe_name

    with open(probe_filename, 'rb') as f:
        probe = pickle.load(f)
        
    for xval in xs:
        
        rates = []
        
        for cell in probe.getCellList():
            rates.append(probe.getCell(cell).lsq(xval))

        # Sort this list...
        rates.sort()
        
        median_n.append(np.median(rates))
        top_ten.append(np.percentile(rates, 85))
        bot_ten.append(np.percentile(rates, 15))
     
    # save the curves
    mid[probe_name] = LSQUnivariateSpline(xs, median_n, knots[1:-1])
    top[probe_name] = LSQUnivariateSpline(xs, top_ten, knots[1:-1])
    bot[probe_name] = LSQUnivariateSpline(xs, bot_ten, knots)


# Plotting median, 85th percentile, and 15th percentile neuron

# Median
for probe_name in probe_names:
    
    probe_filename = MOUSE_ID + "_" + probe_name

    with open(probe_filename, 'rb') as f:
        probe = pickle.load(f)
    
    plt.ylim(0, 5)
    plt.xlim(-20, 500)
    plt.title("Median Neuron Activity for Mouse " + str(MOUSE_ID))
    plt.ylabel('Spikes/second')
    plt.xlabel('Bins')
    plt.plot(xs, mid[probe_name](xs), label=probe.name)

plt.legend()
plt.savefig(SUMMARY_PLOTS_DIRECTORY + str(MOUSE_ID) + "_MEDIAN.png")
plt.clf()

# 85th
for probe_name in probe_names:
    
    probe_filename = MOUSE_ID + "_" + probe_name

    with open(probe_filename, 'rb') as f:
        probe = pickle.load(f)
    
    plt.ylim(0, 20)
    plt.xlim(-20, 500)
    plt.title("85th Percentile Neuron Activity for Mouse " + str(MOUSE_ID))
    plt.ylabel('Spikes/second')
    plt.xlabel('Bins')
    plt.plot(xs, top[probe_name](xs), label=probe.name)

plt.legend()
plt.savefig(SUMMARY_PLOTS_DIRECTORY + str(MOUSE_ID) + "_85TH.png")
plt.clf()


# 15th
for probe_name in probe_names:
    
    probe_filename = MOUSE_ID + "_" + probe_name

    with open(probe_filename, 'rb') as f:
        probe = pickle.load(f)
    
    plt.ylim(0, 1)
    plt.xlim(-2, 500)
    plt.title("15th Percentile Neuron Activity for Mouse " + str(MOUSE_ID))
    plt.ylabel('Spikes/second')
    plt.xlabel('Bins')
    plt.plot(xs, bot[probe_name](xs), label=probe.name)

plt.legend()
plt.savefig(SUMMARY_PLOTS_DIRECTORY + str(MOUSE_ID) + "_15TH.png")
plt.clf()
