# -*- coding: utf-8 -*-
"""
Created by J McCutcheon
22 Feb 2018
To analyze data from Med PC files or text/csv files and calculate/output lick parameters.
"""

# Import statements
from tkinter import *
from tkinter import ttk
from tkinter import filedialog
from tkinter import messagebox
import os
import string
import numpy as np
import scipy.optimize as opt
import scipy.stats as stats
import matplotlib as mpl
mpl.use("TkAgg")
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from matplotlib.figure import Figure
from matplotlib.backends.backend_pdf import PdfPages
import ntpath
import csv
import xlsxwriter as xl
import datetime
from pathlib import Path

from trompy import alert, get_location, lickCalc, sessionlicksFig, iliFig, burstlengthFig, burstprobFig, licklengthFig, isnumeric

# Main class for GUI
class Window_lick(Frame):

    def __init__(self, master=None):
        f1 = ttk.Style()
        f1.configure('.', background='powder blue', padding=5)
        f1.configure('TButton', width=15, sticky=(E,W))
        f1.configure('TMenubutton', background='light cyan', padding=0)
        f1.configure('header.TLabel', font='Helvetica 12')
        f2 = ttk.Style()
        f2.configure('inner.TFrame', background='light cyan')
        
        ttk.Frame.__init__(self, master, style='TFrame', padding=(10, 10, 15, 15))
      
        self.master = master
        
        self.init_window()

    def init_window(self):
        self.master.title('MedfileReader')
        self.pack(fill=BOTH, expand=1)
        
        #Frame for graphs
        self.f2 = ttk.Frame(self, style='inner.TFrame', borderwidth=5,
                            relief="sunken", width=200, height=300)

        #Set up standalone labels
        self.fileparamslbl = ttk.Label(self, text='File Parameters', style='header.TLabel')
        self.calcparamslbl = ttk.Label(self, text='Calculator Parameters', style='header.TLabel')
        self.graphparamslbl = ttk.Label(self, text='Graph Parameters', style='header.TLabel')
        
        self.fileformatlbl = ttk.Label(self, text='File Format')
        self.headerrowslbl = ttk.Label(self, text='Header Rows')
        
        self.onsetlbl = ttk.Label(self, text='Onset')
        self.offsetlbl = ttk.Label(self, text='Offset')
        self.IBthresholdlbl = ttk.Label(self, text='Interburst threshold')
        self.IRthresholdlbl = ttk.Label(self, text='Interrun threshold')
        
        self.nolongILIslbl = ttk.Label(self, text='Ignore Long ILIs')
        self.minburstlengthlbl = ttk.Label(self, text='Minimum burst')
        self.plotburstproblbl = ttk.Label(self, text='Plot burst prob.')
        
        self.outputlbl = ttk.Label(self, text='Output Parameters', style='header.TLabel')
        self.suffixlbl = ttk.Label(self, text='File suffix')
        self.aboutlbl = ttk.Label(self, text='LickCalc-1.4.2 by J McCutcheon')
  
        #Set up Entry variables
        self.shortfilename = StringVar(self.master)
        self.shortfilename.set('No file loaded')
        self.filenamelbl = ttk.Label(self, textvariable=self.shortfilename)
        
        self.headerrows = StringVar(self.master)
        self.headerrowsField = ttk.Entry(self, textvariable=self.headerrows)
        self.headerrowsField.insert(END, '5')
        
        self.IBthreshold = StringVar(self.master)
        self.IRthreshold = StringVar(self.master)
        self.minburst = StringVar(self.master)

        self.IBthresholdField = ttk.Entry(self, textvariable=self.IBthreshold)
        self.IBthresholdField.insert(END,'0.5')
        
        self.IRthresholdField = ttk.Entry(self, textvariable=self.IRthreshold)
        self.IRthresholdField.insert(END,'10')
        
        self.minburstField = ttk.Entry(self, textvariable=self.minburst)
        self.minburstField.insert(END, '1')
        
        self.suffix = StringVar(self.master)
        self.suffixField = ttk.Entry(self, textvariable=self.suffix)
        self.suffixField.insert(END,'')
        
        # Set up Dropdown buttons
        self.FILEOPTIONS = ['Med Associates', '.txt', '.csv', 'DD Lab', 'SF Lab']
        self.fileformat = StringVar(self.master)
        self.fileformatMenu = ttk.OptionMenu(self, self.fileformat, self.FILEOPTIONS[0], *self.FILEOPTIONS)
   
        self.OPTIONS = ['None']
        self.onset = StringVar(self.master)
        self.onsetButton = ttk.OptionMenu(self, self.onset, *self.OPTIONS)    
        self.offset = StringVar(self.master)
        self.offsetButton = ttk.OptionMenu(self, self.offset, *self.OPTIONS)

        #Set up Boolean variables
        self.nolongILIs = BooleanVar(self.master)
        self.nolongILIs.set(False)
        self.nolongILIsButton = ttk.Checkbutton(self, variable=self.nolongILIs, onvalue=True)
        
        self.plotburstprob = BooleanVar(self.master)
        self.plotburstprob.set(False)
        self.plotburstprobButton = ttk.Checkbutton(self, variable=self.plotburstprob, onvalue=True)
        
        # Set up Buttons
        self.loadfileButton = ttk.Button(self, text='Load File', command=self.openfile)
        self.analyzeButton = ttk.Button(self, text='Analyze Data', command=self.analyze)
        self.prevButton = ttk.Button(self, text='Previous', command=lambda: self.load_adj_files(delta=-1))
        self.nextButton = ttk.Button(self, text='Next', command=lambda: self.load_adj_files(delta=1))
        
        self.defaultfolderButton = ttk.Button(self, text='Default Folder', command=self.setsavefolder)
        self.pdfButton = ttk.Button(self, text='PDF', command=self.makePDF)
        self.excelButton = ttk.Button(self, text='Excel', command=self.makeExcel)
        self.textsummaryButton = ttk.Button(self, text='Text Summary', command=self.maketextsummary)
        
        #Place items in grid
        self.fileparamslbl.grid(column=0, row=0, columnspan=2)
        self.calcparamslbl.grid(column=2, row=0, columnspan=2)
        self.graphparamslbl.grid(column=4, row=0, columnspan=2)
 
        self.fileformatlbl.grid(column=0, row=1, sticky=E)
        self.fileformatMenu.grid(column=1, row=1, sticky=(W,E))
        
        self.headerrowslbl.grid(column=0, row=2, sticky=E)
        self.headerrowsField.grid(column=1, row=2, sticky=(W,E))
        
        self.loadfileButton.grid(column=0, row=3, columnspan=2, sticky=(W,E))

        self.prevButton.grid(column=0, row=4, sticky=(W,E))
        self.nextButton.grid(column=1, row=4, sticky=(W,E))

        self.onsetlbl.grid(column=2, row=1, sticky=E)
        self.offsetlbl.grid(column=2, row=2, sticky=E)
        self.IBthresholdlbl.grid(column=2, row=3, sticky=E)
        self.IRthresholdlbl.grid(column=2, row=4, sticky=E)
        
        self.onsetButton.grid(column=3, row=1, sticky=(W,E), pady=5)
        self.offsetButton.grid(column=3, row=2, sticky=(W,E), pady=5)
        self.IBthresholdField.grid(column=3, row=3, sticky=(W,E))
        self.IRthresholdField.grid(column=3, row=4, sticky=(W,E))
        
        self.nolongILIslbl.grid(column=4, row=1)
        self.nolongILIsButton.grid(column=5, row=1)
        
        self.minburstlengthlbl.grid(column=4, row=2)
        self.minburstField.grid(column=5, row=2)
        
        self.plotburstproblbl.grid(column=4, row=3)
        self.plotburstprobButton.grid(column=5, row=3)
        
        self.outputlbl.grid(column=0, row=6, sticky=(W,E), pady=5)
        self.suffixlbl.grid(column=1, row=6, sticky=(E), pady=5)
        self.suffixField.grid(column=2, row=6, sticky=(W,E), pady=5)
        
        self.defaultfolderButton.grid(column=3, row=6, sticky=(W,E), pady=5)
        self.pdfButton.grid(column=4, row=6, sticky=(W,E), pady=5)
        self.excelButton.grid(column=5, row=6, sticky=(W,E), pady=5)
        self.textsummaryButton.grid(column=6, row=6, sticky=(W,E), pady=5)
        
        self.aboutlbl.grid(column=0, row=7, columnspan=7, sticky=W)
        
        self.analyzeButton.grid(column=6, row=1, rowspan=4, sticky=(N, S, E, W))

        self.f2.grid(column=0, row=5, columnspan=7, sticky=(N,S,E,W))
                    
    def openfile(self):
        self.filename = filedialog.askopenfilename(initialdir=os.getcwd(), title='Select a file.')
        self.list_of_files = []
        ff = self.fileformat.get()
        if ff == 'Med Associates':
            self.loadmedfile()
        elif ff == '.txt':
            self.loadcsvfile()
        elif ff == 'DD Lab':
            self.loadDDfile()
        elif ff == 'SF Lab':
            self.loadmedfile()
        else:
            print('Not valid format')
        
    def loadmedfile(self):        
        try:
            if len(checknsessions(self.filename)) > 1:
                alert('More than one session in file. Analysing session 1.')
            else:                    
                self.loaded_vars = medfilereader_licks(self.filename)
        except:
            alert("Problem reading file and extracting data. File may not be properly formatted - see Help for advice.")
            return
        if self.fileformat.get() == 'SF Lab':
            new_vars = {}
            for v in self.loaded_vars:
                new_vars[v] = [x/10 for x in self.loaded_vars[v]]
            self.loaded_vars = new_vars
        try:
            self.updateOptionMenu()
        except TypeError:
            alert("No valid variables to analyze (e.g. arrays with more than one value")
        
        self.currentfiletype = 'med'
        self.shortfilename.set(ntpath.basename(self.filename))
         
    def loadcsvfile(self):               
        try:            
            with open(self.filename, newline='') as myFile:
                reader = csv.DictReader(myFile)
                cols = reader.fieldnames
                self.loaded_vars = {}
                for col in cols:
                    self.loaded_vars[col] = []
                    myFile.seek(0)
                    for row in reader:
                        try:
                            self.loaded_vars[col].append(float(row[col]))
                        except:
                            pass 
        except:
            alert('Cannot load data from selected file. Is it a CSV?')
            return
                        
        try:
            self.updateOptionMenu()
        except TypeError:
            alert("No valid variables to analyze (e.g. arrays with more than one value")
        
        self.currentfiletype = 'csv'
        self.shortfilename.set(ntpath.basename(self.filename))
    
    def loadDDfile(self):
        try:
            header=int(self.headerrows.get())
            f = open(self.filename, 'r')
            vals = f.readlines()[header:]
            f.close()
            ts = [tstamp_to_tdate(val, '%H:%M:%S.%f\n') for val in vals]
            delta = [t-ts[idx-1] for idx, t in enumerate(ts[1:])]
            delta_array = np.array([d.total_seconds() for d in delta])
            if min(delta_array) < 0: #tests if timestamps span midnight
                dayadvance = np.where(delta_array < 0)[0][0]
                ts = ts[:dayadvance+1] + [t + datetime.timedelta(days=1) for t in ts[dayadvance+1:]]
            t0=ts[0]
            self.loaded_vars = {'t': [(t - t0).total_seconds() for t in ts]}

        except:
            alert('Cannot load data from selected file. Is it in the right format?')
            print("Error:", sys.exc_info()[0])
            return
        
        try:
            self.updateOptionMenu()
        except TypeError:
            alert("No valid variables to analyze (e.g. arrays with more than one value")
            
        self.currentfiletype = 'dd'
        self.shortfilename.set(ntpath.basename(self.filename))

    def updateOptionMenu(self):
        options = [x+': '+str(len(self.loaded_vars[x])) for x in self.loaded_vars]
        self.onsetButton = ttk.OptionMenu(self, self.onset, *options).grid(column=3, row=1, sticky=(W,E))
        self.offsetButton = ttk.OptionMenu(self, self.offset, *options).grid(column=3, row=2, sticky=(W,E))

    def load_adj_files(self,delta=1): #delta+1 = next, -1=prev        
        try:
            if not self.list_of_files:
                self.currpath = ntpath.dirname(self.filename)
                self.list_of_files = os.listdir(self.currpath)
            index = [x[0] for x in enumerate(self.list_of_files) if x[1] == self.shortfilename.get()]
            newindex = index[0] + delta
            self.filename = os.path.join(self.currpath, self.list_of_files[newindex])
            
            if self.currentfiletype == 'med':
                self.loadmedfile()
            if self.currentfiletype == 'csv':
                self.loadcsvfile()
            if self.currentfiletype == 'dd':
                self.loadDDfile()
        except:
            alert('Problem loading next file. It might be at the end of the folder or in the wrong format.')
    
    def analyze(self):
        print('Analyzing...')
        
        # Check inputs
        try:
            burstTH = float(self.IBthreshold.get())
        except ValueError:
            alert('Interburst threshold value needs to be numeric')
            return
        
        try:
            runTH = float(self.IRthreshold.get())
        except ValueError:
            alert('Interrun threshold value needs to be numeric')
            return        

        if hasattr(self, 'filename'):            
            try:
                self.onsetArray = self.loaded_vars[self.onset.get().split(':')[0]]
                try:
                    self.offsetArray = self.loaded_vars[self.offset.get().split(':')[0]]
                    self.lickdata = lickCalc(self.onsetArray, offset=self.offsetArray, burstThreshold = burstTH, runThreshold = runTH,
                                             ignorelongilis=self.nolongILIs.get(),
                                             minburstlength=int(self.minburst.get()))
                except:
                    self.lickdata = lickCalc(self.onsetArray, burstThreshold = burstTH, runThreshold = runTH,
                                             ignorelongilis=self.nolongILIs.get(),
                                             minburstlength=int(self.minburst.get()))

            except:
                alert('Have you picked an onset array yet?')
                print("Error:", sys.exc_info()[0])               
                raise
        
        else:
            print('Select a file first')
            messagebox.showinfo("Error", "Select a valid file first.")
               
        self.makegraphs()
                        
    def makegraphs(self):
        
        f = Figure(figsize=(8.27, 5))
        f.suptitle(self.shortfilename.get())
        grid = mpl.gridspec.GridSpec(2, 3, wspace=0.5, hspace=0.5)
        ax1 = f.add_subplot(grid[0,:])
        ax2 = f.add_subplot(grid[1,0])
        ax3 = f.add_subplot(grid[1,1])
        ax4 = f.add_subplot(grid[1,2])

        # Licks over session 
        sessionlicksFig(ax1, self.onsetArray)
        
        # Lick parameter figures
        iliFig(ax2, self.lickdata)
        
        if self.plotburstprob.get() == False:
            burstlengthFig(ax3, self.lickdata)
        else:
            self.weibull_fit = burstprobFig(ax3, self.lickdata)
            
        licklengthFig(ax4, self.lickdata)
        
        canvas = FigureCanvasTkAgg(f, self)
        #canvas.show()
        canvas.get_tk_widget().grid(row=5, column=0, columnspan=7, sticky=(N,S,E,W))
      
        return f
    
    def setsavefolder(self):
        self.savefolder = Path(get_location())
    
    def makePDF(self):
        if not hasattr(self, 'savefolder'):
            self.setsavefolder()
            
        savefile = self.savefolder / f"{self.shortfilename.get()}_{self.suffix.get()}.pdf"
        try:
            pdfFig = self.makegraphs()
            pdf_pages = PdfPages(savefile)
            pdf_pages.savefig(pdfFig)
            pdf_pages.close()
        except:
            print("Error:", sys.exc_info()[0])
            alert('Problem making PDF! Is data loaded and analyzed?')
    
    def makeExcel(self):
        if not hasattr(self, 'savefolder'):
            self.setsavefolder()
            
        savefile = self.savefolder / f"{self.shortfilename.get()}_{self.suffix.get()}.xlsx"

        self.makesummarydictionary()
        
        wb = xl.Workbook(savefile)
        # worksheet with summary data
        sh = wb.add_worksheet('Summary')
        
        bold = wb.add_format({'bold': True})
        
        sh.set_column(0, 1, 20)
        sh.write('A1', 'Parameter', bold)
        sh.write('B1', 'Value', bold)
        for idx, vals in enumerate(self.d):
            sh.write(idx+1, 0, vals[0])
            sh.write(idx+1, 1, vals[1])
        
        # worksheet with lick timestamps
        sh = wb.add_worksheet('Licks')
        for idx, val in enumerate(self.onsetArray):
            sh.write(idx, 0, val)
        
        sh = wb.add_worksheet('ILIs')
        for idx, val in enumerate(self.lickdata['ilis']):
            sh.write(idx, 0, val)

        sh = wb.add_worksheet('Bursts')
        for idx, val in enumerate(self.lickdata['bLicks']):
            sh.write(idx, 0, val)
            
        wb.close()
#        except:
#            alert('Working on making an Excel file')
        
    def maketextsummary(self):
        if not hasattr(self, 'savefolder'):
            self.setsavefolder()
            
        savefile = self.savefolder / f"{self.shortfilename.get()}_{self.suffix.get()}-text_summary.csv"
        try:
            self.makesummarydictionary()
            
            with open(savefile, 'w', newline='') as file:
                csv_out = csv.writer(file)
                csv_out.writerow(['Parameter', 'Value'])
                for row in self.d:
                    csv_out.writerow(row)
        except:
            alert('Problem making text summary!')
            print("Error:", sys.exc_info()[0])
            
    def licksperburstsummary(self):
        self.folder = Path(get_location())
        savefile = self.folder / f"{self.shortfilename.get()}-licks-per-burst.csv"
        try:
            d = self.data['bLicks']
        except:
            alert('Problem making text summary!')
            print("Error:", sys.exc_info()[0])
        
    def makesummarydictionary(self):
        self.d = [('Filename',self.shortfilename.get()),
                 ('Total licks',self.lickdata['total']),
                 ('Frequency',self.lickdata['freq']),
                 ('Number of bursts',self.lickdata['bNum']),
                 ('Licks per burst',self.lickdata['bMean']),
                 ('Licks per burst (first 3)',self.lickdata['bMean-first3']),
                 ('Number of long licks',len(self.lickdata['longlicks'])),
                 ('Weibull: alpha',self.lickdata['weib_alpha']),
                 ('Weibull: beta',self.lickdata['weib_beta']),
                 ('Weibull: rsquared',self.lickdata['weib_rsq'])]

def checknsessions(filename):
    f = open(filename, 'r')
    f.seek(0)
    filerows = f.readlines()[8:]
    datarows = [isnumeric(x) for x in filerows]
    matches = [i for i,x in enumerate(datarows) if x == 0.3]
    return matches

def medfilereader_licks(filename,
                  sessionToExtract = 1,
                  verbose = False,
                  remove_var_header = True):
    '''Gets lick data from Med Associates file.'''
    
    f = open(filename, 'r')
    f.seek(0)
    filerows = f.readlines()[8:]
    datarows = [isnumeric(x) for x in filerows]
    matches = [i for i,x in enumerate(datarows) if x == 0.3]
    if sessionToExtract > len(matches):
        print('Session ' + str(sessionToExtract) + ' does not exist.')
    if verbose == True:
        print('There are ' + str(len(matches)) + ' sessions in ' + filename)
        print('Analyzing session ' + str(sessionToExtract))
    
    varstart = matches[sessionToExtract - 1]    
    medvars = {}
   
    k = int(varstart + 27)
    for i in range(26):
        medvarsN = int(datarows[varstart + i + 1])
        if medvarsN > 1:
            medvars[string.ascii_uppercase[i]] = datarows[k:k + int(medvarsN)]
        k = k + medvarsN
    
    if remove_var_header == True:
        for val in medvars.values():
            val.pop(0)

    return medvars

def tstamp_to_tdate(timestamp, fmt):
    try:
        return datetime.datetime.strptime(timestamp, fmt)
    except ValueError:
        return

def start_lickcalc_gui():
    root = Tk()
    app = Window_lick(root)
    root.lift()
    root.mainloop()
    
if __name__ == '__main__':
    os.chdir("C:\\Github\\Lick-Calc-GUI\\output\\")
    start_lickcalc_gui()

# Files for for testing
#        self.filename = 'C:\\Users\\jaimeHP\\Dropbox\\Python\\cas9\\cas9_medfiles\\!2016-07-19_09h16m.Subject 4'
#        self.filename = 'C:\\Users\\jaimeHP\\Dropbox\\Python\\cas9\\cas9_medfiles\\!2017-06-12_10h53m.Subject thpe1.4'



