# drag coefficient as a function of wind speed and air-sea temperature
# import ../FMT12/MFT5DB_v6 as mft
import MFT23 as mft
import numpy as np
import matplotlib.pyplot as plt


flux_model = -1  # BVW model=0; Smith (1988) = 1; negative numbers use the options below
z0_mom_prm = 6  # https://www.coaps.fsu.edu/~bourassa/MFT_html/ht_adj_docs.php#z0_mom_prm
z0_theta_q_prm = 1  # https://www.coaps.fsu.edu/~bourassa/MFT_html/ht_adj_docs.php#z0_TQ_prm
stable_prm = 0  # https://www.coaps.fsu.edu/~bourassa/MFT_html/ht_adj_docs.php#stable_prm

sst1 = 15.0  # keep SST constant
q = 0.01668835
psfc = 101272.3770478

warn = -1  # warning are given for 1, and hidden for 0    update no warn for 1 and warn for 0
eqv_neut_prm = 0  # output winds are winds rather than equivalent neutral winds
z_wanted = 10.0  # height to which winds, potential temp, and humidity are adjusted
Qnet = 5.0  # usge of this variable is not implemented
sst_prm = 0
oil_fract_area = 0.0
z_over_L = 0.0
zo_m = [0.00001, 0.00001]


dyn_in_prm = 0  # Wind speed, relative to the surface current, m/s
dyn_in_val2 = 0.0
sfc_current1 = 0.0
sfc_current2 = 0.0
ref_ht_wind = 10.0  # Height of the wind observations, m
convect = 0.0  # Convective parameter. Recommended value between 0.7 and 1.25. For details see TOGA NOTES #4 (recommendation comes from no capillary waves)
air_moist_prm = 0  # Specific humidity
air_moist_val = q
sfc_moist_prm = 1  # Relative humidity, fraction
sfc_moist_val = 0.98  # 98% assumed because of salinity obstructing evaporation
salinity = 34.9 / 1000.0  # Salinity, fraction, No salinity in dataset (global average)
ss_prm = 2  # Sea state parameter
ss_val = 28.0  # No sea state data in dataset, can use wave age to set Charnock's constant.
ref_ht_tq = 10.0  # Height of temperature and humidity observations
astab = 1  # Atmospheric stability is calculated
missing = -9999.0

windspeed_list = np.arange(1, 20, 0.2)
temp_list = np.arange(-10.0, 11.0, 1)
ch = np.zeros([len(temp_list), len(windspeed_list)])
z_over_L_list = np.zeros(len(temp_list))

for t, t2 in enumerate(temp_list):
    for w, windspeed in enumerate(windspeed_list):
        # print( 'In main: t = ', t, ' w = ', w )
        pressure = psfc
        dyn_in_val = windspeed
        air_moist_val = q
        t_air = sst1 + t2  # treat as potential temperature
        t_skin = sst1

        if abs(t_air - t_skin) < 0.1:
            astab = 0
            t_air = t_skin + 0.1
        else:
            astab = 1

        # print( 'In main: t = ', t, ' w = ', w, 'dyn_in_val = ', dyn_in_val, dyn_in_val2 )
        (
            count,
            shf,
            lhf,
            tau,
            u_star,
            t_star,
            q_star,
            z_over_L,
            wave_age,
            dom_phase_spd,
            hsig,
            zo_m_out,
            u_at_z,
            t_at_z,
            q_at_z,
        ) = mft.mft_fluxes(
            dyn_in_prm,
            dyn_in_val,
            dyn_in_val2,
            sfc_current1,
            sfc_current2,
            convect,
            pressure,
            air_moist_prm,
            air_moist_val,
            sfc_moist_prm,
            sfc_moist_val,
            salinity,
            ss_prm,
            ss_val,
            t_air,
            sst_prm,
            t_skin,
            ref_ht_wind,
            ref_ht_tq,
            z_wanted,
            astab,
            eqv_neut_prm,
            Qnet,
            warn,
            flux_model,
            z0_mom_prm,
            z0_theta_q_prm,
            stable_prm,
            oil_fract_area,
            z_over_L,
            zo_m,
            missing,
        )
        u_star_mag = (u_star[0] ** 2 + u_star[1] ** 2) ** 0.5
        if abs(t_air - t_skin) < 0.1:
            ch[t, w] = np.nan
        else:
            ch[t, w] = u_star_mag * t_star / dyn_in_val / (t_air + 9.81 * ref_ht_tq / 1004.0 - t_skin)

    z_over_L_list[t] = z_over_L

# Calculate neutral line (overwrite the line where temp_list = 0)
astab = 0
t_skin = sst1
t_air = t_skin
# print( 'Potential Temp diff = ', t_skin - t_air )
for w, windspeed in enumerate(windspeed_list):
    pressure = psfc
    dyn_in_val = windspeed
    air_moist_val = q

    (
        count,
        shf,
        lhf,
        tau,
        u_star,
        t_star,
        q_star,
        z_over_L,
        wave_age,
        dom_phase_spd,
        hsig,
        zo_m_out,
        u_at_z,
        t_at_z,
        q_at_z,
    ) = mft.mft_fluxes(
        dyn_in_prm,
        dyn_in_val,
        dyn_in_val2,
        sfc_current1,
        sfc_current2,
        convect,
        pressure,
        air_moist_prm,
        air_moist_val,
        sfc_moist_prm,
        sfc_moist_val,
        salinity,
        ss_prm,
        ss_val,
        t_air,
        sst_prm,
        t_skin,
        ref_ht_wind,
        ref_ht_tq,
        z_wanted,
        astab,
        eqv_neut_prm,
        Qnet,
        warn,
        flux_model,
        z0_mom_prm,
        z0_theta_q_prm,
        stable_prm,
        oil_fract_area,
        z_over_L,
        zo_m,
        missing,
    )
    u_star_mag = (u_star[0] ** 2 + u_star[1] ** 2) ** 0.5
#    if ( abs( t_air - t_skin) < 0.1 ):
#      ch[t,w]=np.nan
#    else:
#      ch[t,w]=u_star_mag * t_star / dyn_in_val / (t_air - t_skin)
# z_over_L_list[t] = 0.0

# for t, t2 in enumerate(temp_list):
#    print( 'test1: ', t2, ch[t,60] )

ch = np.where(ch < 0.000047, np.nan, ch)

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)

for i in range(0, len(temp_list)):
    xdata = windspeed_list
    ydata = ch[i, :]
    if z_over_L_list[i] < -0.0001:
        ax.plot(xdata, ydata, "r")
    elif z_over_L_list[i] > 0.0001:
        ax.plot(xdata, ydata, "b")
    else:
        ax.plot(xdata, ydata, "k")

ax.set_xlabel("Surface Wind Speed (m/s)")
ax.set_ylabel("Heat Transfer Coefficient")
# ax.set_title('A_oil = 0.0')
fig_title = "plot_test.png"
plt.savefig("%s" % fig_title)
plt.close()
# print(ydata)
# print(udata)