# Mech3410-P1-Graphs
#   Cal Wing - Aug 2024
#

import os, time
import numpy as np
import pandas as pd

from tqdm import tqdm
from numpy import sqrt

from makeGraph import makeGraph, pltKeyClose, UQ_COLOURS as UQC, uq_colour_cycler_factory as uqccf # Custom Graphing Lib

# Override Sin & Cos to use & return degrees
#def sin(angle): return np.sin(np.deg2rad(angle))
#def cos(angle): return np.cos(np.deg2rad(angle))

# Make sure the relevant folders folder exists
#folders = ["./images", "./tmp", "./data"]
folders = ["./images", './images/pressure', './images/cp']
for folder in folders:
    if not os.path.isdir(folder): os.mkdir(folder)


INCH_TO_M = 0.0254
GRAVITY = 9.81 #m/s^2

CHORD_LEN = 90 #mm
PITOT_PLACEMENT = np.array((0,4,8,16,25,34,43,53,61,70,5,9,17,25,34,43,53,61,70)) # mm from base of chord
PITOT_PLACEMENT_CHORD_NORM = PITOT_PLACEMENT / CHORD_LEN

print("="*15, "Loading Data", "="*15)
data_508rpm = {
    "rpm": 508,
    "airSpeed": 10,
    "data": pd.read_excel('.\\data\\508 RPM Results.xlsx', sheet_name=None, header=None),
    "AoA": (0,1,2,3,4,5,6,7,8,9)
}

data_1000rpm = {
    "rpm": 1000,
    "airSpeed": 20.40408122,
    "data": pd.read_excel('.\\data\\1000 RPM Results.xlsx', sheet_name=None, header=None),
    "AoA": (1,2,4,6,8,10,12,14,16)
}

print("="*15, "Loaded Data", "="*15)

def make_pressure_graph(sheet1, aoa, rpm, air_speed, doGraph=True):
    water_density = sheet1.iloc[0, 2] # kg/m^3
    air_density = sheet1.iloc[1, 2] # kg/m^3
    atm_presure_inch = sheet1.iloc[24, 12] # inch

    pitot_height_inch = sheet1.iloc[4:23, 11]
    pitot_height_m = (pitot_height_inch - atm_presure_inch)*INCH_TO_M

    pressure = water_density * GRAVITY * pitot_height_m

    #print(pressure)
    #print(pressure.min(), pressure.max())

    # Extrapolate Aerofoil Tip
    p1 = (PITOT_PLACEMENT[9], pressure.iloc[9])
    p2 = (PITOT_PLACEMENT[18], pressure.iloc[18])
    pN = (0.9*CHORD_LEN, p2[1] + (p1[1]-p2[1])/2)
    #print(aoa, rpm, p1[1], p2[1], (p1[1]-p2[1])/2, pN[1], pN[1] > p1[1])


    graph = {
            "title":  f"Pressure vs Pitot Placement along Chord\nfor a Clark Y 14% Aerofoil at:\nα = {int(aoa):d}° at {rpm:d} RPM ({air_speed:.1f}m/s)", 
            "windowTitle": f"Pressure along Clark Y 14% Airfoil - Alpha {int(aoa):d} Degrees - {rpm:d}rpm - {air_speed:.1f}m_s",
            "xLabel": "Pitot Placement [mm]", 
            "yLabel": "Pressure [Pa]",
            "grid": True,
            "yLim": (pressure.min()-10, pressure.max()+10),
            "plots": [
                # Draw Lines
                {"x": PITOT_PLACEMENT[:10], "y":pressure[:10], "colour": UQC["purple"]},
                {"x": PITOT_PLACEMENT[10:], "y":pressure[10:], "colour": UQC["purple"]},
                {"x": [PITOT_PLACEMENT[0], PITOT_PLACEMENT[10]],  "y": [pressure.iloc[0], pressure.iloc[10]], "colour": UQC["purple"]},
                #{"x": [PITOT_PLACEMENT[9], PITOT_PLACEMENT[18]], "y": [pressure.iloc[9], pressure.iloc[18]], "colour": UQC["purple"]},
                
                # Draw Extrapolated Airfoil Tip
                {"x": [PITOT_PLACEMENT[9], pN[0]], "y": [pressure.iloc[9], pN[1]], "colour": UQC["aqua"], "label": "Extrapolated Airfoil Tip"},
                {"x": [PITOT_PLACEMENT[9], pN[0]], "y": [pressure.iloc[18], pN[1]], "colour": UQC["aqua"]},

                # Draw Colour Shading
                {"x": PITOT_PLACEMENT[:10], "y":pressure[:10], "colour": UQC["purple"], "alpha":0.2, "type":"fill"},
                {"x": PITOT_PLACEMENT[10:], "y":pressure[10:], "colour": "w", "type":"fill"},
                {"x":[PITOT_PLACEMENT[0], PITOT_PLACEMENT[10]], "y":[pressure.iloc[0], pressure.iloc[10]], "colour": "w", "type":"fill"},
                {"x": [PITOT_PLACEMENT[9], pN[0]], "y": [pressure.iloc[9], pN[1]], "colour": UQC["aqua"], "type":"fill", "alpha":0.2},
                {"x":[PITOT_PLACEMENT[9], pN[0]], "y": [pressure.iloc[18], pN[1]], "colour": "w", "type":"fill"},
                
                # Draw Points & text
                {"x":PITOT_PLACEMENT, "y":pressure, "label":"Pressure Data", "type":"scatter", "args":{"zorder":2}},
                {"label":[str(i) for i in range(1, len(PITOT_PLACEMENT)+1)], "x":PITOT_PLACEMENT, "y":pressure, "type":"annotate"},
                {"type":"text", "x": 0.98, "y": 0.02, "text": f"Min: {pressure.min():.3f} Pa\nMax: {pressure.max():.3f} Pa", "align": ('bottom', 'right')}
            ]
        }
    
    if doGraph:
        makeGraph(graph, False, figSavePath=f'./images/pressure/{{0}}.png', closeFig=True)
    
    return pressure, graph, (water_density, air_density, atm_presure_inch, pitot_height_inch), (p1, p2, pN)


def make_rpm_graph():
    air_density = data_508rpm["data"]["0 AoA"].iloc[1, 2] # kg/m^3
    rpm = np.concat((np.array([0]), np.array(data_508rpm["data"]["0 AoA"].iloc[4:14, 2])))
    pressure = np.concat((np.array([0]), np.array(data_508rpm["data"]["0 AoA"].iloc[4:14, 4], dtype=np.float64)))

    airSpeed = sqrt(-pressure * 2 / air_density)
    
    x = np.linspace(rpm[0], rpm[-1], 1000)
    y = 0.0212*x - 0.7751

    graph = {
                "title":  f"RPM vs Velocity for Wind Tunnel 3", 
                "xLabel": "Fan RPM", 
                "yLabel": "Airspeed [m/s]",
                "grid": True,
                "plots": [
                    {"x": x, "y":y, "label":"Linear Estimate", "args":{"linestyle":"--"}, "colour": UQC["neutral"]},

                    {"type": "scatter", "x":rpm, "y": airSpeed, "label": "Measured Velocity", "colour":UQC["purple"], "args":{"zorder":2}},
                    {"type": "point", "x":508, "y": 10, "label": "Selected RPM (508)", "colour":UQC["red"], "zorder":3},

                    {"type":"text", "x": 0.98, "y": 0.02, "text": f"Velocity at 508 RPM: {10:.3f} m/s\nVelocity at 1000 RPM: {airSpeed[-1]:.3f} m/s", "align": ('bottom', 'right')},
                    {"type":"text", "x": 0.80, "y": 0.85, "text": f"$y = 0.0212 \\cdot x - 0.7751$", "align": ('bottom', 'right'), "facecolour": UQC["neutral"]}
                ]
            }

    makeGraph(graph, False, figSavePath='./images/{0}.png')

if __name__ == '__main__':
    print("Generating RPM Graph")
    #make_rpm_graph()
    print("Generated")

    print("Loading Data & Generating Pressure Graphs")
    data = {}
    for raw_data in tqdm((data_508rpm, data_1000rpm), position=0):
        aoa_data = {}
        for aoa in tqdm(raw_data["AoA"], position=1):
            sheet = raw_data["data"][f"{aoa} AoA"]

            aoa_data[aoa] = make_pressure_graph(sheet, aoa, raw_data["rpm"], raw_data["airSpeed"], True)

        data[raw_data["rpm"]] = aoa_data

        if True:
            # All
            graph = {
                "title":  f"Pressure vs Pitot Placement along Chord\nfor a Clark Y 14% Aerofoil at:\n{raw_data["rpm"]:d} RPM ({raw_data["airSpeed"]:.1f}m/s)", 
                "windowTitle": f"All Pressure along Clark Y 14% Airfoil - {raw_data["rpm"]:d}rpm - {raw_data["airSpeed"]:.1f}m_s",
                "xLabel": "Pitot Placement [mm]", 
                "yLabel": "Pressure [Pa]",
                "grid": True,
                "ledgLoc": 1,
                "plots": []
            }

            colour_cycle = uqccf()
            for aoa, c in zip(aoa_data, colour_cycle):
                this_data = aoa_data[aoa]
                pressure = this_data[0]

                # Extrapolate Aerofoil Tip
                p1 = (PITOT_PLACEMENT[9], pressure.iloc[9])
                p2 = (PITOT_PLACEMENT[18], pressure.iloc[18])
                pN = (0.9*CHORD_LEN, p2[1] + (p1[1]-p2[1])/2)

                plts = (
                    {"x": PITOT_PLACEMENT[:10], "y":pressure[:10], "colour": c["color"]},
                    {"x": PITOT_PLACEMENT[10:], "y":pressure[10:], "colour": c["color"]},
                    
                    {"x": [PITOT_PLACEMENT[0], PITOT_PLACEMENT[10]],  "y": [pressure.iloc[0], pressure.iloc[10]], "colour": c["color"]},
                    
                    #{"x": [PITOT_PLACEMENT[9], PITOT_PLACEMENT[18]], "y": [pressure.iloc[9], pressure.iloc[18]], "colour": c["color"]},
                    
                    {"x": [PITOT_PLACEMENT[9], pN[0]], "y": [pressure.iloc[9], pN[1]], "colour": c["color"].lighten(0.4)},
                    {"x": [PITOT_PLACEMENT[9], pN[0]], "y": [pressure.iloc[18], pN[1]], "colour": c["color"].lighten(0.4)},

                    {"x":PITOT_PLACEMENT, "y":pressure, "label":f"α = {int(aoa):d}°", "type":"scatter", "colour": c["color"], "args":{"zorder":2}},
                )
                for plt in plts:
                    graph["plots"].append(plt)
            
            if raw_data["rpm"] == 508:
                graph['xLim'] = (
                    min(PITOT_PLACEMENT) - 5,
                    max(PITOT_PLACEMENT) + 20
                )
                graph["figSize"] = (8, 6)


            makeGraph(graph, False, figSavePath="./images/pressure/__{0}.png")

            graph["plots"] = []
            graph["windowTitle"] = f"Pressure along Clark Y 14% Airfoil - {raw_data["rpm"]:d}rpm - {raw_data["airSpeed"]:.1f}m_s"

            colour_cycle = uqccf()
            for aoa, c in list(zip(aoa_data, colour_cycle))[1::2]:
                this_data = aoa_data[aoa]
                pressure = this_data[0]

                # Extrapolate Aerofoil Tip
                p1 = (PITOT_PLACEMENT[9], pressure.iloc[9])
                p2 = (PITOT_PLACEMENT[18], pressure.iloc[18])
                pN = (0.9*CHORD_LEN, p2[1] + (p1[1]-p2[1])/2)

                plts = (
                    {"x": PITOT_PLACEMENT[:10], "y":pressure[:10], "colour": c["color"]},
                    {"x": PITOT_PLACEMENT[10:], "y":pressure[10:], "colour": c["color"]},
                    
                    {"x": [PITOT_PLACEMENT[0], PITOT_PLACEMENT[10]],  "y": [pressure.iloc[0], pressure.iloc[10]], "colour": c["color"]},
                    #{"x": [PITOT_PLACEMENT[9], PITOT_PLACEMENT[18]], "y": [pressure.iloc[9], pressure.iloc[18]], "colour": c["color"]},

                    {"x": [PITOT_PLACEMENT[9], pN[0]], "y": [pressure.iloc[9], pN[1]], "colour": c["color"].lighten(0.4)},
                    {"x": [PITOT_PLACEMENT[9], pN[0]], "y": [pressure.iloc[18], pN[1]], "colour": c["color"].lighten(0.4)},
                    
                    {"x":PITOT_PLACEMENT, "y":pressure, "label":f"α = {int(aoa):d}°", "type":"scatter", "colour": c["color"]},
                )
                for plt in plts:
                    graph["plots"].append(plt)
            
            makeGraph(graph, False, figSavePath="./images/pressure/__{0}.png")

    print("Data Loaded")