import csv
import os
import matplotlib.pylab as plt

AirlineDict = {"BA": "British Airways"}
ModelDict = {"A319-100": "A319", "A320-200": "A320", "A320neo": "A32N",
             "A321-200": "A321", "A321neo": "A21N"}
VariableDict1 = {"acc": ["Acc. (m/s2)", "Acceleration"],
                "bf": ["Breakaway (N)", "Breakaway Force"],
                "atf": ["Acc. Trac. (N)", "Acceleration Traction Force"],
                "ad": ["Aerod. Drag (N)", "Aerodynamic Drag"], 
                "tt": ["Total Traction (N)", "Total Traction"], 
                "T": ["Torque (Nm)", "Torque"], 
                "P": ["Power (kW)", "Power"]}
VariableDict2 = {"Time": "sec", "Energy": "kWh"}
#Data obtained previously in Airport Operation Frequency Studies (in mins)
AirportOperationFreq = {"BA": [17.47442, 6.8660247],
                        "W6": [11.4252, 5.8024524],
                        "EC": [14.12387, 5.8095266]}

# ---------------------- PLOT SINGLE FLIGHT CASE FILE ---------------------- #

def SingleFlightGraph():
    
    def PlotGraph(FlightFile, case, variable):
        
        try:
            
            #Read FileList and read the file names
            FileList = []
            FileString = FlightFile.split("_")
            FlightNumber = FileString[0]
            
            if case == "Taxi In":
                with open ('X(IN)_'+FlightFile,'r') as file:
                    tempFileList = csv.reader(file)
                    for row in tempFileList:
                        FileList.append(row)
            else:
                with open ('X(OUT)_'+FlightFile,'r') as file:
                    tempFileList = csv.reader(file)
                    for row in tempFileList:
                        FileList.append(row)
            
            #Extract X-coordinate (Taxi Time) and Y1-coordinate (Ground Speed)
            XList = []
            GSpeed = []
            for i in range (1,len(FileList)):
                XList.append(float(FileList[i][7]))
                GSpeed.append(float(FileList[i][8]))
            
            #Extract Y2-coordinate (User Input)
            YList = []
            Column = 0
            for j in range (10,17):
                if VariableDict1[Variable][0] == FileList[0][j]:
                    Column = j
            for i in range (1,len(FileList)):
                YList.append(float(FileList[i][Column]))
        
            #Plot Graph
            fig,ax1 = plt.subplots(figsize=(12, 8))
            ax2 = ax1.twinx()
        
            ax1.plot(XList, GSpeed, color = "#B6AD7C")
            ax2.plot(XList, YList, color = "#215868")
            
            ax1.set_xlabel("Taxi Time (sec)")
            ax1.set_ylabel("Ground Speed (m/s)", color = "#B6AD7C")
        
            ax2.set_ylabel(VariableDict1[Variable][0], color = "#215868")
            
            fig.suptitle(FlightNumber + " " + case + " Ground Speed and " + 
                         VariableDict1[Variable][1] + " Relation", 
                         fontsize=20)
            fig.autofmt_xdate()
            fig = plt.gcf()
            plt.show()
        
        except IOError:
            
            print("\nWarning: Cannot locate file.")
    
    #Input User Demands
    FlightFile = input ("\nWhich flight file do you want to generate a graph"\
                        ": ")
    
    Case = input ("\nTaxi In or Taxi Out: ")
    while Case.lower() != "taxi in" and Case.lower() != "taxi out":
        Case = input ("Input error. Select Again: ")
    
    print ("\nWhat variable do you want to add?")
    Variable = input ("Please choose from Acceleration (acc), Breakaway Forc"\
                      "e (bf), Acceleration Traction Force (atf), Aerodynami"\
                      "c Drag (ad), Total Traction (tt), Torque (T) and Powe"\
                      "r (P): ")
    while Variable not in VariableDict1:
        Variable = input ("Input error. Select Again: ")
    
    #Execute Program
    if Case.lower() == "taxi in":
        PlotGraph(FlightFile,"Taxi In", Variable)
    else:
        PlotGraph(FlightFile,"Taxi Out", Variable)
    

# ---------------------- PLOT TAXI TIME VS TAXI ENERGY --------------------- #

def TaxiTimeVersusTaxiEnergy():
    
    def PlotGraph(ICAO, Model, Case, Elim):
        
        #Read FileList and read the file names
        FileList = []
        EnergyList = []
        TaxiTimeList = []
        ElimTaxiTimeList = []
        ElimEnergyList = []
        CorrAvgTaxiTimeList = []
        CorrAvgEnergyList = []
        
        if Case == "Taxi In":
            with open ('#' + ICAO + ' ' + Model +
                       ' Taxi In Flight List.csv','r') as file:
                tempFileList = csv.reader(file)
                for row in tempFileList:
                    FileList.append(row)
        else:
            with open ('#' + ICAO + ' ' + Model +
                       ' Taxi Out Flight List.csv','r') as file:
                tempFileList = csv.reader(file)
                for row in tempFileList:
                    FileList.append(row)
        
        for i in range(1, len(FileList)-2):
            if len(FileList[i]) < 10:
                TaxiTimeList.append(float(FileList[i][7]))
                EnergyList.append(float(FileList[i][8]))
            else:
                ElimTaxiTimeList.append(float(FileList[i][7]))
                ElimEnergyList.append(float(FileList[i][8]))
        
        CorrAvgTaxiTimeList.append(float(FileList[len(FileList)-1][7]))
        CorrAvgEnergyList.append(float(FileList[len(FileList)-1][8]))
    
        plt.figure ("Time VS Energy", figsize=(12, 8))
        plt.clf()
        plt.title(AirlineDict[ICAO] + " " + Model + " " + Case + 
                  " - Taxi Time Versus Taxi Energy Requirement Graph")
        plt.plot(EnergyList, TaxiTimeList, color = "#B6AD7C", marker = ".",
                 linestyle = "", label = "Useful Flights")
        if Elim.lower() == "yes":
            plt.plot(ElimEnergyList, ElimTaxiTimeList, color = "#B30000", 
                     marker = ".", linestyle = "", 
                     label = "Eliminated Flights")
        plt.plot(CorrAvgEnergyList, CorrAvgTaxiTimeList, color = "#215868",
                 marker = "X", markersize = 8, linestyle  = "",
                 label = "Corrected Average")
        if Case == "Taxi In":
            plt.axhline(y = float(AirportOperationFreq[ICAO][1]*60),
                        color = "#C49108",
                        label = "Airport Opt. Freq. Studies Avg.")
            plt.axvline(x = 8.81, color = "#006666",
                        label = "2020 Ariadna F. M.")
            plt.axvline(x = 9.98, color = "#006666", linestyle = "--", 
                        label = "2020 Manousos L.")
            #plt.axvline(x = 31.687, color = "#006666",  linestyle = "-.", 
            #            label = "2020 Thibault S.")
        else:
            plt.axhline(y = float(AirportOperationFreq[ICAO][0]*60),
                        color = "#C49108", 
                        label = "Airport Opt. Freq. Studies Avg.")
            plt.axvline(x = 31.17, color = "#006666",
                        label = "2020 Ariadna F. M.")
            plt.axvline(x = 28.24, color = "#006666",  linestyle = "--",
                        label = "2020 Manousos L.")
            plt.axvline(x = 22.3389, color = "#006666",  linestyle = "-.", 
                        label = "2020 Thibault S.")
        plt.xlabel("Taxi Energy Required (kWh)")
        plt.ylabel("Taxi Time (sec)")
        plt.legend()
        fig = plt.gcf()
        plt.show()
        
        save = input("\nSave the plot?\nSelect 'Yes' or 'No': ")
        if save.lower() == "yes" or save.lower() == "save":
            fig.savefig("./#Plots/" + AirlineDict[ICAO] + " " + Model + " " +
                        Case + " - Taxi Time Versus Taxi Energy Requirement"\
                        " Graph.png", dpi = 300)
                    
    #Input User Demands
    print ("\nWhich airline are you looking into?")
    Question = "Please choose from the following: "
    for i in AirlineDict.keys():
        Question += i + "  "
    ICAO = input (Question + "\nYour selection: ")
    while ICAO.upper() not in AirlineDict:
        ICAO = input ("Input error. Select Again: ")
    
    print ("\nWhat model do you want to show?")
    Question = "Please choose from the following: "
    for i in ModelDict.keys():
        Question += i + "  "
    Model = input (Question + "\nYour selection: ")
    while Model not in ModelDict:
        Model = input ("Input error. Select Again: ")
    
    Case = input ("\nTaxi In or Taxi Out: ")
    while Case.lower() != "taxi in" and Case.lower() != "taxi out":
        Case = input ("Input error. Select Again: ")
        
    Elim = input ("\nShow eliminated cases: ")
    while Elim.lower() != "yes" and Elim.lower() != "no":
        Elim = input ("Input error. Select Again: ")
    
    #Execute Program
    if Case.lower() == "taxi in":
        PlotGraph(ICAO.upper(), Model, "Taxi In", Elim.lower())
    else:
        PlotGraph(ICAO.upper(), Model, "Taxi Out", Elim.lower())

# ---------------------- PLOT TAXI TIME OR TAXI ENERGY --------------------- #
# -------------------------------- BOX GRAPH ------------------------------- #

def TaxiInfoBoxGraph():
    
    def PlotGraph(ICAO, Variable, Case):
        
        TheDict = {}
        
        for Model in ModelDict.keys():
            try:
                with open ('#' + ICAO + ' ' + Model + ' ' +\
                           Case + ' Flight List.csv','r') as file:
                    tempFileList = csv.reader(file)
                    count = 0
                    for row in tempFileList:
                        count += 1
                        if count != 1:
                            if Variable == "Time":
                                if ModelDict[Model] not in TheDict and\
                                    len(row) < 10:
                                    TheDict[ModelDict[Model]] = [float(row[7])]
                                elif ModelDict[Model] in TheDict and\
                                    len(row) < 10:
                                    TheDict[ModelDict[Model]].append(float(row[7]))
                            elif Variable == "Energy":
                                if ModelDict[Model] not in TheDict and\
                                    len(row) < 10:
                                    TheDict[ModelDict[Model]] = [float(row[8])]
                                elif ModelDict[Model] in TheDict and\
                                    len(row) < 10:
                                    TheDict[ModelDict[Model]].append(float(row[8]))
            except:
                print("Cannot find " + Model + " files.")
        
        TheList = []
        NameList = []
        NameCountList = []
        NameCount = 0
        for Model in TheDict.keys():
            NameCount += 1
            NameList.append(Model)
            NameCountList.append(NameCount)
            TheList.append(TheDict[Model])
        
        plt.figure ("Single Info", figsize=(12, 8))
        plt.clf()
        plt.title(AirlineDict[ICAO] + " " + Case +\
                  " Taxi " + Variable + " Box Graph")
        plt.boxplot(TheList)
        plt.xticks(NameCountList, NameList)
        if Variable == "Time":
            plt.ylabel("Taxi Time (secs)")
        elif Variable == "Energy":
            plt.ylabel("Taxi Energy (kWh)")
        plt.xlabel("Aircraft Models")
        fig = plt.gcf()
        plt.show()
        
        save = input("\nSave the plot?\nSelect 'Yes' or 'No': ")
        if save.lower() == "yes" or save.lower() == "save":
            fig.savefig("./#Plots/" + AirlineDict[ICAO] + " " + Case +\
                      " Taxi " + Variable + " Box Graph.png", dpi = 300)
    
    #Input User Demands
    print ("\nWhich airline are you looking into?")
    Question = "Please choose from the following: "
    for i in AirlineDict.keys():
        Question += i + "  "
    ICAO = input (Question + "\nYour selection: ")
    while ICAO.upper() not in AirlineDict:
        ICAO = input ("Input error. Select Again: ")
        
    Case = input ("\nTaxi In or Taxi Out: ")
    if Case.lower() == "taxi in" or Case.lower() == "in":
        Case = "Taxi In"
    elif Case.lower() == "taxi out" or Case.lower() == "out":
        Case = "Taxi Out"
    while Case != "Taxi In" and Case != "Taxi Out":
        Case = input ("Input error. Select Again: ")
        if Case.lower() == "taxi in" or Case.lower() == "in":
            Case = "Taxi In"
        elif Case.lower() == "taxi out" or Case.lower() == "out":
            Case = "Taxi Out"
          
    Variable = input ("\nTaxi Time or Taxi Energy: ")
    while Variable.lower() != "taxi time"\
        and Variable.lower() != "taxi energy"\
        and Variable.lower() != "time"\
        and Variable.lower() != "energy":
        Variable = input ("Input error. Select Again: ")
        
    #Execute Program
    if Variable.lower() == "taxi time" or Variable.lower() == "time":
        PlotGraph(ICAO.upper(), "Time", Case)
    else:
        PlotGraph(ICAO.upper(), "Energy", Case)

# ---------------------- PLOT TAXI TIME OR TAXI ENERGY --------------------- #
# -------------------------------- HISTOGRAM ------------------------------- #

def Histogram():
    
    def PlotGraph(ICAO, Model, Case, Variable, Elim):
        
        #Read FileList and read the file names
        FileList = []
        VariableList = []
        
        if Case == "Taxi In":
            with open ('#' + ICAO + ' ' + Model +\
                       ' Taxi In Flight List.csv','r') as file:
                tempFileList = csv.reader(file)
                for row in tempFileList:
                    FileList.append(row)
        else:
            with open ('#' + ICAO + ' ' + Model +\
                       ' Taxi Out Flight List.csv','r') as file:
                tempFileList = csv.reader(file)
                for row in tempFileList:
                    FileList.append(row)
        
        for i in range(1, len(FileList)-2):
            if Elim == "no":
                if len(FileList[i]) < 10:
                    if Variable == "Time":
                        VariableList.append(float(FileList[i][7]))
                    elif Variable == "Energy":
                        VariableList.append(float(FileList[i][8]))
            elif Elim == "yes":
                if Variable == "Time":
                    VariableList.append(float(FileList[i][7]))
                elif Variable == "Energy":
                    VariableList.append(float(FileList[i][8]))
    
        plt.figure ("Histogram", figsize=(12, 8))
        plt.clf()
        plt.title(AirlineDict[ICAO] + " " + Model + " " + Case + " - Taxi " + 
                  Variable + " Histogram")
        plt.hist(VariableList, bins = 50, color = "#215868")
        
        plt.xlabel("Taxi " + Variable + " (" + VariableDict2[Variable] + ")")
        plt.ylabel("Frequency")
        fig = plt.gcf()
        plt.show()
        
        save = input("\nSave the plot?\nSelect 'Yes' or 'No': ")
        if save.lower() == "yes" or save.lower() == "save":
            fig.savefig("./#Plots/" + AirlineDict[ICAO] + " " + Model + " " +
                        Case + " - Taxi " + Variable + " Histogram.png",
                        dpi = 300)
                    
    #Input User Demands
    print ("\nWhich airline are you looking into?")
    Question = "Please choose from the following: "
    for i in AirlineDict.keys():
        Question += i + "  "
    ICAO = input (Question + "\nYour selection: ")
    while ICAO.upper() not in AirlineDict:
        ICAO = input ("Input error. Select Again: ")
    
    print ("\nWhat model do you want to show?")
    Question = "Please choose from the following: "
    for i in ModelDict.keys():
        Question += i + "  "
    Model = input (Question + "\nYour selection: ")
    while Model not in ModelDict:
        Model = input ("Input error. Select Again: ")
    
    Case = input ("\nTaxi In or Taxi Out: ")
    while Case.lower() != "taxi in" and Case.lower() != "taxi out":
        Case = input ("Input error. Select Again: ")
        
    Variable = input ("\nTaxi Time or Taxi Energy: ")
    while Variable.lower() != "taxi time"\
        and Variable.lower() != "taxi energy"\
        and Variable.lower() != "time"\
        and Variable.lower() != "energy":
        Variable = input ("Input error. Select Again: ")
    if Variable.lower() == "taxi time" or Variable.lower() == "time":
        Variable = "Time"
    else:
        Variable = "Energy"
        
    Elim = input ("\nShow eliminated cases: ")
    while Elim.lower() != "yes" and Elim.lower() != "no":
        Elim = input ("Input error. Select Again: ")
    
    #Execute Program
    if Case.lower() == "taxi in":
        PlotGraph(ICAO.upper(), Model, "Taxi In", Variable, Elim.lower())
    else:
        PlotGraph(ICAO.upper(), Model, "Taxi Out", Variable, Elim.lower())
        
# ---------------------------- PLOT TAXI SUMMARY --------------------------- #

def Summary():
    
    def PlotGraph(ICAO, Variable):
        
        FileList = []
        with open ('#' + ICAO + ' Final List.csv','r') as file:
            tempFileList = csv.reader(file)
            for row in tempFileList:
                FileList.append(row)
        
        ModelList = []
        InList = []
        OutList = []
        
        if Variable == "Time":
            for i in range (1, len(FileList)):
                if FileList[i][2] not in ModelList:
                    ModelList.append(FileList[i][2])
                if FileList[i][3] == "Taxi Out":
                    OutList.append(float(FileList[i][5]))
                elif FileList[i][3] == "Taxi In":
                    InList.append(float(FileList[i][5]))
        else:
            for i in range (1, len(FileList)):
                if FileList[i][2] not in ModelList:
                    ModelList.append(FileList[i][2])
                if FileList[i][3] == "Taxi Out":
                    OutList.append(float(FileList[i][6]))
                elif FileList[i][3] == "Taxi In":
                    InList.append(float(FileList[i][6]))
                    
        plt.figure ("Summary", figsize=(12, 8))
        plt.clf()
        plt.title(AirlineDict[ICAO] + " Taxi " + 
                  Variable + " Summary")                      
        plt.bar(ModelList, InList, color = "#215868", label = "Taxi In")
        plt.bar(ModelList, OutList, bottom = InList, color = "#DDD9C3", 
                label = "Taxi Out")
        if Variable == "Time":
            plt.ylabel("Taxi Time (mins)")
        else:
            plt.ylabel("Taxi Energy (kWh)")
        plt.xlabel("Aircraft Models")
        plt.legend()
        fig = plt.gcf()
        plt.show()
        
        save = input("\nSave the plot?\nSelect 'Yes' or 'No': ")
        if save.lower() == "yes" or save.lower() == "save":
            fig.savefig("./#Plots/" + "#" + AirlineDict[ICAO] + " Taxi " + 
                      Variable + " Summary.png", dpi = 300)
    
    #Input User Demands
    print ("\nWhich airline are you looking into?")
    Question = "Please choose from the following: "
    for i in AirlineDict.keys():
        Question += i + "  "
    ICAO = input (Question + "\nYour selection: ")
    while ICAO.upper() not in AirlineDict:
        ICAO = input ("Input error. Select Again: ")
          
    Variable = input ("\nTaxi Time or Taxi Energy: ")
    while Variable.lower() != "taxi time"\
        and Variable.lower() != "taxi energy"\
        and Variable.lower() != "time"\
        and Variable.lower() != "energy":
        Variable = input ("Input error. Select Again: ")
        
    #Execute Program
    if Variable.lower() == "taxi time" or Variable.lower() == "time":
        PlotGraph(ICAO.upper(), "Time")
    else:
        PlotGraph(ICAO.upper(), "Energy")

# ------------------------- MASTER GRAPH SELECTION ------------------------- #

cont = True

while cont == True:

    print("\nThis is the graph plotting program for Chi Hin Pang's Electrica"\
          "l Taxiing System Research Project. Select your desired program ty"\
          "pe and enter the respective code: \n\n  A. Single flight file - P"\
          "arameter versus ground speed graph \n  B. Taxi time versus taxi e"\
          "nergy requirement scatter type graph \n  C. Taxi time or taxi ene"\
          "rgy box type graph \n  D. Taxi time or taxi energy histogram \n  "\
          "E. Taxi Summary \n")
    
    MasterChoice = input("Your selection: ")
    
    SelectionList = ["a", "b", "c", "d", "e"]
    while MasterChoice.lower() not in SelectionList:
        MasterChoice = input("Input error. Select Again: ")
    
    if MasterChoice.lower() == "a":
        print("\n--------------- THE FOLLOWING GENERATES THE SINGLE FLIGHT F"\
              "ILE DATA GRAPH ---------------")
        SingleFlightGraph()
        print("\n------------------------ SINGLE FLIGHT FILE DATA GENERATED "\
              "ABOVE ------------------------")
    elif MasterChoice.lower() == "b":
        print("\n------ THE FOLLOWING GENERATES THE TAXI TIME VERSUS TAXI EN"\
              "ERGY SCATTER TYPE GRAPH ------")
        TaxiTimeVersusTaxiEnergy()
        print("\n------------ TAXI TIME VERSUS TAXI ENERGY SCATTER TYPE GRAP"\
              "H GENERATED ABOVE ------------")
    elif MasterChoice.lower() == "c":
        print("\n---------- THE FOLLOWING GENERATES THE TAXI TIME OR TAXI EN"\
              "ERGY BOX TYPE GRAPH ----------")
        TaxiInfoBoxGraph()
        print("\n---------------- TAXI TIME OR TAXI ENERGY BOX TYPE GRAPH GE"\
              "NERATED ABOVE ----------------")
    elif MasterChoice.lower() == "d":
        print("\n------------ THE FOLLOWING GENERATES THE TAXI TIME OR TAXI "\
              "ENERGY HISTOGRAM -------------")
        Histogram()
        print("\n------------------ TAXI TIME OR TAXI ENERGY HISTOGRAM GENER"\
              "ATED ABOVE -------------------")
    elif MasterChoice.lower() == "e":
        print("\n----------------------- THE FOLLOWING GENERATES THE TAXI SU"\
              "MMARY ------------------------")
        Summary()
        print("\n----------------------------- TAXI SUMMARY GENERATED ABOVE "\
              "------------------------------")
        
    
    another = input("\nRun another plot?\nSelect 'Yes' or 'No': ")
    if another.lower() == "yes" or another.lower() == "run":
        cont = True
    else:
        cont = False

print("\n--------------------------- THANK YOU FOR USING THIS PROGRAM ------"\
      "----------------------\n----------------------------------- HAVE A NI"\
      "CE DAY! ------------------------------------")