# 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

CONV_CRIT = 0.00005  # convergence criterion (fractional change)  []
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  # usAge of this variable is not implemented
sst_prm = 0
A_oil = 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)
cd = 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
        t_skin = sst1

        (
            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,
            A_oil,
            z_over_L,
            zo_m,
            missing,
        )
        u_star_mag = (u_star[0] ** 2 + u_star[1] ** 2) ** 0.5
        cd[t, w] = (u_star_mag / dyn_in_val) ** 2  # calculates the drag coefficient

    z_over_L_list[t] = z_over_L

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

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

for i in range(0, len(temp_list)):
    xdata = windspeed_list
    ydata = cd[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("Drag Coefficient")
# ax.set_title('A_oil = 0.0')
fig_title = "plot_test.png"
plt.savefig("%s" % fig_title)
plt.close()
# print(ydata)
# print(udata)