try:
    # for Python2
    from Tkinter import Tk, messagebox, Label, Button, Entry, FileDialog, RIGHT, YES, X, DISABLED, StringVar, Frame, TOP, LEFT
except ImportError:
    # for Python3
    from tkinter import Tk, messagebox, Label, Button, Entry, filedialog, RIGHT, YES, X, DISABLED, StringVar, Frame, TOP, LEFT

import numpy as np
import os
from itertools import groupby
from operator import itemgetter


#Input Form
fields = ('Vehicles Ignored', 'Saturation Flow Headway', 'Min Vehs per Measurement', 'Minimum Valid Cycles')

class SatFlowsCalc:
    def check_params(self, *args):
        for p in self.params:
            try:
                t = int(p.get())
                if t>0:
                    self.run_button.config(state='normal')
                else:
                    self.run_button.config(state='disabled')
                    break
            except ValueError:
                break

    def __init__(self, master):
        self.master = master
        master.title("Saturation Flows Calculator")

        vcmd    = master.register(self.validate) # we have to wrap the command
        vcmd_fp = master.register(self.validate_inputFolder) # we have to wrap the command


        self.fp_dis = StringVar()
        e = Entry(master, textvariable=self.fp_dis,validate="key", validatecommand=(vcmd_fp,'%P'))
        b = Button(master, text="Browse",
                   command=lambda:self.fp_dis.set(filedialog.askdirectory()))
        e.pack()
        b.pack()

        self.vehsIgnored = StringVar()
        self.satFlowHead = StringVar()
        self.minVehsPerMeas = StringVar()
        self.minValidCycles = StringVar()

        self.params = [self.vehsIgnored, self.satFlowHead, self.minVehsPerMeas, self.minValidCycles]

        i=0
        for p in self.params:
            row = Frame(master)
            lab = Label(row, width=22, text=fields[i]+": ", anchor='w')
            self.ent = Entry(row, textvariable=p, validate="key", validatecommand=(vcmd, '%P'))
            self.ent.insert(0,'0')
            row.pack(side=TOP, fill=X, padx=5, pady=5)
            lab.pack(side=LEFT)
            self.ent.pack(side=RIGHT, expand=YES, fill=X)
            p.trace("w", self.check_params)
            i+=1

        self.run_button = Button(master, text='Calculate Saturation Flow',
                                 state=DISABLED, command=lambda: self.Main(self.params))
        self.run_button.pack(padx=5, pady=5,expand=YES, fill=X)


    def validate_inputFolder(self, fp):
        print('in')
        if os.listdir(fp):
            return False
        else:
            return True

    def validate(self, new_text):
        if not new_text: # the field is being cleared
            self.entered_number = 0
            return True
        try:
            self.entered_number = int(new_text)
            return True
        except ValueError:
            return False

    def Main(self,params):
        #Close the window
        self.master.withdraw()
        vehsIgnored = self.vehsIgnored.get()
        satFlowHead = self.satFlowHead.get()
        minValidCycles = self.minValidCycles.get()
        minVehsPerMeas = self.minVehsPerMeas.get()

        csv_output_fp = r'Output_SatFlows.csv'
        dis_input_fp = self.fp_dis.get()
        #--------------------------------------------------------------------------------------------------------------#
        satFlow_data = self.GetDisData(dis_input_fp)

        #Keep the original Data in memory
        wt = satFlow_data
        #Remove the rows of the data where crossVehIdx <= VehsIgnored
        wt = self.RemoveIgnoredVehs(wt,vehsIgnored)
        #Return a dictionary with the total valid crossing time of vehicles for each Signal Head and Simulation
        satFlows = self.CalculateSatFlows(wt,satFlowHead,minValidCycles,minVehsPerMeas)
        #For the valid Signal Heads we can calculate the average of Saturation Flows
        #For the rest we print the list
        avgSatFlows = self.CalculateAvgSatFlows(satFlows)
        self.Output(csv_output_fp,avgSatFlows)
        completion_msg = 'End of execution - Results in {}'.format(csv_output_fp)
        print(completion_msg )
        messagebox.showinfo('Result', completion_msg )

        self.master.destroy()

    def GetDisData(self,dis_input_fp):
        #Function to read the dis files generated by Vissim
        #The output is a table with columns['sigHead','SimRun','cycleStart','timeGap']
        import glob

        sigHeads = []
        Simulation_Runs = []
        cycleStarts = []
        crossVehIdxs = []
        timeGaps = []

        path = dis_input_fp + r'\*.dis'
        for dis_file in glob.glob(path):
            with open(dis_file) as f:
                simRun = int(dis_file.rsplit('_',1)[1].replace('.dis',''))
                lines = f.readlines()
                for line in lines:
                    line = line.strip()
                    if line.startswith('Discharge at'):
                        sigHead = int(line.rsplit(' ', 1)[1])
                    else:
                        first_val = line.split(' ',1)[0]
                        #Try to convert value in float
                        try:
                            cycleStart = float(first_val)

                            #remove the consecutive spaces from the string
                            line_tmp = " ".join(line.split())

                            #First remove everything after the parenthesis
                            timeGaps_inLine = line_tmp.split('(')[0].strip()
                            timeGaps_inLine = timeGaps_inLine.split(' ')
                            del timeGaps_inLine[0]

                            crossVehIdx = 1
                            for timeGap in timeGaps_inLine:
                                sigHeads.append(sigHead)
                                Simulation_Runs.append(simRun)
                                cycleStarts.append(cycleStart)
                                crossVehIdxs.append(crossVehIdx)
                                timeGaps.append(timeGap)
                                crossVehIdx += 1
                        except:
                            continue
        #Create the numpy array
        #np_array = np.array(sigHeads,Simulation_Runs,cycleStarts,crossVehIdxs,timeGaps)
        wtype=np.dtype([('sigHead','i4'),('SimRun','i4'),
                        ('cycleStart','f4'),('crossVehIdx','f4'),
                        ('timeGap','f4')])

        wt=np.empty(len(sigHeads),dtype=wtype)
        wt['sigHead'] = sigHeads
        wt['SimRun'] = Simulation_Runs
        wt['cycleStart'] = cycleStarts
        wt['crossVehIdx'] = crossVehIdxs
        wt['timeGap'] = timeGaps

        wt.sort(order='sigHead')
        return wt

    def RemoveIgnoredVehs(self,inp_tbl,vehsIgnored):
        inds_to_del = np.where(inp_tbl['crossVehIdx'] <= vehsIgnored)[0]
        out_tbl = np.delete(inp_tbl,inds_to_del,axis=0)
        return out_tbl

    def CalculateSatFlows(self,wt,satFlowHead,minValidCycles,minVehsPerMeas):
        """ Return a dictionary with keys [Signal_Head_Simulation_Run]
            and values the total time of vehicles crossing (number of subsequent vehicles with
            acceptable timeGap, greater than minVehsPerMeas) and their count
        """
        satFlows = {}
        for signalHead in np.unique(wt['sigHead']):
            #dictionary to store the sat flows of the signal head for each iteration
            d = {}
            for simRun in np.unique(wt['SimRun']):
                #Find the unique values of cycleStarts for this signalHead and simRun
                sliced_tbl = wt[(wt['sigHead'] == signalHead) & (wt['SimRun'] == simRun)]

                unique_cycleStarts = np.unique(sliced_tbl['cycleStart'])

                if len(unique_cycleStarts) >= minValidCycles:
                    for cycleStart in unique_cycleStarts:
                        #Indices of the vehicles with low (acceptable) relative gap
                        idx_validGaps = np.where((wt['sigHead'] == signalHead) &
                                                 (wt['SimRun'] == simRun) &
                                                 (wt['cycleStart'] == cycleStart) &
                                                 (wt['timeGap'] <= satFlowHead))[0]

                        #Check if there is a sequence of acceptable indices with more element than the limit (minVehsPerMeas)

                        #Group the array in groups of consecutive integers
                        for k, g in groupby(enumerate(idx_validGaps), lambda ix: ix[0]-ix[1]):
                            idxs = map(itemgetter(1), g)

                            if len(idxs) >= minVehsPerMeas:
                                #We have found a sequence of vehicles with more elements than minVehsPerMeas so we can move on
                                crossTime = 0
                                for idx in idxs:
                                    #Add the sum to a dictionary with the number of vehicles that were counted
                                    crossTime += wt['timeGap'][idx]

                                sf = round((3600 * len(idxs)/crossTime),2)
                                d.setdefault(simRun,[]).append(sf)
                                break

                satFlows[signalHead] = d
        return satFlows


    def CalculateAvgSatFlows(self,satFlows):
        #Create a dictionary with the average sat flows for each Signal Head and each iteration
        avgSatFlows = {}

        for signalHead, simRun_SatFlows in satFlows.iteritems():
            #dictionary with the avg sat flows for each iteration
            d = {}
            for simRun, satFlows in simRun_SatFlows.iteritems():
                if len(satFlows) > 0:
                    avg = sum(satFlows)/len(satFlows)
                else:
                    avg=-1

                d[simRun] = round(avg,2)

            avgSatFlows[signalHead] = d

            #Calculate also the avg of all simRuns
            tot = 0
            for k, v in d.iteritems():
                tot += v
            if len(d) > 0:
                d['AllSimRuns'] = round(tot/len(d),2)
            else:
                d['AllSimRuns'] = -1

        return avgSatFlows

    def Output(self,filepath_out,results):
        import csv
        import os
        import collections as cl

        directory = os.path.dirname(filepath_out)

        if not os.path.exists(directory):
            os.makedirs(directory)

        with open(filepath_out,'wb') as f:
            w = csv.writer(f)
            w.writerow(['Signal_Head','Avg_Saturation_Flow'])
            results_ordered = cl.OrderedDict(sorted(results.items()))
            for k, v in results_ordered.iteritems():
                dout = v
                avgAllSimRuns = dout.pop('AllSimRuns',-1)
                l = [[k],[avgAllSimRuns],dout.values()]
                #flatten the list
                row = [i for slst in l for i in slst]
                w.writerow(row)

if __name__ == "__main__":
    root = Tk()
    my_gui = SatFlowsCalc(root)
root.mainloop()
