# -*- coding: utf-8 -*-
"""
Created on Thu May 27 12:00:27 2021
Last updated on Tue Oct 19 16:00:23 2021

ENGG4552 Heat Shield Trajectory Code

Authors: Matthew Rickerby, Daniel Burgess, Clifford Asmussen, Nathan Holyoak,
         Daniel Kearney, Shanil Panchal, Tharang Addepalli

"""

# Import required modules
import numpy as np
import matplotlib.pyplot as plt

"""Initial Conditions"""
init_alt = 200E3            #[m]            Initial altitude
init_disp = 0               #[m]            Initial displacement
init_v = 7800               #[m/s]          Initial velocity

"""Constants"""
G = 6.67430E-11             #[Nm^2/kg^2]    Gravitational constant
Me = 5.97219E24             #[kg]           Mass of Earth
Re = 6371E3                 #[m]            Radius of Earth
y0 = 6620                   #[m]            Scale Height - AERO4800 Textbook
beta = 1/y0                 #               Scaling factor
surface_rho = 1.225         #[kg/m^3]       Surface air density
pica_rho = 1155.54          #[kg/m^3]       Density of PICA
char_rho = 280              #[kg/m^3]       Density of Char (from research)
ins_rho = 280               #[kg/m^3]       Carbon felt density
al_rho = 2710               #[kg/m^3]       Aluminium density
fo_rho = 96                 #[kg/m^3]       Internal foam density
mass_cube = 1.33            #[kg]           CubeSat mass
mass_struc = 2              #[kg]           Aluminium structure mass (both sheets included)
mass_bolts = 1              #[kg]           Assumed bolt mass
mass_other = 5              #[kg]           Assumed mass of parachute and other components

"""Vehicle Properties"""
radius = 0.14586            #[m]            Radius of capsule
ang = 87                    #[deg]          Angle prototype wraps to
scale = 1                   #               Steps per degree
step = ((ang * scale) + 1)  #               Steps in angles
angles = np.linspace(0, ang, step) #        Angle steps generation

noser = 0.07927             #[m]            Nose radius
Cd = 1.4                    #               Coefficient of drag (from cart 3d)
Cl = 0                      #               Coefficient of lift (ballistic entry)
S = 0.05106375              #[m^2]          Reference area

# Radiative properties
em = 0.9                    #               Emissivity (approximated as carbon emissivity)
sig = 5.67*(10**-8)         #[W/(m^2 k^4)]  Stefan-Boltzmann constant

# Specific Heat
cpc = 1925.56               #[J/kgK]        PICA (approximated as specific heat of carbon)
cpch = 4605.48              #[J/kgK]        Charr (from research)
cpin = 730                  #[J/kgK]        Carbon felt
cpfo = 1050                 #[J/kgK]        Insulative foam
cpal = 1050                  #[J/kgK]        Aluminium

# Thermal conductivity
k_pica = 0.577              #[W/mK]         PICA
k_char = 0.577              #[W/mK]         Charr (from research)
k_ins = 0.13                #[W/mK]         Carbon felt
k_foa = 0.05                #[W/mK]         Insulative foam

"""Thicknesses"""
pthc = 0.010                #[m]            PICA
sthc = 0.026                #[m]            Carbon felt
fthc = 0.0135               #[m]            Insulative foam
althc = 0.016               #[m]            Aluminium

#Open lists and append thickness around front heat shield
thic_p = []
thic_s = []
thic_f = fthc # Insulative foam
thic_al = althc
for i in range(step):
    thicccc = pthc # Constant PICA thickness
    thic_p.append(thicccc)

    # Calculate extra thickness in nose
    if i < 30.5: # In nose radius
        thicss = (-6.9921 * (10**-6) * i**2) + (1.9423*i*10**-6) + (0.032418 - sthc)
    else: # In flat part
        thicss = 0.00
    thic_s.append(thicss + sthc)

"""Loop Properties"""
dt = 0.05                   #[s]            Timestep size
time = 2000                 #[s]            Time of simulation
steps = time/dt + 1         #               Number of timesteps

###############################################################################
###############################################################################

"""Area Calculations"""
# Area iterative generator - Calculates sliced spherical area based on angle and radius
# Used for nose
def area_slice(r, t0, t1):
    A = 2 * np.pi * r * r * (1 - np.cos(np.radians(t1)) - (1 - np.cos(np.radians(t0))))
    return(A)

# Area iterative generator - Calculates sliced conical area based on angle and radius
# Used for flat front sides of shield
def areacone_slice(t1, t2):
    r1 = radius * (1 - np.cos(np.radians(t1)))
    r2 = radius * (1 - np.cos(np.radians(t2)))
    h1 = (r1 / np.tan(np.radians(t1))) - 0.0032
    h2 = (r2 / np.tan(np.radians(t2))) - 0.0032
    A1 = np.pi * r1  * np.sqrt((h1**2) + (r1**2))
    A2 = np.pi * r2  * np.sqrt((h2**2) + (r2**2))
    return(A2-A1)

# Start iterations at 0
ar_it = 0
vol_it = 0
vol_s = 0
mass_it = 0
i = 0
areas = []
conestart = 30.502          #[mm]           Heigh where curved nose transitiones to flat sides

# Iterate through angles appending area, volume and masses
for i in range(step - 1):
    if i <= conestart:
        x = area_slice(radius, angles[i], angles[i+1])
        y = x * thic_p[i]
        ar_it += x
        vol_it += y
        mass_it += y * pica_rho
        vol_s += x * thic_s[i]
        areas.append(x)
    if i > conestart:
        x = areacone_slice(angles[i], angles[i+1])
        y = x * thic_p[i]
        ar_it += x
        vol_it += y
        mass_it += y * pica_rho
        vol_s += x * thic_s[i]
        areas.append(x)

# Extract masses and areas
mass_in = vol_s * ins_rho
ar_ac_cir = 2 * np.pi * noser * noser * (1 - np.cos(np.radians(30.52)))
cone1 = np.pi * radius * np.sqrt((0.0859181**2) + (radius**2))
cone2 = np.pi * 0.04023 * np.sqrt((0.04023**2) + (0.0208181**2))
ar_ac_con = cone1 - cone2
ar_ac = ar_ac_cir + ar_ac_con

###############################################################################
###############################################################################

# Sum masses for total vehicle mass
mass = 4.32 #mass_it + mass_cube + mass_struc + mass_bolts + mass_in + mass_other
BC = mass/(S*Cd)            #               Ballistic coefficient

def g_acc(alt):

    """Calculates acceleration due to gravity at a certain altitude.
    params
    input
    alt [float] Altitude in m
    output
    acceleration in m/s^2
    """
    g = (G*Me)/((Re + alt)**2)
    return g

def g_mod(v):

    """Calculates gravitational acceleration modifier based on velocity.
    params
    input
    vx [float] horizontal velocity in m/s
    output
    modifier
    """
    if v >= 7800:
        g_mod = 0
    else:
        g_mod = (7800 - v)/7800
    return g_mod

def density(alt):

    """Calculates density at a certain altitude.
    params
    input
    alt [float] Altitude in m
    output
    density in kg/m^3
    """
    rho = surface_rho*np.exp(-beta*alt)
    return rho

def drag(v,rho):

    """Calculates drag force at certain flow conditions.
    params
    input
    v [float] velocity in m/s
    rho [float] density in kg/m^3
    output
    drag force in N
    """

    Fd = 0.5*Cd*(v**2)*rho*S
    return Fd

def lift(v,rho):

    """Calculates lift force at certain flow conditions.
    params
    input
    v [float] velocity in m/s
    rho [float] density in kg/m^3
    output
    lfit force in N
    """

    Fl = 0.5*Cl*(v**2)*rho*S
    return Fl

def TrajectorySolver(gamma):

    """Calculates Trajectory of Re-Entry Vehicle based on vehicle paramaters.
    params
    input
    gamma - initial flight path angle in degrees
    variables at top of code
    output
    altitude array in m
    velocity array in m/s
    time at which parachute deployed in s
    """

    init_fpa = np.radians(gamma)
    t_vals = np.linspace(0,time,int(steps))

    #sets up required values for loop with initial parameters
    vx = vx_prev = init_v*np.cos(init_fpa)
    vy = vy_prev = init_v*np.sin(init_fpa)
    fpa = init_fpa
    alt = alt_prev = init_alt
    disp = disp_prev = init_disp

    #initialises arrays
    v_vals = np.zeros(int(steps))
    alt_vals = np.zeros(int(steps))
    disp_vals = np.zeros(int(steps))
    i = 0

    #for loop calculates trajectory
    for i,t in enumerate(t_vals):

        v = np.sqrt((vx**2)+(vy**2)) #Gets total velocity vector
        alt_vals[i] = alt            #Adds altitude val to array
        disp_vals[i] = disp          #Adds displacement val to array
        v_vals[i] = v                #Adds velocity val to array
        rho = density(alt)           #Calculates density at this step
        ga = g_acc(alt)              #Calculates grav acc at this step
        gm = g_mod(v)                #Calculates grav modifier at this step
        g = ga * gm                  #Applies grav modifier
        liftval = lift(v,rho)        #Calculates lift (should be zero)
        dragval = drag(v,rho)        #Calculates Drag at this step
        Fx = liftval*np.sin(fpa) - dragval*np.cos(fpa)#X Direction Force Calcs
        Fy = mass*g - liftval*np.cos(fpa) \
            - dragval*np.sin(fpa)   #Y Direction Force Calcs

        vx = vx_prev + ((Fx/mass) * dt) #Calcs v in x direction
        if vx <= 0:
            vx = 0
        vy = vy_prev + ((Fy/mass) * dt) #Calcs v in y direction

        if vx == 0:         #Calcs new flight path angle based on vx and vy
            fpa = np.radians(90)
        else:
            fpa = np.arctan(vy/vx)

        alt = alt_prev - (vy*dt + 0.5*(Fy/mass)*(dt**2)) #Calcs new altitude
        disp = disp_prev + (vx*dt + 0.5*(Fx/mass)*(dt**2)) #Calcs new x disp
        vx_prev = vx   #Sets vx val for next loop
        vy_prev = vy   #Sets vy val for next loop
        alt_prev = alt #Sets alt val for next loop
        disp_prev = disp
        tlim = t
        if alt <= 2000: #Breaks loop at parachute deployment
        #    print("Parachute deployed at",round(alt,2), "m and Velocity = ",\
        #          round(v,2),"m/s after", t, "seconds of flight time.")
            break
    return alt_vals,disp_vals,v_vals,tlim

def a_val(v_vals):

    """Produces acceleration array based on velocity array.
    params
    input
    v_vals array [floats] velocity in m/s
    output
    Acceleration array in g's'
    """

    a = np.zeros(int(steps))
    for i in range(int(steps)):
        if i == 0:
            a[i] = 0
        else:
            a[i] = v_vals[i-1] - v_vals[i]
    a_vals = a/(9.81*dt)
    return a_vals

def array_cleaner(gamma):

    """Cleans extra zero values from TrajectorySolver arrays.
    params
    input
    gamma - initial flight path angle in degrees
    variables at top of code
    output
    TrajectorySolver and Acc arrays pruned of excess end values.
    """

    t_vals = np.linspace(0,time,int(steps))
    alt_vals,disp_vals,v_vals,tlim = TrajectorySolver(gamma)
    a_vals = a_val(v_vals)
    i = 0
    check = 0

    for i in range(int(steps)):

        #If the parachute hasn't deployed, do nothing and go to next timestep
        if i <= (tlim/dt):
            check += 1
            continue

        #If parachute has deployed, delete unnecessary values
        else:
            alt_vals = np.delete(alt_vals,check)
            disp_vals = np.delete(disp_vals,check)
            v_vals = np.delete(v_vals,check)
            a_vals = np.delete(a_vals,check)
            t_vals = np.delete(t_vals,check)

    rho_vals = density(alt_vals) #Creates density array

    return alt_vals,disp_vals,rho_vals,v_vals,a_vals,t_vals,gamma

"""Running the Code"""
alt_vals,disp_vals,rho_vals,v_vals,a_vals,t_vals,gamma = array_cleaner(10)

# Itertive Numerical Process
"""The for-loop cycles first through 'p' lists of
velocity and density values, each list being a
different flight path angle (FPA). For each FPA,
the loop then cycles through time steps, and for
each time step the area slicing method is conducted
and summed."""

# Open arrays for heat loads and heat fluxes
Heat_Load_in = []
Heat_flux_list = []
Heat_load_list = []
q_conv_list = []
q_rerad_list = []
q_net_list = []

# Open matrices from temperature of each material
# Appends temperature for all angles for each timestep
Ts = np.ones((len(v_vals), step))
Tc = np.ones((len(v_vals), step))
Ti = np.ones((len(v_vals), step))
Ta = np.ones((len(v_vals), step))
Tal = np.ones((len(v_vals), step))

# Open matrix for ablation
# Appends ablation for all angles for each timestep
ab_mount = np.zeros((len(v_vals), step)) #[m]

# Set intial temperatures
Ts[:,:] = 300               #[K]    PICA
Tc[:,:] = 300               #[K]    Cubesat
Ti[:,:] = 300               #[K]    Carbon felt
Ta[:,:] = 300               #[K]    Foam
Tal[:,:] = 300              #[K]    Aluminium

# Open lists to append temperatures at different points
# Nose
Ti0_List = []
Tc0_List = []
Ts0_List = []
Ta0_List = []
Tal0_List = []

# 21 degrees
Ti21_List = []
Tc21_List = []
Ts21_List = []
Ta21_List = []
Tal21_List = []

# 42 degrees
Ti42_List = []
Tc42_List = []
Ts42_List = []
Ta42_List = []
Tal42_List = []

# 64 degrees
Ti64_List = []
Tc64_List = []
Ts64_List = []
Ta64_List = []
Tal64_List = []

# 86 degrees (Edge)
Ti86_List = []
Tc86_List = []
Ts86_List = []
Ta86_List = []
Tal86_List = []

# Open arrays for recorded values
# heatrad = []
# heatcond = []
Q_abl = []              #[MW/m^2]           Ablated heat flux
Q_abtot = []            #[MW]               Ablated heat flux
mass_abl = []           #[kg]               Ablated mass
line_ab = []            #[m]                Linear ablation
ab_rate = []            #[m/s]              Linear ablation rate at any angle
ab_rate1 = []           #[m/s]              Linear ablation rate at nose
ab_tot = np.zeros(step) #[m]                Total ablation
ab_time = []            #[s]                Ablation time

# Start max heating time
heattime = 0            #[s]
maxtime = 3 * dt        #[s]
mgtime = 3 * dt         #[s]
mTtime = 3 * dt         #[s]

# Iterate through each timestep
for i in range(len(v_vals)): # i is timestep
    # Set flux and load for this slice to 0
    Heat_flux_slices = 0
    Heat_load_slices = 0

    # Calculate stagnation point convection heat flux (W/m^2)
    q_conv_stag = 7.455*(10**(-9))*(rho_vals[i]**0.4705)*(v_vals[i]**(3.089))*(noser**(-0.52))*10000

    # If stagnation point greater than 0.5 MW/m^2, add dt to max heating time
    if q_conv_stag >= 0.5 * (10**6):
        heattime += dt

    # Copy and append stagnation convective heat flux
    q_conv_prev = q_conv_stag.copy()
    q_conv_list.append(q_conv_stag/(10**6)) # MW/m^2

    # For each timestep, iterate through each angle and calcualte fluxes and temperatures
    for j in range(step - 1): # j is angle
        # Set area for this angle
        area = areas[j]

        # Set heat flux at each angle using trend from stardust data
        if j < 15.75:
            q_conv_ang = q_conv_stag * (((1.01*(10**-5))*(j**4)) - (0.000293 * (j**3)) + (0.0022571 * (j**2)) - (0.019333 * j) + 1)
        if j >= 15.75:
            q_conv_ang = q_conv_stag * ((-0.001289 * j) + 0.75435)

        # Set convective heat flux using average over this area
        q_conv_slice = (q_conv_ang + q_conv_prev)/2

        # If first timestep, set all loads to starting
        if i == 0:
            q_rerad = sig*em*((300)**4)
            q_cond = 0
            q_cond_c = 0
            Q_abla = 0
            thic_char = 0
            abrate = 0
            if j == 0:
                Q_abl.append(Q_abla)
                mass_abl.append(0)

        # If not first timestep, calculate all loads and fluxes
        if i != 0:
            # Starting temperature is set as the last temperature for this slice
            Temp = Ts[i-1, j]

            # If temperature is less than ablation temperature, no ablation
            if Temp < (1100 + 273.15):
                Q_abla = 0                      #[MW/m^2] No ablated heat flux
                if j == 0:
                    Q_abl.append(Q_abla)        # Append heat flux per area
                mass_abl.append(0)              #[kg]     No mass ablated
                ab_tot[j] += 0                  #[m]      No ablation
                abrate = 0                      #[m/s]    No ablation rate

                # If temperature is more than ablation temperature, find ablation
            if Temp >= (1100 + 273.15):
                abrate = ( 0.00000001 * Temp - 0.000013732) #[m/s]   Ablation rate
                ab_rate.append(abrate)                      #        Append ablation rate
                abm = abrate * dt                           #[m]     Ablation amount
                ab_mount[i,j] = abm                         #[m]     Append ablation amount
                ab_tot[j] += abm                            #[m]     Add to this slices ablation amount
                Q_abla = pica_rho * abrate * 116299895.55   #[W/m^2] Ablative heat flux per unit area (from research)

                # Append values at nose
                if j == 0:
                    ab_time.append(dt)                      #[s]      Ablation time
                    Q_abl.append(Q_abla * 10**-6)           #[MW/m^2] Ablation rate
                    line_ab.append(abrate*dt)               #[m]      Linear ablation

                # Append total ablation amount
                Q_abtot.append(area*Q_abla*dt)                    #[J]  Total heat ablated
                mass_abl.append(dt * area* Q_abla / 116299895.55) #[kg] Mass ablated

            # Charr thickness
            thic_char = ab_tot[j]

            # Calculate other heat fluxes
            # Reradiated heat flux
            q_rerad = sig*em*((Temp)**4) #[W/m^2]

            # Conduction from top of PICA to top of carbon felt
            q_cond = ((Ts[i-1,j]) - (Ti[i-1,j]))*(((k_pica * (1 - thic_char)) + (k_char * thic_char)) * area)/(thic_p[j])

            # Conduction from top of carbon felt to top of insulative foam
            q_cond_ins = ((Ti[i-1,j]) - (Ta[i-1,j]))*((k_ins * area)/thic_s[j])

            # Conduction from top of insulative foam to cubesat
            q_cond_fo = ((Ta[i-1,j]) - (Tc[i-1,j]))*((k_foa * area)/thic_f)

            # Calculate internal temperatures at end of timestep from resulting net heat flux
            # Carbon felt
            Ti[i,j] = Ti[i-1,j] + (((q_cond - q_cond_ins) * dt) / (area * pica_rho * thic_p[j] * cpc))

            # Insulative foam
            Ta[i,j] = Ta[i-1,j]  + (((q_cond_ins) * dt) / (area * ins_rho * thic_s[j] * cpin))

            # Aluminium structure
            Tal[i,j] = Tal[i-1,j]  + (((q_cond_ins) * dt) / (area * ins_rho * thic_al * cpal))

            # CubeSat
            Tc[i,j] = Tc[i-1,j]  + (((q_cond_fo) * dt) / (area * fo_rho * thic_f * cpfo))

            # To account for overreradiation at beginning set all temperatures to always above 300K
            if Ti[i,j] < 300:
                Ti[i,j] = 300
            if Ta[i,j] < 300:
                Ta[i,j] = 300
            if Tc[i,j] < 300:
                Tc[i,j] = 300
            if Tal[i,j] < 300:
                Tal[i,j] = 300

        # Calculate heat fluxes for area slice
        Heat_flux_conv = q_conv_slice * area    # W
        Heat_flux_rerad = q_rerad * area        # W
        Heat_flux_cond = q_cond                 # W
        Heat_flux_abla = Q_abla * area          # W

        # If not first timestep, find resulting surface temperature for next timestep from net heat flux
        if i != 0:
            Ts[i,j] = Ts[i-1,j] + (((Heat_flux_conv - Heat_flux_abla - Heat_flux_rerad - Heat_flux_cond) * dt) / (area * pica_rho * thic_p[j] * cpc))
            if Ts[i,j] < 300:
                Ts[i,j] = 300

        # Add loads and fluxes at surface
        Heat_Load_in += Heat_flux_conv * dt                                 #[J] Convective heat load
        Heat_flux_slices += area * (q_conv_slice - q_rerad - Q_abla)        #[W] Net heat flux
        Heat_load_slices += area * (q_conv_slice - q_rerad - Q_abla) * dt   #[J] Net heat load

        # Set previous convection for calculating average flux over area at next angle
        q_conv_prev = q_conv_ang.copy()

        # Append temperatures at different angles for plotting
        # Nose (Stagnation point)
        if j == 0:
            q_rerad_list.append(-q_rerad*(10**(-6)))                    #[MW/m^2] Reradiated heat flux
            q_net_list.append((q_conv_slice-q_rerad-Q_abla)*10**(-6))   #[MW/m^2] Net heat flux
            ab_rate1.append(abrate)                                     #[m/s]    Ablation rate
            Ti0_List.append(Ti[i,j])
            Tc0_List.append(Tc[i,j])
            Ts0_List.append(Ts[i,j])
            Ta0_List.append(Ta[i,j])
            Tal0_List.append(Tal[i,j])
        if j == 21:
            Ti21_List.append(Ti[i,j])
            Tc21_List.append(Tc[i,j])
            Ts21_List.append(Ts[i,j])
            Ta21_List.append(Ta[i,j])
            Tal21_List.append(Tal[i,j])
        if j == 42:
            Ti42_List.append(Ti[i,j])
            Tc42_List.append(Tc[i,j])
            Ts42_List.append(Ts[i,j])
            Ta42_List.append(Ta[i,j])
            Tal42_List.append(Tal[i,j])
        if j == 64:
            Ti64_List.append(Ti[i,j])
            Tc64_List.append(Tc[i,j])
            Ts64_List.append(Ts[i,j])
            Ta64_List.append(Ta[i,j])
            Tal64_List.append(Tal[i,j])
        if j == 86:
            Ti86_List.append(Ti[i,j])
            Tc86_List.append(Tc[i,j])
            Ts86_List.append(Ts[i,j])
            Ta86_List.append(Ta[i,j])
            Tal86_List.append(Tal[i,j])

    if i > 3:
        if q_conv_list[-1] > q_conv_list[-2]:
            maxtime += dt
        if a_vals[i] > a_vals[i-1]:
            # print("AAAA")
            mgtime += dt
        if Ts[i,0] > Ts[i-1,0]:
            # print("BBBB")
            mTtime += dt
    # Append heat flux and load for each timesep
    Heat_flux_list.append(Heat_flux_slices/(10**6)) #[MW/m^2]
    Heat_load_list.append(Heat_load_slices)         #[J]

# Open lists to plot temperatures and values of interest
surftemp_plot = []      #[K]    Surface temperatures
sttemp_plot = []        #[K]    Structure temperatures
cubetemp_plot = []      #[K]    Cubesat temperatures
intemp_plot = []        #[K]    Carbon felt temperatures
altemp_plot = []        #[K]    Foam temperatures
ablation_plot = []      #[mm]   Ablation amount
angle_plot = []         #[deg]  Angles

# For each angle append: the max temperature, ablation amount, and angle
for n in range(step-1):
    surftemp_plot.append(max(Ts[:,n]))
    cubetemp_plot.append(max(Tc[:,n]))
    intemp_plot.append(max(Ti[:,n]))
    altemp_plot.append(max(Ta[:,n]))
    sttemp_plot.append(max(Tal[:,n]))
    ablation_plot.append(sum(ab_mount[:,n]) * 10**3)
    angle_plot.append(n)

# Set edge values in accordance with Stardust data
for n in range(1):
    angle_plot.append(86)
    surftemp_plot.append(300)
    cubetemp_plot.append(300)
    intemp_plot.append(300)
    altemp_plot.append(300)
    sttemp_plot.append(300)
    ablation_plot.append(0)

# Extract relevant data for results
altop = np.max(Ta) - 273.15       #[C]      Maximum insulative foam temperature
Talll = np.max(Ti) - 273.15       #[C]      Maximum carbon felt temperature
stTalll = np.max(Tal) - 273.15    #[C]      Maximum structure temperature
Hottt = np.max(Ts) - 273.15       #[C]      Maximum surface temperature
Tcube = np.max(Tc) - 273.15       #[C]      Maximum CubeSat temperature
maxcon = max(q_conv_list)         #[MW/m^2] Maximum convective heat flux
QMAX = max(q_net_list)            #[MW/m^2] Maximum net heat flux
HLOAD = (sum(Heat_flux_list)*dt)  #[MJ]     Total Heat Load
ab_len = sum(line_ab) * 10**3     #[mm/s]   Length of ablation
ab_tim = sum(ab_time)             #[s]      Time ablating
ab_rates = Q_abl                  #[MW/m^2] Ablation heat fluxes at nose

# Plot PICA thickness as a function of angle
plt.figure(453)
plt.plot(angle_plot, thic_p)
plt.title("PICA Thickness Plotted as a Function of Angle")
plt.xlabel('Angle (degree)')
plt.ylabel('PICA Thickness (mm)')
plt.show()

# Plot carbon felt thickness as a function of angle
plt.figure(457)
plt.plot(angle_plot, thic_s)
plt.title("Carbon Felt Thickness Plotted as a Function of Angle")
plt.xlabel('Angle (degree)')
plt.ylabel('PICA Thickness (mm)')
plt.show()

# Plot linear ablation amount as a function of angle
plt.figure(124242)
plt.plot(angle_plot, ablation_plot)
plt.title("Linear Ablation Amount Plotted as a Function of Angle")
plt.xlabel('Angle (degree)')
plt.ylabel('Linear Ablation (mm)')
plt.show()

# Plot maximum surface temperature as a function of angle
plt.figure(124243)
plt.plot(angle_plot, surftemp_plot)
plt.title("Maximum Surface Temperature Plotted as a Function of Angle")
plt.xlabel('Angle (degree)')
plt.ylabel('Temperature (K)')
plt.show()

# Plot maximum carbon felt temperature as a function of angle
plt.figure(124249)
plt.plot(angle_plot, intemp_plot)
plt.title("Maximum Carbon Felt Temperature Plotted as a Function of Angle")
plt.xlabel('Angle (degree)')
plt.ylabel('Temperature (K)')
plt.show()

# Plot maximum insulative foam temperature as a function of angle
plt.figure(124244)
plt.plot(angle_plot, altemp_plot)
plt.title("Maximum Insulative Foam Temperature Plotted as a Function of Angle")
plt.xlabel('Angle (degree)')
plt.ylabel('Temperature (K)')
plt.show()

# Plot maximum CubeSat temperature as a function of angle
plt.figure(1243)
plt.plot(angle_plot, cubetemp_plot)
plt.title("Maximum CubeSat Temperature Plotted as a Function of Angle")
plt.xlabel('Angle (degree)')
plt.ylabel('Temperature (K)')
plt.show()

# Plot nose temperatures as a function of time
plt.figure(0)
plt.plot(t_vals,Ts0_List, "r", label="PICA")
plt.plot(t_vals,Tal0_List, "k", label="Structure")
plt.plot(t_vals,Ti0_List, "b", label="Carbon Felt")
plt.plot(t_vals,Ta0_List, "c-.", label="Foam")
plt.plot(t_vals,Tc0_List, "g", label="CubeSat")
plt.title("Nose Temperatures Plotted as a Function of Time")
plt.xlabel('Time (s)')
plt.ylabel('Temperature (K)')
plt.legend()
plt.show()

# Plot temperatures at and anlge 21 degrees from nose as a function of time
plt.figure(21)
plt.plot(t_vals,Ts21_List, "r", label="PICA")
plt.plot(t_vals,Tal21_List, "k", label="Structure")
plt.plot(t_vals,Ta21_List, "c-.", label="Foam")
plt.plot(t_vals,Ti21_List, "b", label="Carbon Felt")
plt.plot(t_vals,Tc21_List, "g", label="CubeSat")
plt.title("21 Degrees Temperatures Plotted as a Function of Time")
plt.xlabel('Time (s)')
plt.ylabel('Temperature (K)')
plt.legend()
plt.show()

# Plot temperatures at and anlge 42 degrees from nose as a function of time
plt.figure(42)
plt.plot(t_vals,Ts42_List, "r", label="PICA")
plt.plot(t_vals,Tal42_List, "k", label="Structure")
plt.plot(t_vals,Ta42_List, "c-.", label="Foam")
plt.plot(t_vals,Ti42_List, "b", label="Carbon Felt")
plt.plot(t_vals,Tc42_List, "g", label="CubeSat")
plt.title("42 Degrees Temperatures Plotted as a Function of Time")
plt.xlabel('Time (s)')
plt.ylabel('Temperature (K)')
plt.legend()
plt.show()

# Plot temperatures at and anlge 64 degrees from nose as a function of time
plt.figure(64)
plt.plot(t_vals,Ts64_List, "r", label="PICA")
plt.plot(t_vals,Tal64_List, "k", label="Structure")
plt.plot(t_vals,Ta64_List, "c:", label="Foam")
plt.plot(t_vals,Ti64_List, "b", label="Carbon Felt")
plt.plot(t_vals,Tc64_List, "g", label="CubeSat")
plt.title("64 Degrees Temperatures Plotted as a Function of Time")
plt.xlabel('Time (s)')
plt.ylabel('Temperature (K)')
plt.legend()
plt.show()

# Plot temperatures at and anlge 86 degrees from nose as a function of time
plt.figure(86)
plt.plot(t_vals,Ts86_List, "r", label="PICA")
plt.plot(t_vals,Tal86_List, "k", label="Structure")
plt.plot(t_vals,Ta86_List, "c-.", label="Foam")
plt.plot(t_vals,Ti86_List, "b", label="Carbon Felt")
plt.plot(t_vals,Tc86_List, "g", label="CubeSat")
plt.title("86 Degrees Temperatures Plotted as a Function of Time")
plt.xlabel('Time (s)')
plt.ylabel('Temperature (K)')
plt.legend()
plt.show()

# Plot ablation heat flux at nose as a function of time
plt.figure(111)
plt.plot(t_vals, ab_rates )
plt.title("Ablation Heat Flux at Nose Plotted as a Function of Time")
plt.xlabel('Time (s)')
plt.ylabel('Heat Flux (MW/m^2)')
plt.show()

# Plot stagnation point convective heat flux as a function of time
plt.figure(1)
plt.plot(t_vals,q_conv_list)
plt.title("Stagnation Point Convective Heat Flux Plotted as a Function of Time")
plt.xlabel('Time (s)')
plt.ylabel('Heat Flux (MW/m^2)')
plt.show()

# Plot stagnation point reradiative heat flux as a function of time
plt.figure(2)
plt.plot(t_vals,q_rerad_list)
plt.title("Stagnation Point Reradiative Heat Flux Plotted as a Function of Time")
plt.xlabel('Time (s)')
plt.ylabel('Heat Flux (MW/m^2)')
plt.show()

# Plot stagnation point net heat flux as a function of time
plt.figure(6)
plt.plot(t_vals,q_net_list)
plt.title("Stagnation Point Net Heat Flux Plotted as a Function of Time")
plt.xlabel('Time (s)')
plt.ylabel('Heat Flux (MW/m^2)')
#plt.legend()
plt.show()

# Plot total heat flux as a function of time
plt.figure(3)
plt.plot(t_vals,Heat_flux_list)
plt.title("Total Heat Flux Plotted as a Function of Time")
plt.xlabel('Time (s)')
plt.ylabel('Heat Flux (MW)')
plt.show()

# Plot ablation rate at nose as a function of convective heat flux
plt.figure(143320094)
plt.plot(q_conv_list, ab_rate1)
plt.xlabel("Convective Heat Flux at Nose (MW/m^2)")
plt.ylabel("Ablation Rate (mm/s)")
plt.title("Nose Ablation Rate Plotted as a Function of Convective Heat Flux")
plt.show()

# Plot stagnation point convective heat flux as a function of altitude
plt.figure(1235234234234)
plt.plot(alt_vals / 1000, q_conv_list)
plt.title("Stagnation Point Convective Heat Flux Plotted as a Function of Altitude")
plt.xlabel('Altitude (km)')
plt.ylabel('Heat Flux (MW/m^2)')
plt.show()

# Print results
print("\n\nMass = {g:0.5f} kg".format(g=mass))
print("PICA thickness = {g:0.3f} mm\nInsulator thickness = {gg:0.3f} mm\nFoam thickness = {ggg:0.3f} mm".format(g = pthc*10**3, gg = sthc*10**3, ggg = fthc*10**3))
print("\nThe maximum temperature of the CubeSat = {g:0.3f} C, {gg:0.3f} K".format(g = Tcube, gg = Tcube+273.15))
print("\nThe maximum temperature of the foam = {g:0.3f} C, {gg:0.3f} K".format(g = altop, gg = altop+273.15))
print("\nThe maximum temperature of the insulation = {g:0.3f} C, {gg:0.3f} K".format(g = Talll, gg = Talll+273.15))
print("\nThe maximum temperature of the structure = {g:0.3f} C, {gg:0.3f} K".format(g = stTalll, gg = stTalll+273.15))
print("\nThe maximum surface temperature of the heat shield = {g:0.3f} C, {gg:0.3f} K".format(g = Hottt, gg = Hottt+273.15))
print("\nThe maximum instantaneous convective heat flux through the shield is {g:0.6f} MW/m^2".format(g = maxcon))
print("\nThe maximum instantaneous net heat flux throught the shield is {g:0.6f} MW/m^2".format(g = QMAX))
print("\nThe total heat load is {g:0.6f} MJ".format(g = HLOAD))
print("\nThe nose ablates {g:0.6f} mm".format(g = ab_len))
print("\nThe nose ablates for {g:0.4f} seconds".format(g = ab_tim))
print("\nTotal mass ablated =", sum(mass_abl), "kg")
print("\nMaximum ablation rate =", max(ab_rate)*10**3, "mm/s")
print("\nTotal heat ablated =", sum(Q_abtot)*10**-6, "MJ")
print("\nMaximum heat ablation rate =", max(Q_abl), "MW/m^2\n\n")

plt.plot(t_vals, alt_vals/1000, 'k', label = "Trajectory")
plt.plot(t_vals[int(round(maxtime, 3)/0.05)], alt_vals[int(round(maxtime, 3)/0.05)]/1000, 'b^', label = "Max Convective Heat Flux")
plt.plot(t_vals[int(round(mgtime, 3)/0.05)], alt_vals[int(round(mgtime, 3)/0.05)]/1000, 'go', label = "Max G Load")
plt.plot(t_vals[int(round(mTtime, 3)/0.05)], alt_vals[int(round(mTtime, 3)/0.05)]/1000, 'rX', label = "Max Temperature")
plt.title("Altitude Plotted as a Function of Time")
plt.ylabel('Altitude (km)')
plt.xlabel('Time (s)')
plt.legend()
plt.show()

plt.plot(t_vals, v_vals/1000, 'k', label = "Trajectory")
plt.plot(t_vals[int(round(maxtime, 3)/0.05)], v_vals[int(round(maxtime, 3)/0.05)]/1000, 'b^', label = "Max Convective Heat Flux")
plt.plot(t_vals[int(round(mgtime, 3)/0.05)], v_vals[int(round(mgtime, 3)/0.05)]/1000, 'go', label = "Max G Load")
plt.plot(t_vals[int(round(mTtime, 3)/0.05)], v_vals[int(round(mTtime, 3)/0.05)]/1000, 'rX', label = "Max Temperature")
plt.title("Velocity Plotted as a Function of Altitude")
plt.ylabel('Velocity (m/s)')
plt.xlabel('Time (s)')
plt.legend()
plt.show()

plt.plot(disp_vals/1000, alt_vals/1000, 'k', label = "Trajectory")
plt.plot(disp_vals[int(round(maxtime, 3)/0.05)]/1000, alt_vals[int(round(maxtime, 3)/0.05)]/1000, 'b^', label = "Max Convective Heat Flux")
plt.plot(disp_vals[int(round(mgtime, 3)/0.05)]/1000, alt_vals[int(round(mgtime, 3)/0.05)]/1000, 'go', label = "Max G Load")
plt.plot(disp_vals[int(round(mTtime, 3)/0.05)]/1000, alt_vals[int(round(mTtime, 3)/0.05)]/1000, 'rX', label = "Max Heat Flux")
plt.title("Altitude Plotted as a Function of Downrange Displacement")
plt.ylabel('Altitude (km)')
plt.xlabel('Downrange Displacement (km)')
plt.legend()
plt.show()

# Plot nose temperatures as a function of time
plt.figure(121213224314324423)
plt.plot(t_vals,np.array(Tal0_List)-273.15, "k", label="Structure")
plt.title("Structure Maximum Temperature Plotted as a Function of Time")
plt.xlabel('Time (s)')
plt.ylabel('Temperature (C)')
plt.legend()
plt.show()

Trajectory Analysis and Requirements Generation/Aerothermodynamic Analysis Code.py


from functools import partial
from multiprocessing import Pool
from pathlib import Path
from sklearn.model_selection import ParameterGrid  # type: ignore
import argparse
import json
import logging
import matplotlib.pyplot as plt  # type: ignore
import numpy as np  # type: ignore
import tensorflow as tf  # type: ignore
import typing as ty

from tensorflow.keras.callbacks import EarlyStopping  # type: ignore
from tensorflow.keras.callbacks import LearningRateScheduler  # type: ignore
from tensorflow.keras.callbacks import ModelCheckpoint  # type: ignore
from tensorflow.keras.optimizers import Adam  # type: ignore
from tensorflow.keras.optimizers import RMSprop  # type: ignore

from area_model import ActualAreaNet
from area_model import AreaNet
from area_model import SimpleNet
from area_model import VerticalAreaNet
from augment_data import do_augmentation
from clr_callback import CyclicLR
from lr_finder import LearningRateFinder
from plot_utils import plot_confusion_matrix
from preprocess_data import load_processed
from preprocess_data import preprocess_spec
from schedules import exp_decay_smooth
from schedules import exp_decay_step
from utils import analyze_confusion
from utils import pred_hot_2_cm
from utils import setup_gpus
from utils import setup_logger
from utils import words_types

def parse_arguments() -> argparse.Namespace:
    """Setup CLI interface"""
    parser = argparse.ArgumentParser(description="")

    parser.add_argument(
        "-tt",
        "--training_type",
        type=str,
        default="hypa_tune",
        choices=["hypa_tune"],
        help="Which training to execute",
    )

    tra_types = [w for w in words_types.keys() if not w.startswith("_")]
    parser.add_argument(
        "-wt",
        "--words_type",
        type=str,
        default="f2",
        choices=tra_types,
        help="Words to preprocess",
    )

    parser.add_argument(
        "-ft",
        "--force_retrain",
        dest="force_retrain",
        action="store_true",
        help="Force the training and overwrite the previous results",
    )

    parser.add_argument(
        "-nv",
        "--no_use_validation",
        dest="use_validation",
        action="store_false",
        help="Do not use validation data while training",
    )

    parser.add_argument(
        "-dr",
        "--dry_run",
        action="store_true",
        help="Do a dry run for the hypa grid",
    )

    parser.add_argument(
        "-lrf",
        "--do_find_best_lr",
        dest="do_find_best_lr",
        action="store_true",
        help="Find the best values for the learning rate",
    )

    # last line to parse the args
    args = parser.parse_args()
    return args

def setup_env() -> argparse.Namespace:
    setup_logger("DEBUG")
    args = parse_arguments()
    # build command string to repeat this run
    # FIXME if an option is a flag this does not work, sorry
    recap = "python3 train_area.py"
    for a, v in args._get_kwargs():
        recap += f" --{a} {v}"
    logmain = logging.getLogger(f"c.{__name__}.setup_env")
    logmain.info(recap)
    return args

def train_model_area_dry(
    hypa: ty.Dict[str, str], use_validation: bool, trained_folder: Path
) -> str:
    """MAKEDOC: what is train_model_area_dry doing?"""
    # logg = logging.getLogger(f"c.{__name__}.train_model_area_dry")
    # logg.setLevel("INFO")
    # logg.debug("Start train_model_area_dry")

    model_name = build_area_name(hypa, use_validation)
    placeholder_path = trained_folder / f"{model_name}.txt"

    if placeholder_path.exists():
        return "already_trained"

    return "to_train"

def hyper_train_area(
    words_type: str,
    force_retrain: bool,
    use_validation: bool,
    dry_run: bool,
    do_find_best_lr: bool,
) -> None:
    """MAKEDOC: what is hyper_train_area doing?"""
    logg = logging.getLogger(f"c.{__name__}.hyper_train_area")
    # logg.setLevel("INFO")
    logg.debug("Start hyper_train_area")

    ##########################################################
    #   Hyper-parameters grid
    ##########################################################

    hypa_grid: ty.Dict[str, ty.List[str]] = {}

    ###### the net type
    nt = []
    nt.append("SIM")
    nt.append("SI2")
    nt.append("AAN")
    nt.append("VAN")
    hypa_grid["net_type"] = nt

    ###### the words to train on
    wtl = []
    # wtl.extend([words_type])
    # wtl.extend(["LTnum", "LTall", "yn"])
    # wtl.extend(["LTnum", "yn"])
    # wtl.extend(["LTnum", "yn", "f1", "k1"])
    # wtl.extend(["LTBnum", "LTBall"])
    # wtl.extend(["LTBnum"])
    # wtl.extend(["num"])
    # wtl.extend(["LTBnumLS", "LTBallLS"])
    wtl.extend(["LTBnumLS"])
    # wtl.extend(["lr"])
    # wtl.extend(["similar", "all"])
    # wtl.extend(["all"])

    hypa_grid["words_type"] = wtl

    ###### the dataset to train on
    ds = []

    # TODO VAN on LTall
    # TODO AAN/SIM/VAN on LTnum
    # TODO AAN on LTnum for all datasets, only one lr
    # ds.extend(["mel04"])
    # ds.extend(["mela1"])
    # ds.extend(["aug07"])
    # ds.extend(["aug14"])
    # ds.extend(["aug15"])

    # TODO auL6789 auL18901 on all net_type
    # TODO auA5678 on VAN (on LTnumLS)
    # TODO auA5678 with lr04 (on LTnumLS to complete lr03 is done)
    # MAYBE start removing ARN, too much time
    # ds.extend(["auA01", "auA02", "auA03", "auA04"])
    # ds.extend(["auA05", "auA06", "auA07", "auA08"])
    # ds.extend(["auA04"])
    ds.extend(["mel04L"])

    # TODO just the 3 best per architecture on noval

    hypa_grid["dataset_name"] = ds

    ###### the learning rates for the optimizer
    lr = []
    # lr.extend(["01", "02"])  # fixed
    lr.extend(["03"])  # exp_decay_step_01
    lr.extend(["04"])  # exp_decay_smooth_01
    # lr.extend(["05"])  # clr_triangular2_01
    # lr.extend(["06"])  # clr_triangular2_02
    hypa_grid["learning_rate_type"] = lr

    ###### which optimizer to use
    # hypa_grid["optimizer_type"] = ["a1", "r1"]
    hypa_grid["optimizer_type"] = ["a1"]

    ###### the batch size (the key is converted to int)
    bs = []
    bs.extend(["32"])
    # bs.extend(["16"])
    hypa_grid["batch_size_type"] = bs

    ###### the number of epochs (the key is converted to int)
    en = []
    en.extend(["15"])
    # en.extend(["13"])
    # en.extend(["10"])
    hypa_grid["epoch_num_type"] = en

    ###### build the combinations
    logg.debug(f"hypa_grid = {hypa_grid}")
    the_grid = list(ParameterGrid(hypa_grid))

    # hijack the grid
    # hypa = {
    #     "batch_size_type": "32",
    #     "dataset_name": "auA07",
    #     "epoch_num_type": "15",
    #     "learning_rate_type": "03",
    #     "net_type": "VAN",
    #     "optimizer_type": "a1",
    #     "words_type": "LTnumLS",
    # }
    # the_grid = [hypa]

    num_hypa = len(the_grid)
    logg.debug(f"num_hypa: {num_hypa}")

    ##########################################################
    #   Setup pre train
    ##########################################################

    train_type_tag = "area"

    # where to save the trained models
    trained_folder = Path("trained_models") / train_type_tag
    if not trained_folder.exists():
        trained_folder.mkdir(parents=True, exist_ok=True)

    # where to put the info folders
    root_info_folder = Path("info") / train_type_tag
    if not root_info_folder.exists():
        root_info_folder.mkdir(parents=True, exist_ok=True)

    # count how many models are left to train
    if dry_run:
        tra_info = {"already_trained": 0, "to_train": 0}
        for hypa in the_grid:
            train_status = train_model_area_dry(hypa, use_validation, trained_folder)
            tra_info[train_status] += 1
        logg.debug(f"tra_info: {tra_info}")
        return

    # check that the data is available
    for dn in hypa_grid["dataset_name"]:
        for wt in hypa_grid["words_type"]:
            logg.debug(f"\nwt: {wt} dn: {dn}\n")
            if dn.startswith("mel"):
                preprocess_spec(dn, wt)
            elif dn.startswith("aug"):
                do_augmentation(dn, wt)

    # useful to pick good values of lr
    if do_find_best_lr:
        hypa = the_grid[0]
        logg.debug(f"\nSTARTING find_best_lr with hypa: {hypa}")
        find_best_lr(hypa)
        return

    ##########################################################
    #   Train all hypas
    ##########################################################

    for i, hypa in enumerate(the_grid):
        logg.debug(f"\nSTARTING {i+1}/{num_hypa} with hypa: {hypa}")
        with Pool(1) as p:
            p.apply(
                train_area,
                (hypa, force_retrain, use_validation, trained_folder, root_info_folder),
            )

def build_area_name(hypa: ty.Dict[str, str], use_validation: bool) -> str:
    """MAKEDOC: what is build_area_name doing?"""
    logg = logging.getLogger(f"c.{__name__}.build_area_name")
    logg.setLevel("INFO")
    logg.debug("Start build_area_name")

    # model_name = "ARN"
    # model_name = "SIM"
    # model_name = "AAN"
    model_name = hypa["net_type"]
    model_name += f"_op{hypa['optimizer_type']}"
    model_name += f"_lr{hypa['learning_rate_type']}"
    model_name += f"_bs{hypa['batch_size_type']}"
    model_name += f"_en{hypa['epoch_num_type']}"

    model_name += f"_ds{hypa['dataset_name']}"
    model_name += f"_w{hypa['words_type']}"

    if not use_validation:
        model_name += "_noval"

    return model_name

def get_model_param_area(
    hypa: ty.Dict[str, str], num_labels: int, input_shape: ty.Tuple[int, int, int]
) -> ty.Dict[str, ty.Any]:
    """MAKEDOC: what is get_model_param_area doing?"""
    # logg = logging.getLogger(f"c.{__name__}.get_model_param_area")
    # logg.setLevel("INFO")
    # logg.debug("Start get_model_param_area")

    model_param: ty.Dict[str, ty.Any] = {}

    model_param["num_classes"] = num_labels
    model_param["input_shape"] = input_shape

    return model_param

def get_training_param_area(
    hypa: ty.Dict[str, str],
    use_validation: bool,
    model_path: ty.Optional[Path],
    num_samples: int,
) -> ty.Dict[str, ty.Any]:
    """MAKEDOC: what is get_training_param_area doing?"""
    logg = logging.getLogger(f"c.{__name__}.get_training_param_area")
    # logg.setLevel("INFO")
    # logg.debug("Start get_training_param_area")

    # TODO extract epoch/batch/opt/lr
    training_param: ty.Dict[str, ty.Any] = {}

    training_param["callbacks"] = []

    # it seems silly but the key must be 2 char long for a consistent model name
    # batch_size_types = {"32": 32, "16": 16}
    # batch_size = batch_size_types[hypa["batch_size_type"]]
    # training_param["batch_size"] = batch_size
    training_param["batch_size"] = int(hypa["batch_size_type"])

    # epochs_types = {"15": 15, "30": 30, "02": 2, "04": 4}
    # epochs = epochs_types[hypa["epochs"]]
    # training_param["epochs"] = epochs
    training_param["epochs"] = int(hypa["epoch_num_type"])

    # translate from short key to long name
    learning_rate_types = {
        "01": "fixed01",
        "02": "fixed02",
        "03": "exp_decay_step_01",
        "04": "exp_decay_smooth_01",
        "05": "clr_triangular2_01",
        "06": "clr_triangular2_02",
    }
    learning_rate_type = hypa["learning_rate_type"]
    lr_name = learning_rate_types[learning_rate_type]
    training_param["lr_name"] = lr_name

    if lr_name.startswith("fixed"):
        if lr_name == "fixed01":
            lr = 1e-3
        elif lr_name == "fixed02":
            lr = 1e-4
    else:
        lr = 1e-3

    optimizer_types = {"a1": Adam(learning_rate=lr), "r1": RMSprop(learning_rate=lr)}
    training_param["opt"] = optimizer_types[hypa["optimizer_type"]]

    callbacks = []

    if lr_name.startswith("exp_decay"):
        if lr_name == "exp_decay_step_01":
            exp_decay_part = partial(exp_decay_step, epochs_drop=5)
        elif lr_name == "exp_decay_smooth_01":
            exp_decay_part = partial(exp_decay_smooth, epochs_drop=5)
        lrate = LearningRateScheduler(exp_decay_part)
        callbacks.append(lrate)

    # setup cyclic learning rate
    elif lr_name.startswith("clr_triangular2"):

        # target_cycles = the number of cycles we want in those epochs
        # it_per_epoch = num_samples // batch_size
        # total_iterations = it_per_epoch * epoch_num
        # step_size = total_iterations // target_cycles

        if lr_name == "clr_triangular2_01":
            target_cycles = 2
            it_per_epoch = num_samples // training_param["batch_size"]
            total_iterations = it_per_epoch * training_param["epochs"]
            step_size = total_iterations // (target_cycles * 2)
            base_lr = 1e-5
            max_lr = 1e-3
        elif lr_name == "clr_triangular2_02":
            target_cycles = 8
            it_per_epoch = num_samples // training_param["batch_size"]
            total_iterations = it_per_epoch * training_param["epochs"]
            step_size = total_iterations // (target_cycles * 2)
            base_lr = 1e-6
            max_lr = 1e-3

        logg.debug(f"target_cycles: {target_cycles}")
        logg.debug(f"it_per_epoch: {it_per_epoch}")
        logg.debug(f"total_iterations: {total_iterations}")
        logg.debug(f"num_samples: {num_samples}")
        logg.debug(f"CLR is using step_size: {step_size}")

        mode = "triangular2"
        cyclic_lr = CyclicLR(base_lr, max_lr, step_size, mode)
        callbacks.append(cyclic_lr)

    # which metric to monitor for early_stop and model_checkpoint
    metric_to_monitor = "val_loss" if use_validation else "loss"

    if lr_name.startswith("fixed") or lr_name.startswith("exp_decay"):
        early_stop = EarlyStopping(
            monitor=metric_to_monitor,
            patience=4,
            restore_best_weights=True,
            verbose=1,
        )
        callbacks.append(early_stop)

    # to inhibit checkpointing by passing None
    if model_path is not None:
        model_checkpoint = ModelCheckpoint(
            str(model_path), monitor=metric_to_monitor, verbose=1, save_best_only=True
        )
        callbacks.append(model_checkpoint)

    training_param["callbacks"] = callbacks

    return training_param

def train_area(
    hypa: ty.Dict[str, str],
    force_retrain: bool,
    use_validation: bool,
    trained_folder: Path,
    root_info_folder: Path,
) -> None:
    """MAKEDOC: what is train_area doing?"""
    logg = logging.getLogger(f"c.{__name__}.train_area")
    # logg.setLevel("INFO")
    logg.debug("Start train_area")

    ##########################################################
    #   Setup folders
    ##########################################################

    # name the model
    model_name = build_area_name(hypa, use_validation)
    logg.debug(f"model_name: {model_name}")

    # save the trained model here
    model_path = trained_folder / f"{model_name}.h5"
    placeholder_path = trained_folder / f"{model_name}.txt"

    # check if this model has already been trained
    if placeholder_path.exists():
        if force_retrain:
            logg.warn("\nRETRAINING MODEL!!\n")
        else:
            logg.debug("Already trained")
            return

    # save info regarding the model training in this folder
    model_info_folder = root_info_folder / model_name
    if not model_info_folder.exists():
        model_info_folder.mkdir(parents=True, exist_ok=True)

    # magic to fix the GPUs
    setup_gpus()

    ##########################################################
    #   Load data
    ##########################################################

    # get the words
    words = words_types[hypa["words_type"]]
    num_labels = len(words)

    # load data
    processed_folder = Path("data_proc")
    processed_path = processed_folder / f"{hypa['dataset_name']}"
    data, labels = load_processed(processed_path, words)

    # concatenate train and val for final train
    val_data = None
    if use_validation:
        x = data["training"]
        y = labels["training"]
        val_data = (data["validation"], labels["validation"])
        logg.debug("Using validation data")
    else:
        x = np.concatenate((data["training"], data["validation"]))
        y = np.concatenate((labels["training"], labels["validation"]))
        logg.debug("NOT using validation data")

    ##########################################################
    #   Setup model
    ##########################################################

    # the shape of each sample
    input_shape = data["training"][0].shape

    # from hypa extract model param
    model_param = get_model_param_area(hypa, num_labels, input_shape)

    # get the model with the chosen params
    net_type = hypa["net_type"]
    if net_type == "ARN":
        model = AreaNet.build(**model_param)
    elif net_type == "AAN":
        model = ActualAreaNet.build(**model_param)
    elif net_type == "VAN":
        model = VerticalAreaNet.build(**model_param)
    elif net_type.startswith("SI"):
        if net_type == "SIM":
            sim_type = "1"
        elif net_type == "SI2":
            sim_type = "2"
        model = SimpleNet.build(sim_type=sim_type, **model_param)

    num_samples = x.shape[0]
    logg.debug(f"num_samples: {num_samples}")

    # from hypa extract training param (epochs, batch, opt, ...)
    training_param = get_training_param_area(
        hypa, use_validation, model_path, num_samples
    )

    # a few metrics to track
    metrics = [
        tf.keras.metrics.CategoricalAccuracy(),
        tf.keras.metrics.Precision(),
        tf.keras.metrics.Recall(),
    ]

    # compile the model
    model.compile(
        optimizer=training_param["opt"],
        loss=tf.keras.losses.CategoricalCrossentropy(),
        metrics=metrics,
    )

    # recap
    recap: ty.Dict[str, ty.Any] = {}
    recap["model_name"] = model_name
    recap["words"] = words
    recap["hypa"] = hypa
    recap["model_param"] = model_param
    recap["use_validation"] = use_validation
    recap["batch_size"] = training_param["batch_size"]
    recap["epochs"] = training_param["epochs"]
    recap["lr_name"] = training_param["lr_name"]
    recap["version"] = "002"

    # logg.debug(f"recap: {recap}")
    recap_path = model_info_folder / "recap.json"
    recap_path.write_text(json.dumps(recap, indent=4))

    # https://stackoverflow.com/a/45546663/2237151
    model_summary_path = model_info_folder / "model_summary.txt"
    with model_summary_path.open("w") as msf:
        model.summary(line_length=150, print_fn=lambda x: msf.write(x + "\n"))

    ##########################################################
    #   Fit model
    ##########################################################

    results = model.fit(
        x,
        y,
        validation_data=val_data,
        epochs=training_param["epochs"],
        batch_size=training_param["batch_size"],
        callbacks=training_param["callbacks"],
    )

    ##########################################################
    #   Save results, history, performance
    ##########################################################

    # results_recap
    results_recap: ty.Dict[str, ty.Any] = {}
    results_recap["model_name"] = model_name
    results_recap["results_recap_version"] = "001"

    # evaluate performance
    eval_testing = model.evaluate(data["testing"], labels["testing"])
    for metrics_name, value in zip(model.metrics_names, eval_testing):
        logg.debug(f"{metrics_name}: {value}")
        results_recap[metrics_name] = value

    # confusion matrix
    y_pred = model.predict(data["testing"])
    cm = pred_hot_2_cm(labels["testing"], y_pred, words)
    results_recap["cm"] = cm.tolist()

    # fscore
    fscore = analyze_confusion(cm, words)
    logg.debug(f"fscore: {fscore}")
    results_recap["fscore"] = fscore

    # save the histories
    results_recap["history_train"] = {
        mn: results.history[mn] for mn in model.metrics_names
    }
    if use_validation:
        results_recap["history_val"] = {
            f"val_{mn}": results.history[f"val_{mn}"] for mn in model.metrics_names
        }

    # save the results
    res_recap_path = model_info_folder / "results_recap.json"
    res_recap_path.write_text(json.dumps(results_recap, indent=4))

    # plot the cm
    fig, ax = plt.subplots(figsize=(12, 12))
    plot_confusion_matrix(cm, ax, model_name, words, fscore)
    plot_cm_path = model_info_folder / "test_confusion_matrix.png"
    fig.savefig(plot_cm_path)
    plt.close(fig)

    # save the trained model
    model.save(model_path)

    # save the placeholder
    placeholder_path.write_text(f"Trained. F-score: {fscore}")

def find_best_lr(hypa: ty.Dict[str, str]) -> None:
    """MAKEDOC: what is find_best_lr doing?"""
    logg = logging.getLogger(f"c.{__name__}.find_best_lr")
    # logg.setLevel("INFO")
    logg.debug("Start find_best_lr")

    # get the word list
    words = words_types[hypa["words_type"]]
    num_labels = len(words)

    # no validation just find the LR
    use_validation = False

    # name the model
    model_name = build_area_name(hypa, use_validation)
    logg.debug(f"model_name: {model_name}")

    # load data
    processed_folder = Path("data_proc")
    processed_path = processed_folder / f"{hypa['dataset_name']}"
    data, labels = load_processed(processed_path, words)

    # the shape of each sample
    input_shape = data["training"][0].shape

    # from hypa extract model param
    model_param = get_model_param_area(hypa, num_labels, input_shape)

    # no need for validation
    x = np.concatenate((data["training"], data["validation"]))
    y = np.concatenate((labels["training"], labels["validation"]))

    # magic to fix the GPUs
    setup_gpus()

    # get the model with the chosen params
    net_type = hypa["net_type"]
    if net_type == "ARN":
        model = AreaNet.build(**model_param)
    elif net_type == "AAN":
        model = ActualAreaNet.build(**model_param)
    elif net_type == "VAN":
        model = VerticalAreaNet.build(**model_param)
    elif net_type.startswith("SI"):
        if net_type == "SIM":
            sim_type = "1"
        elif net_type == "SI2":
            sim_type = "2"
        model = SimpleNet.build(sim_type=sim_type, **model_param)

    num_samples = x.shape[0]
    logg.debug(f"num_samples: {num_samples}")

    # from hypa extract training param (epochs, batch, opt, ...)
    training_param = get_training_param_area(
        hypa, use_validation, model_path=None, num_samples=num_samples
    )

    # a few metrics to track
    metrics = [
        tf.keras.metrics.CategoricalAccuracy(),
        tf.keras.metrics.Precision(),
        tf.keras.metrics.Recall(),
    ]

    # compile the model
    model.compile(
        optimizer=training_param["opt"],
        loss=tf.keras.losses.CategoricalCrossentropy(),
        metrics=metrics,
    )

    # boundary values
    start_lr = 1e-9
    end_lr = 1e1

    # find the best values
    lrf = LearningRateFinder(model)
    lrf.find(
        (x, y),
        start_lr,
        end_lr,
        epochs=training_param["epochs"],
        batchSize=training_param["batch_size"],
    )

    fig_title = "LR_sweep"
    fig_title += f"__{model_name}"

    fig, ax = plt.subplots(figsize=(8, 8))

    # get the plot
    lrf.plot_loss(ax=ax, title=fig_title)

    # save the plot
    plot_fol = Path("plot_results") / "area" / "find_best_lr"
    if not plot_fol.exists():
        plot_fol.mkdir(parents=True, exist_ok=True)
    fig_name = fig_title + ".{}"
    fig.savefig(plot_fol / fig_name.format("png"))
    fig.savefig(plot_fol / fig_name.format("pdf"))

    recap_loss = {}
    recap_loss["lrs"] = [float(lr) for lr in lrf.lrs[:]]
    recap_loss["losses"] = [float(loss) for loss in lrf.losses[:]]
    loss_path = plot_fol / f"loss_{fig_title}.json"
    loss_path.write_text(json.dumps(recap_loss, indent=4))

    plt.show()

def run_train_area(args: argparse.Namespace) -> None:
    """MAKEDOC: What is train_area doing?"""
    logg = logging.getLogger(f"c.{__name__}.run_train_area")
    logg.debug("Starting run_train_area")

    training_type = args.training_type
    words_type = args.words_type
    use_validation = args.use_validation
    dry_run = args.dry_run
    force_retrain = args.force_retrain
    do_find_best_lr = args.do_find_best_lr

    if training_type == "hypa_tune":
        hyper_train_area(
            words_type, force_retrain, use_validation, dry_run, do_find_best_lr
        )

if __name__ == "__main__":
    args = setup_env()
    run_train_area(args)


import serial
import csv
from datetime import datetime

ser = serial.Serial(
        # Serial Port to read the data from
        port='COM6',

        #Rate at which the information is shared to the communication channel
        baudrate = 9600,

        #Applying Parity Checking (none in this case)
        parity=serial.PARITY_NONE,

       # Pattern of Bits to be read
        stopbits=serial.STOPBITS_ONE,

        # Total number of bits to be read
        bytesize=serial.EIGHTBITS,

        # Number of serial commands to accept before timing out
        timeout=1
)

# Pause the program for 1 second to avoid overworking the serial port
while 1:
        f = open("log.txt", "a")
        x = str(ser.readline())[2:-5]
        if len(x) != 1:
            now = datetime.now()
            f = open("log.txt", "a")
            print(x)
            f.write(str(now)[:10]+","+str(now)[10:-5]+","+x+"\n")
            f.close()


"""Functions that help you work with QMK keymaps.
"""
import json
import sys
from pathlib import Path
from subprocess import DEVNULL

import argcomplete
from milc import cli
from pygments.lexers.c_cpp import CLexer
from pygments.token import Token
from pygments import lex

import qmk.path
from qmk.keyboard import find_keyboard_from_dir, keyboard_folder
from qmk.errors import CppError
from qmk.info import info_json

# The `keymap.c` template to use when a keyboard doesn't have its own
DEFAULT_KEYMAP_C = """#include QMK_KEYBOARD_H
__INCLUDES__

/* THIS FILE WAS GENERATED!
 *
 * This file was generated by qmk json2c. You may or may not want to
 * edit it directly.
 */
__KEYCODE_OUTPUT_GOES_HERE__

const uint16_t PROGMEM keymaps[][MATRIX_ROWS][MATRIX_COLS] = {
__KEYMAP_GOES_HERE__
};

#if defined(ENCODER_ENABLE) && defined(ENCODER_MAP_ENABLE)
const uint16_t PROGMEM encoder_map[][NUM_ENCODERS][NUM_DIRECTIONS] = {
__ENCODER_MAP_GOES_HERE__
};
#endif // defined(ENCODER_ENABLE) && defined(ENCODER_MAP_ENABLE)

__MACRO_OUTPUT_GOES_HERE__

"""

def _generate_keymap_table(keymap_json):
    lines = []
    for layer_num, layer in enumerate(keymap_json['layers']):
        if layer_num != 0:
            lines[-1] = lines[-1] + ','
        layer = map(_strip_any, layer)
        layer_keys = ', '.join(layer)
        lines.append('\t[%s] = %s(%s)' % (layer_num, keymap_json['layout'], layer_keys))
    return lines

def _generate_encodermap_table(keymap_json):
    lines = []
    for layer_num, layer in enumerate(keymap_json['encoders']):
        if layer_num != 0:
            lines[-1] = lines[-1] + ','
        encoder_keycode_txt = ', '.join([f'ENCODER_CCW_CW({_strip_any(e["ccw"])}, {_strip_any(e["cw"])})' for e in layer])
        lines.append('\t[%s] = {%s}' % (layer_num, encoder_keycode_txt))
    return lines

def _generate_macros_function(keymap_json):
    macro_txt = [
        'bool process_record_user(uint16_t keycode, keyrecord_t *record) {',
        '    if (record->event.pressed) {',
        '        switch (keycode) {',
    ]

    for i, macro_array in enumerate(keymap_json['macros']):
        macro = []

        for macro_fragment in macro_array:
            if isinstance(macro_fragment, str):
                macro_fragment = macro_fragment.replace('\\', '\\\\')
                macro_fragment = macro_fragment.replace('\r\n', r'\n')
                macro_fragment = macro_fragment.replace('\n', r'\n')
                macro_fragment = macro_fragment.replace('\r', r'\n')
                macro_fragment = macro_fragment.replace('\t', r'\t')
                macro_fragment = macro_fragment.replace('"', r'\"')

                macro.append(f'"{macro_fragment}"')

            elif isinstance(macro_fragment, dict):
                newstring = []

                if macro_fragment['action'] == 'delay':
                    newstring.append(f"SS_DELAY({macro_fragment['duration']})")

                elif macro_fragment['action'] == 'beep':
                    newstring.append(r'"\a"')

                elif macro_fragment['action'] == 'tap' and len(macro_fragment['keycodes']) > 1:
                    last_keycode = macro_fragment['keycodes'].pop()

                    for keycode in macro_fragment['keycodes']:
                        newstring.append(f'SS_DOWN(X_{keycode})')

                    newstring.append(f'SS_TAP(X_{last_keycode})')

                    for keycode in reversed(macro_fragment['keycodes']):
                        newstring.append(f'SS_UP(X_{keycode})')

                else:
                    for keycode in macro_fragment['keycodes']:
                        newstring.append(f"SS_{macro_fragment['action'].upper()}(X_{keycode})")

                macro.append(''.join(newstring))

        new_macro = "".join(macro)
        new_macro = new_macro.replace('""', '')
        macro_txt.append(f'            case QK_MACRO_{i}:')
        macro_txt.append(f'                SEND_STRING({new_macro});')
        macro_txt.append('                return false;')

    macro_txt.append('        }')
    macro_txt.append('    }')
    macro_txt.append('\n    return true;')
    macro_txt.append('};')
    macro_txt.append('')
    return macro_txt

def _generate_keycodes_function(keymap_json):
    """Generates keymap level keycodes.
    """
    lines = []
    lines.append('enum keymap_keycodes {')

    for index, item in enumerate(keymap_json.get('keycodes', [])):
        key = item["key"]
        if index == 0:
            lines.append(f'  {key} = QK_USER_0,')
        else:
            lines.append(f'  {key},')

    lines.append('};')

    for item in keymap_json.get('keycodes', []):
        key = item["key"]
        for alias in item.get("aliases", []):
            lines.append(f'#define {alias} {key}')

    return lines

def template_json(keyboard):
    """Returns a `keymap.json` template for a keyboard.

    If a template exists in `keyboards/<keyboard>/templates/keymap.json` that text will be used instead of an empty dictionary.

    Args:
        keyboard
            The keyboard to return a template for.
    """
    template_file = Path('keyboards/%s/templates/keymap.json' % keyboard)
    template = {'keyboard': keyboard}
    if template_file.exists():
        template.update(json.load(template_file.open(encoding='utf-8')))

    return template

def template_c(keyboard):
    """Returns a `keymap.c` template for a keyboard.

    If a template exists in `keyboards/<keyboard>/templates/keymap.c` that text will be used instead of an empty dictionary.

    Args:
        keyboard
            The keyboard to return a template for.
    """
    template_file = Path('keyboards/%s/templates/keymap.c' % keyboard)
    if template_file.exists():
        template = template_file.read_text(encoding='utf-8')
    else:
        template = DEFAULT_KEYMAP_C

    return template

def _strip_any(keycode):
    """Remove ANY() from a keycode.
    """
    if keycode.startswith('ANY(') and keycode.endswith(')'):
        keycode = keycode[4:-1]

    return keycode

def find_keymap_from_dir():
    """Returns `(keymap_name, source)` for the directory we're currently in.

    """
    relative_cwd = qmk.path.under_qmk_firmware()

    if relative_cwd and len(relative_cwd.parts) > 1:
        # If we're in `qmk_firmware/keyboards` and `keymaps` is in our path, try to find the keyboard name.
        if relative_cwd.parts[0] == 'keyboards' and 'keymaps' in relative_cwd.parts:
            current_path = Path('/'.join(relative_cwd.parts[1:]))  # Strip 'keyboards' from the front

            if 'keymaps' in current_path.parts and current_path.name != 'keymaps':
                while current_path.parent.name != 'keymaps':
                    current_path = current_path.parent

                return current_path.name, 'keymap_directory'

        # If we're in `qmk_firmware/layouts` guess the name from the community keymap they're in
        elif relative_cwd.parts[0] == 'layouts' and is_keymap_dir(relative_cwd):
            return relative_cwd.name, 'layouts_directory'

        # If we're in `qmk_firmware/users` guess the name from the userspace they're in
        elif relative_cwd.parts[0] == 'users':
            # Guess the keymap name based on which userspace they're in
            return relative_cwd.parts[1], 'users_directory'

    return None, None

def keymap_completer(prefix, action, parser, parsed_args):
    """Returns a list of keymaps for tab completion.
    """
    try:
        if parsed_args.keyboard:
            return list_keymaps(parsed_args.keyboard)

        keyboard = find_keyboard_from_dir()

        if keyboard:
            return list_keymaps(keyboard)

    except Exception as e:
        argcomplete.warn(f'Error: {e.__class__.__name__}: {str(e)}')
        return []

    return []

def is_keymap_dir(keymap, c=True, json=True, additional_files=None):
    """Return True if Path object `keymap` has a keymap file inside.

    Args:
        keymap
            A Path() object for the keymap directory you want to check.

        c
            When true include `keymap.c` keymaps.

        json
            When true include `keymap.json` keymaps.

        additional_files
            A sequence of additional filenames to check against to determine if a directory is a keymap. All files must exist for a match to happen. For example, if you want to match a C keymap with both a `config.h` and `rules.mk` file: `is_keymap_dir(keymap_dir, json=False, additional_files=['config.h', 'rules.mk'])`
    """
    files = []

    if c:
        files.append('keymap.c')

    if json:
        files.append('keymap.json')

    for file in files:
        if (keymap / file).is_file():
            if additional_files:
                for additional_file in additional_files:
                    if not (keymap / additional_file).is_file():
                        return False

            return True

def generate_json(keymap, keyboard, layout, layers, macros=None):
    """Returns a `keymap.json` for the specified keyboard, layout, and layers.

    Args:
        keymap
            A name for this keymap.

        keyboard
            The name of the keyboard.

        layout
            The LAYOUT macro this keymap uses.

        layers
            An array of arrays describing the keymap. Each item in the inner array should be a string that is a valid QMK keycode.

        macros
            A sequence of strings containing macros to implement for this keyboard.
    """
    new_keymap = template_json(keyboard)
    new_keymap['keymap'] = keymap
    new_keymap['layout'] = layout
    new_keymap['layers'] = layers
    if macros:
        new_keymap['macros'] = macros

    return new_keymap

def generate_c(keymap_json):
    """Returns a `keymap.c`.

    `keymap_json` is a dictionary with the following keys:

        keyboard
            The name of the keyboard

        layout
            The LAYOUT macro this keymap uses.

        layers
            An array of arrays describing the keymap. Each item in the inner array should be a string that is a valid QMK keycode.

        macros
            A sequence of strings containing macros to implement for this keyboard.
    """
    new_keymap = template_c(keymap_json['keyboard'])
    layer_txt = _generate_keymap_table(keymap_json)
    keymap = '\n'.join(layer_txt)
    new_keymap = new_keymap.replace('__KEYMAP_GOES_HERE__', keymap)

    encodermap = ''
    if 'encoders' in keymap_json and keymap_json['encoders'] is not None:
        encoder_txt = _generate_encodermap_table(keymap_json)
        encodermap = '\n'.join(encoder_txt)
    new_keymap = new_keymap.replace('__ENCODER_MAP_GOES_HERE__', encodermap)

    macros = ''
    if 'macros' in keymap_json and keymap_json['macros'] is not None:
        macro_txt = _generate_macros_function(keymap_json)
        macros = '\n'.join(macro_txt)
    new_keymap = new_keymap.replace('__MACRO_OUTPUT_GOES_HERE__', macros)

    hostlang = ''
    if 'host_language' in keymap_json and keymap_json['host_language'] is not None:
        hostlang = f'#include "keymap_{keymap_json["host_language"]}.h"\n#include "sendstring_{keymap_json["host_language"]}.h"\n'
    new_keymap = new_keymap.replace('__INCLUDES__', hostlang)

    keycodes = ''
    if 'keycodes' in keymap_json and keymap_json['keycodes'] is not None:
        keycodes_txt = _generate_keycodes_function(keymap_json)
        keycodes = '\n'.join(keycodes_txt)
    new_keymap = new_keymap.replace('__KEYCODE_OUTPUT_GOES_HERE__', keycodes)

    return new_keymap

def write_file(keymap_filename, keymap_content):
    keymap_filename.parent.mkdir(parents=True, exist_ok=True)
    keymap_filename.write_text(keymap_content)

    cli.log.info('Wrote keymap to {fg_cyan}%s', keymap_filename)

    return keymap_filename

def write_json(keyboard, keymap, layout, layers, macros=None):
    """Generate the `keymap.json` and write it to disk.

    Returns the filename written to.

    Args:
        keyboard
            The name of the keyboard

        keymap
            The name of the keymap

        layout
            The LAYOUT macro this keymap uses.

        layers
            An array of arrays describing the keymap. Each item in the inner array should be a string that is a valid QMK keycode.
    """
    keymap_json = generate_json(keyboard, keymap, layout, layers, macros=None)
    keymap_content = json.dumps(keymap_json)
    keymap_file = qmk.path.keymaps(keyboard)[0] / keymap / 'keymap.json'

    return write_file(keymap_file, keymap_content)

def write(keymap_json):
    """Generate the `keymap.c` and write it to disk.

    Returns the filename written to.

    `keymap_json` should be a dict with the following keys:
        keyboard
            The name of the keyboard

        keymap
            The name of the keymap

        layout
            The LAYOUT macro this keymap uses.

        layers
            An array of arrays describing the keymap. Each item in the inner array should be a string that is a valid QMK keycode.

        macros
            A list of macros for this keymap.
    """
    keymap_content = generate_c(keymap_json)
    keymap_file = qmk.path.keymaps(keymap_json['keyboard'])[0] / keymap_json['keymap'] / 'keymap.c'

    return write_file(keymap_file, keymap_content)

def locate_keymap(keyboard, keymap):
    """Returns the path to a keymap for a specific keyboard.
    """
    if not qmk.path.is_keyboard(keyboard):
        raise KeyError('Invalid keyboard: ' + repr(keyboard))

    # Check the keyboard folder first, last match wins
    checked_dirs = ''
    keymap_path = ''

    for dir in keyboard_folder(keyboard).split('/'):
        if checked_dirs:
            checked_dirs = '/'.join((checked_dirs, dir))
        else:
            checked_dirs = dir

        keymap_dir = Path('keyboards') / checked_dirs / 'keymaps'

        if (keymap_dir / keymap / 'keymap.c').exists():
            keymap_path = keymap_dir / keymap / 'keymap.c'
        if (keymap_dir / keymap / 'keymap.json').exists():
            keymap_path = keymap_dir / keymap / 'keymap.json'

    if keymap_path:
        return keymap_path

    # Check community layouts as a fallback
    info = info_json(keyboard)

    for community_parent in Path('layouts').glob('*/'):
        for layout in info.get("community_layouts", []):
            community_layout = community_parent / layout / keymap
            if community_layout.exists():
                if (community_layout / 'keymap.json').exists():
                    return community_layout / 'keymap.json'
                if (community_layout / 'keymap.c').exists():
                    return community_layout / 'keymap.c'

def list_keymaps(keyboard, c=True, json=True, additional_files=None, fullpath=False):
    """List the available keymaps for a keyboard.

    Args:
        keyboard
            The keyboards full name with vendor and revision if necessary, example: clueboard/66/rev3

        c
            When true include `keymap.c` keymaps.

        json
            When true include `keymap.json` keymaps.

        additional_files
            A sequence of additional filenames to check against to determine if a directory is a keymap. All files must exist for a match to happen. For example, if you want to match a C keymap with both a `config.h` and `rules.mk` file: `is_keymap_dir(keymap_dir, json=False, additional_files=['config.h', 'rules.mk'])`

        fullpath
            When set to True the full path of the keymap relative to the `qmk_firmware` root will be provided.

    Returns:
        a sorted list of valid keymap names.
    """
    names = set()

    keyboards_dir = Path('keyboards')
    kb_path = keyboards_dir / keyboard

    # walk up the directory tree until keyboards_dir
    # and collect all directories' name with keymap.c file in it
    while kb_path != keyboards_dir:
        keymaps_dir = kb_path / "keymaps"

        if keymaps_dir.is_dir():
            for keymap in keymaps_dir.iterdir():
                if is_keymap_dir(keymap, c, json, additional_files):
                    keymap = keymap if fullpath else keymap.name
                    names.add(keymap)

        kb_path = kb_path.parent

    # Check community layouts as a fallback
    info = info_json(keyboard)

    for community_parent in Path('layouts').glob('*/'):
        for layout in info.get("community_layouts", []):
            cl_path = community_parent / layout
            if cl_path.is_dir():
                for keymap in cl_path.iterdir():
                    if is_keymap_dir(keymap, c, json, additional_files):
                        keymap = keymap if fullpath else keymap.name
                        names.add(keymap)

    return sorted(names)

def _c_preprocess(path, stdin=DEVNULL):
    """ Run a file through the C pre-processor

    Args:
        path: path of the keymap.c file (set None to use stdin)
        stdin: stdin pipe (e.g. sys.stdin)

    Returns:
        the stdout of the pre-processor
    """
    cmd = ['cpp', str(path)] if path else ['cpp']
    pre_processed_keymap = cli.run(cmd, stdin=stdin)
    if 'fatal error' in pre_processed_keymap.stderr:
        for line in pre_processed_keymap.stderr.split('\n'):
            if 'fatal error' in line:
                raise (CppError(line))
    return pre_processed_keymap.stdout

def _get_layers(keymap):  # noqa C901 : until someone has a good idea how to simplify/split up this code
    """ Find the layers in a keymap.c file.

    Args:
        keymap: the content of the keymap.c file

    Returns:
        a dictionary containing the parsed keymap
    """
    layers = list()
    opening_braces = '({['
    closing_braces = ')}]'
    keymap_certainty = brace_depth = 0
    is_keymap = is_layer = is_adv_kc = False
    layer = dict(name=False, layout=False, keycodes=list())
    for line in lex(keymap, CLexer()):
        if line[0] is Token.Name:
            if is_keymap:
                # If we are inside the keymap array
                # we know the keymap's name and the layout macro will come,
                # followed by the keycodes
                if not layer['name']:
                    if line[1].startswith('LAYOUT') or line[1].startswith('KEYMAP'):
                        # This can happen if the keymap array only has one layer,
                        # for macropads and such
                        layer['name'] = '0'
                        layer['layout'] = line[1]
                    else:
                        layer['name'] = line[1]
                elif not layer['layout']:
                    layer['layout'] = line[1]
                elif is_layer:
                    # If we are inside a layout macro,
                    # collect all keycodes
                    if line[1] == '_______':
                        kc = 'KC_TRNS'
                    elif line[1] == 'XXXXXXX':
                        kc = 'KC_NO'
                    else:
                        kc = line[1]
                    if is_adv_kc:
                        # If we are inside an advanced keycode
                        # collect everything and hope the user
                        # knew what he/she was doing
                        layer['keycodes'][-1] += kc
                    else:
                        layer['keycodes'].append(kc)

        # The keymaps array's signature:
        # const uint16_t PROGMEM keymaps[][MATRIX_ROWS][MATRIX_COLS]
        #
        # Only if we've found all 6 keywords in this specific order
        # can we know for sure that we are inside the keymaps array
            elif line[1] == 'PROGMEM' and keymap_certainty == 2:
                keymap_certainty = 3
            elif line[1] == 'keymaps' and keymap_certainty == 3:
                keymap_certainty = 4
            elif line[1] == 'MATRIX_ROWS' and keymap_certainty == 4:
                keymap_certainty = 5
            elif line[1] == 'MATRIX_COLS' and keymap_certainty == 5:
                keymap_certainty = 6
        elif line[0] is Token.Keyword:
            if line[1] == 'const' and keymap_certainty == 0:
                keymap_certainty = 1
        elif line[0] is Token.Keyword.Type:
            if line[1] == 'uint16_t' and keymap_certainty == 1:
                keymap_certainty = 2
        elif line[0] is Token.Punctuation:
            if line[1] in opening_braces:
                brace_depth += 1
                if is_keymap:
                    if is_layer:
                        # We found the beginning of a non-basic keycode
                        is_adv_kc = True
                        layer['keycodes'][-1] += line[1]
                    elif line[1] == '(' and brace_depth == 2:
                        # We found the beginning of a layer
                        is_layer = True
                elif line[1] == '{' and keymap_certainty == 6:
                    # We found the beginning of the keymaps array
                    is_keymap = True
            elif line[1] in closing_braces:
                brace_depth -= 1
                if is_keymap:
                    if is_adv_kc:
                        layer['keycodes'][-1] += line[1]
                        if brace_depth == 2:
                            # We found the end of a non-basic keycode
                            is_adv_kc = False
                    elif line[1] == ')' and brace_depth == 1:
                        # We found the end of a layer
                        is_layer = False
                        layers.append(layer)
                        layer = dict(name=False, layout=False, keycodes=list())
                    elif line[1] == '}' and brace_depth == 0:
                        # We found the end of the keymaps array
                        is_keymap = False
                        keymap_certainty = 0
            elif is_adv_kc:
                # Advanced keycodes can contain other punctuation
                # e.g.: MT(MOD_LCTL | MOD_LSFT, KC_ESC)
                layer['keycodes'][-1] += line[1]

        elif line[0] is Token.Literal.Number.Integer and is_keymap and not is_adv_kc:
            # If the pre-processor finds the 'meaning' of the layer names,
            # they will be numbers
            if not layer['name']:
                layer['name'] = line[1]

        else:
            # We only care about
            # operators and such if we
            # are inside an advanced keycode
            # e.g.: MT(MOD_LCTL | MOD_LSFT, KC_ESC)
            if is_adv_kc:
                layer['keycodes'][-1] += line[1]

    return layers

def parse_keymap_c(keymap_file, use_cpp=True):
    """ Parse a keymap.c file.

    Currently only cares about the keymaps array.

    Args:
        keymap_file: path of the keymap.c file (or '-' to use stdin)

        use_cpp: if True, pre-process the file with the C pre-processor

    Returns:
        a dictionary containing the parsed keymap
    """
    if keymap_file == '-':
        if use_cpp:
            keymap_file = _c_preprocess(None, sys.stdin)
        else:
            keymap_file = sys.stdin.read()
    else:
        if use_cpp:
            keymap_file = _c_preprocess(keymap_file)
        else:
            keymap_file = keymap_file.read_text(encoding='utf-8')

    keymap = dict()
    keymap['layers'] = _get_layers(keymap_file)
    return keymap

def c2json(keyboard, keymap, keymap_file, use_cpp=True):
    """ Convert keymap.c to keymap.json

    Args:
        keyboard: The name of the keyboard

        keymap: The name of the keymap

        layout: The LAYOUT macro this keymap uses.

        keymap_file: path of the keymap.c file

        use_cpp: if True, pre-process the file with the C pre-processor

    Returns:
        a dictionary in keymap.json format
    """
    keymap_json = parse_keymap_c(keymap_file, use_cpp)

    dirty_layers = keymap_json.pop('layers', None)
    keymap_json['layers'] = list()
    for layer in dirty_layers:
        layer.pop('name')
        layout = layer.pop('layout')
        if not keymap_json.get('layout', False):
            keymap_json['layout'] = layout
        keymap_json['layers'].append(layer.pop('keycodes'))

    keymap_json['keyboard'] = keyboard
    keymap_json['keymap'] = keymap
    return keymap_json


# Tyler Sharretts
# Week 07: SciPy Assignment

### Set-up

# Import
import matplotlib.pyplot as plt
import cartopy
import cartopy.feature as cf
import cartopy.crs as ccrs
import xarray as xr
import numpy as np

# Load in nc file dataset
OceanSST = xr.open_dataset("/Users/tshar/Documents/GEOG_673/geog473-673-master/geog473-673-master/datasets/OR_ABI-L2-SSTF-M3_G16_s20192081300453_e20192081400161_c20192081406297.nc")
OceanSST

# Check out DQF
OceanSST['DQF']

# Define new SST
# Use DQF as the threshold
# Converting all DQF values equal to 0 to nans and masking them
OceanDQF = OceanSST['SST'].where(OceanSST['DQF'] == 0)
OceanDQF

### Plotting

# Plot using Xarray function
fig = plt.figure(figsize = (16, 12))
ax = fig.add_subplot(1, 1, 1, projection = ccrs.PlateCarree())
ax.set_extent((-100, -60, 16, 52), crs = ccrs.PlateCarree())
im = ax.pcolormesh(OceanSST['longitude'].values, OceanSST['latitude'].values, OceanSST['DQF'][0].values,
                   cmap = 'jet', transform = ccrs.PlateCarree(), vmin = 0, vmax = 3)
ax.add_feature(cf.COASTLINE)
ax.add_feature(cf.BORDERS)
plt.colorbar(im, label = "Data Quality Flag Grade")
plt.xticks(np.arange(np.min(OceanSST['longitude'].values), np.max(OceanSST['longitude'].values), 10))
plt.yticks(np.arange(np.min(OceanSST['latitude'].values), np.max(OceanSST['latitude'].values), 10))
plt.xlabel("Longitude")
plt.ylabel("Latitude")
plt.title("Data Quality Flag Grade")


"""OpenAPI v3 Specification"""

# apispec via OpenApi
from apispec import APISpec
from apispec.ext.marshmallow import MarshmallowPlugin
from apispec_webframeworks.flask import FlaskPlugin
from marshmallow import Schema, fields

# Create an APISpec
spec = APISpec(
    title="String Chopper",
    version="1.0.0",
    openapi_version="3.0.3",
    plugins=[FlaskPlugin(), MarshmallowPlugin()],
)

# Define schemas
class InputSchema(Schema):
    raw_string = fields.String(description="A string to be processed", required=True)

class OutputSchema(Schema):
    chopped_string = fields.String(description="A string containing every third letter from the raw_string parameter")

class ErrorSchema(Schema):
    message = fields.String(description="An error message")

# register schemas with spec
spec.components.schema("Input", schema=InputSchema)
spec.components.schema("Output", schema=OutputSchema)

# add swagger tags that are used for endpoint annotation
tags = [
            {"name": "string chopper",
             "description": "Funtion for processing a string"
            },
       ]

for tag in tags:
    print(f"Adding tag: {tag['name']}")
    spec.tag(tag)


from telethon import TelegramClient
import time
import branca.colormap as cmp
from datetime import date, timedelta
from natasha import (
    Segmenter,
    MorphVocab,
    NewsEmbedding,
    NewsMorphTagger,
    NewsSyntaxParser,
    NewsNERTagger,
    Doc)
import re
import folium
import pandas as pd
import geopandas as gpd
from geopandas.tools import geocode
import csv
import geopy.distance
from shapely.geometry import Point
from branca.element import Template, MacroElement

# данные для телеграмм бота
api_id = 'your api_id'
api_hash = 'your api_hash'
client = TelegramClient('anon', api_id, api_hash, system_version="4.16.30-vxCUSTOM")

date_of_post=date.fromisoformat('2020-01-01')

yesterday = date.today() - timedelta(days=1)
yesterday.strftime('%m%d%y')
print(yesterday)

list2=[]
news = []
async def main():
    n=0

    # You can print the message history of any chat:
    async for message in client.iter_messages('@papagaz', reverse = True, offset_date = yesterday):
        mess = message.sender.username, message.text, message.date
        m = str(mess[1]) + str(message.date) #+ ';' +'\n'
        news.append(m)
    print(news)
with client:
    client.loop.run_until_complete(main())
#
# ---------------------------------------------------------------------------------------------------------------
#
def main():
#  разпор полученных сообщений
    segmenter = Segmenter()

    emb = NewsEmbedding()
    morph_tagger = NewsMorphTagger(emb)
    syntax_parser = NewsSyntaxParser(emb)

#     #импортируем данные
    words = r"your path...words_true.txt".replace('\\','/')
    w1 = open(words, 'r', encoding="utf-8", errors='ignore')
    wtrue = w1.read().split(',')
    news_need = []
    n=0
    #сопоставляем данные со словарями
    for i in news:
        new = str(i)
        for s in wtrue:
            if new.find(str(s).lower()) != -1:
                # print("True")
                n=n+1
                news_need.append(new)

    print('done')

    #NLP анализ
    sentt = []
    if len(news_need)>0:
        for sen in news_need[0].split('\n\n'):
            if re.findall(r"\xa0", sen):
                sen = sen.replace("\xa0", " ")
            sentt.append(sen)

        sen = []
        loc = []
        loc1 = []
        loc2 = []

        print('done2')
        for i in news_need:
            print('----------------------')
            doc = Doc(i)
            doc.segment(segmenter)
            doc.tag_morph(morph_tagger)
            doc.parse_syntax(syntax_parser)

            ner_tagger = NewsNERTagger(emb)
            doc.tag_ner(ner_tagger)

            morph_vocab = MorphVocab()
            for span in doc.spans:
                span.normalize(morph_vocab)

            temp1=[]
            for j in doc.spans[:]:
                if j.type =='LOC':
                    temp1.append(j.normal)

            if len(temp1)>1:
                sen.append(i)

                loc1.append(temp1[0])
                loc2.append(temp1[1])

            elif len(temp1)==1:
                sen.append(i)
                loc1.append(temp1[0])
                loc2.append(temp1[0])

        print('done3')

        # геокодирование
        zipped_loc = zip(loc1,loc2, sen)
        zipped_list_loc= list(zipped_loc)
        df_loc = pd.DataFrame(zipped_list_loc, columns=["11", "22",'3'])
        lon_1 = []
        n='Nan'
        lon_2 = []
        news1=[]
        news2=[]
        for index, row in df_loc.iterrows():
            dataframe = geocode(row[0] , provider="nominatim" , user_agent = 'my_request')
            if str(dataframe.iloc[0,1]) == 'None':
                lon_1.append(n)
                news1.append( row[2])
                pass

            else:
                point = dataframe.geometry.iloc[0]
                lon_1.append (str(point.y) + ',' + str(point.x))
                news1.append( row[2])

        for index, row in df_loc.iterrows():
            dataframe = geocode(row[1] , provider="nominatim" , user_agent = 'my_request')
            if str(dataframe.iloc[0,1]) == 'None':
                lon_2.append(n)
                news2.append( row[2])
                pass

            else:
                point = dataframe.geometry.iloc[0]
                lon_2.append (str(point.y) + ',' +str(point.x))
                news2.append( row[2])

        print(loc1, lon_1, loc2, lon_2)
        loc11 = list(zip(loc1, lon_1))
        print(loc11)
        loc22 = list(zip(loc2, lon_2))
        print(loc22)
        dict = {'news1':news1, 'name1': loc1, 'lon1': lon_1, 'news2':news2, 'name2': loc2, 'lon2': lon_2}
        df_loc = pd.DataFrame(dict)
        df_loc.to_csv (r'your path... /temp_lonlat2.csv',  'a', index= False )
        print('done3-------------')

    else:
        print('no new events found')

    # пространственный анализ
    df = pd.read_csv(r'your path... /temp_lonlat2.csv')
    neew=[]
    lat=[]
    lon=[]
    st=[]

    for index, row in df.iterrows():
        if row[2] !='Nan' and row[5] !='Nan':
            new1=row[0]
            new2=row[3]
            n1 = row[2]
            n2 =row[5]
            arr = n1.split(',')
            arr2 = n2.split(',')

            dis = geopy.distance.geodesic(arr, arr2).km

            if dis >2000:
                lon.append(arr2[0])
                lat.append(arr2[1])
                neew.append(row[0])
                sta=1
                st.append(sta)
            else:
                lon.append(arr[0])
                lat.append(arr[1])
                neew.append(row[0])
                sta=0
                st.append(sta)

    data1 = gpd.datasets.get_path('naturalearth_lowres')
    gdf = gpd.read_file(data1)
    zipped_values = zip(lat, lon, neew, st)
    zipped_list = list(zipped_values)

    df = pd.DataFrame(zipped_list, columns=["latitude", "longitude", "new", 'status'])

    gdf_points = gpd.GeoDataFrame(
        df, geometry=gpd.points_from_xy( df.latitude,df.longitude)).set_crs("EPSG:4326")
    gdf_points = gdf_points.to_crs(gdf.crs)

    # визуализация
    def color_producer(charec):
        if charec == '1':
            return 'orange'
        elif charec == '0':
            return 'green'
        else:
            return 'yellow'

    loc = "События в нефтрегазовой сфере с 2020 года"
    title_html = '''

{}

                '''.format(loc)

    template = """

Надежность результата привязки

надежный

ненадежный

    """

    data_merged = gpd.sjoin( gdf_points, gdf, how="inner", predicate='intersects')
    data_m = data_merged.groupby(['name'], as_index=False,)['new'].count()

    data_m.columns = ['name', 'count_new']

    data_m.reset_index(inplace= True )
    g = data_m.merge(gdf, how='left', on='name')
    j = gdf.to_json()

    data_m2 = data_merged.groupby('name').agg('count')
    data_m2.reset_index(inplace= True )
    g2 = data_m2.merge(gdf, how='left', on='name')

    mapit = folium.Map( location=[0, 0], zoom_start=2, max_zoom=6, min_zoom=2 )

    step = cmp.StepColormap(
    ['#F2D7D5', '#D98880', '#CD6155','#C0392B'],
    vmin=1, vmax=15,
    index=[3, 5, 9, 15],
    caption='количество событий'
    )

    map_dict = g.set_index('name')['count_new'].to_dict()

    def pol_color(value):
        if int(value)<2:
            return '#F2D7D5'
        elif int(value)<5 and int(value)>=2:
            return '#D98880'
        elif int(value)<9 and int(value)>=5:
            return'#CD6155'
        elif int(value)>=9:
            return '#C0392B'
        else:
            return '#C0392B'

    style_function=lambda feature: {
                    'fillOpacity': 0.5,
                    'weight': 0.5,
                    'color': pol_color( feature['properties']['name'])}

    for _, r in g2.iterrows():
        sim_geo = gpd.GeoSeries(r['geometry_y']).simplify(tolerance=0.001)
        ne = pd.Series(r['new'])
        ht = gpd.GeoDataFrame({'name': ne, 'geometry':sim_geo})
        geo_j = ht.to_json()

        geo_j = folium.GeoJson(data=geo_j,
                            style_function=  style_function)
        geo_j.add_to(mapit)
    step.add_to(mapit)

    for lat , lon, list_loc,charec in zip(  df.latitude , df.longitude, df.new, df.status ): folium.CircleMarker( location=[ lon, lat ], popup= folium.Popup(list_loc, max_width=200, max_height=300), fill_color=color_producer(str(charec)), color=color_producer(str(charec)), radius=5 ).add_to( mapit )
    mapit.get_root().html.add_child(folium.Element(title_html))
    mapit.get_root().html.add_child(folium.Element(template))
    mapit.show_in_browser()
    mapit.save("map.html")

if __name__=="__main__":
    main()


from generator.dataIter import DataIter
import cv2
import numpy as np

class Trainer:
    def __init__(self, config, g_net, d_net, g_train_model, d_train_model, preprocess=True, train_on_gan=False,
                 ):
        self.g_net = g_net
        self.d_net = d_net
        self.g_model = g_train_model
        self.d_model = d_train_model
        self.preprocess = preprocess
        self.config = config
        self.train_on_gan = train_on_gan

    def train(self, img_dir, keypoint_dir, csv_dir):
        # train using image_generator
        if self.preprocess == True:
            dataIter = DataIter(img_dir, keypoint_dir, csv_dir, batch_size=self.config['batch_size'])
        else:
            dataIter = DataIter(img_dir, keypoint_dir, csv_dir, batch_size=self.config['batch_size'])
        train_generator = dataIter.gen_next()
        if self.train_on_gan:
            for epoch in range(self.config['epochs']):
                for step in range(self.config['steps']):
                    g_loss = 0
                    img, pose_img, target_img, target_pose = next(train_generator)
                    x_real = target_img
                    d_loss = self.d_model.train_on_batch([img, target_pose, x_real], None)
                    for i in range(2):
                        g_loss = self.g_model.train_on_batch([img, target_pose, x_real], None)
                        img, pose_img, target_img, target_pose = next(train_generator)
                        x_real = target_img
                    print("Epoch: {}, steps: {}, g_loss: {}, d_loss: {}".format(epoch, step, g_loss, d_loss))
                print("saving model......")
                self.save()
                test_img, test_pose_img, test_target_img, test_target_pose = next(train_generator)
                predicted = self.g_net.predict([test_img, test_target_pose])
                predicted = np.array(predicted * 255, dtype=np.uint8)
                cv2.imshow('pred', predicted[0])
                cv2.waitKey(1000)
                cv2.imshow('target', test_target_img[0])
                cv2.waitKey(1000)
                cv2.imshow('input', test_img[0])
                cv2.waitKey(1000)
                cv2.destroyAllWindows()
        else:
            for epoch in self.config['epochs']:
                for step in self.config['steps']:
                    img, pose_img, target_img, target_pose = next(train_generator)
                    x_real = img
                    # x_fake = self.g_model.predict(x_real)
                    g_loss = self.g_model.train_on_batch([img, target_pose, x_real], None)
                    print("Epoch: {}, steps: {}, g_loss: {}".format(epoch, step, g_loss))

    def save(self):
        self.g_net.save(self.config['save_dir'] + 'g_model.h5')
        self.d_net.save(self.config['save_dir'] + 'd_model.h5')
        print("Model saved Successfully!")