import matplotlib.pyplot as plt
import numpy as np

"""
Explanation: HydroTrend, Ganges Basin, Q0 Climate Scenario
Created By: Abby Eckland and Irina Overeem, March 2020

About this notebook
This notebook replicates and improves upon simulations originally run by Frances Dunn and Stephen Darby, reported in Darby et al. 2015.
This simulation is driven by climate predictions (daily temperature and precipitation) obtained from the Hadley Centre (HadRM3P) Regional Climate Model. The Q0 realization is utilized in this notebook.
At the end of this notebook, questions are available to help analyze the model outputs.
See the following links for more informaton and resources regarding the model HydroTrend: <br>
https://csdms.colorado.edu/wiki/Model_help:HydroTrend <br>
https://csdms.colorado.edu/wiki/Model:HydroTrend <br>

125 year simulation: 1975-2100
End of explanation
"""


import pymt.models

hydrotrend_GQ0 = pymt.models.Hydrotrend()

# Add directory for output files
config_file, config_folder = "hydro_config.txt", "Ganges"

"""
Explanation: Import the pymt package. Create a new instance. With each new run, it is wise to rename the instance
End of explanation
"""


cat Ganges/HYDRO.IN

cat Ganges/HYDRO0.HYPS

cat Ganges/HYDRO.CLIMATE

"""
Explanation: Note that the following "cat" commands will allow you to view your input files. They are not required to run HydroTrend, but it is good measure to make sure the model you are running is using the correct basin information.
End of explanation
"""


hydrotrend_GQ0.output_var_names

"""
Explanation: In pymt one can always find out what output a model generates by using the .output_var_names method.
End of explanation
"""


hydrotrend_GQ0.initialize(config_file, config_folder)

"""
Explanation: Now we initialize the model with the configure file and in the configure folder.
End of explanation
"""


(
    hydrotrend_GQ0.start_time,
    hydrotrend_GQ0.time,
    hydrotrend_GQ0.end_time,
    hydrotrend_GQ0.time_step,
    hydrotrend_GQ0.time_units,
)

"""
Explanation: This line of code lists time parameters: when, how long, and at what timestep the model simulation will work.
End of explanation
"""


n_days_GQ0 = int(hydrotrend_GQ0.end_time)
q_GQ0 = np.empty(n_days_GQ0)  # river discharge at the outlet
qs_GQ0 = np.empty(n_days_GQ0)  # sediment load at the outlet
cs_GQ0 = np.empty(
    n_days_GQ0
)  # suspended sediment concentration for different grainsize classes at the outlet
qb_GQ0 = np.empty(n_days_GQ0)  # bedload at the outlet

"""
Explanation: This code declares numpy arrays for several important parameters we want to save.
Empty is declaring space in memory for info to go in.
End of explanation
"""


for i in range(n_days_GQ0):
    hydrotrend_GQ0.update()
    q_GQ0[i] = hydrotrend_GQ0.get_value("channel_exit_water__volume_flow_rate")
    qs_GQ0[i] = hydrotrend_GQ0.get_value(
        "channel_exit_water_sediment~suspended__mass_flow_rate"
    )
    cs_GQ0[i] = hydrotrend_GQ0.get_value(
        "channel_exit_water_sediment~suspended__mass_concentration"
    )
    qb_GQ0[i] = hydrotrend_GQ0.get_value(
        "channel_exit_water_sediment~bedload__mass_flow_rate"
    )

"""
Explanation: Here we have coded up the time loop using i as the index.
We update the model with one timestep at the time, untill we reach the end time. 
For each time step, we also get the values for the output parameters we wish to.
End of explanation
"""


plt.plot(q_GQ0, color="blue")
plt.title("Simulated water discharge, Ganges River (1975-2100)", y=1.05)
plt.xlabel("Day in simulation")
plt.ylabel("Water discharge (m^3/s)")
plt.show()

"""
Explanation: Plot water discharge (q)
End of explanation
"""


plt.plot(qs_GQ0, color="tab:brown")
plt.title("Simulated suspended sediment flux, Ganges River (1975-2100)", y=1.05)
plt.xlabel("Day in simulation")
plt.ylabel("Sediment discharge (kg/s)")
plt.show()

"""
Explanation: Plot suspended sediment discharge (qs)
End of explanation
"""


# Reshape data array to find mean yearly water discharge
q_reshape_GQ0 = q_GQ0.reshape(124, 365)
q_mean_rows_GQ0 = np.mean(q_reshape_GQ0, axis=1)
q_y_vals = np.arange(124)

# Plot data, add trendline
plt.plot(q_y_vals, q_mean_rows_GQ0, color="blue")
plt.xlabel("Year in simulation")
plt.ylabel("Discharge (m^3/s)")
plt.title("Simulated Mean Annual Water Discharge, Ganges River (1975-2100)", y=1.05)

z = np.polyfit(q_y_vals.flatten(), q_mean_rows_GQ0.flatten(), 1)
p = np.poly1d(z)
plt.plot(q_y_vals, p(q_y_vals), "r--")
plt.suptitle("y={:.6f}x+{:.6f}".format(z[0], z[1]), y=0.8)
plt.show()

"""
Explanation: Explore mean annual water discharge trends through time
End of explanation
"""


q_GQ0_daily = np.mean(q_reshape_GQ0, axis=0)

plt.plot(q_GQ0_daily, color="blue")
plt.xlabel("Day of Year")
plt.ylabel("Discharge (m^3/s)")
plt.title("Simulated Mean Daily Water Discharge, Ganges River, 1975-2100")
plt.show()

"""
Explanation: Plot mean daily discharge over 125 year period
End of explanation
"""


# Reshape data array to find mean yearly sediment discharge
qs_reshape_GQ0 = qs_GQ0.reshape(124, 365)
qs_mean_rows_GQ0 = np.mean(qs_reshape_GQ0, axis=1)
qs_y_vals = np.arange(124)

# Plot data, add trendline
plt.plot(qs_y_vals, qs_mean_rows_GQ0, color="tab:brown")
plt.xlabel("Year in simulation")
plt.ylabel("Discharge (kg/s)")
plt.title("Simulated Mean Annual Sediment Discharge, Ganges River (1975-2100)", y=1.05)

z = np.polyfit(qs_y_vals.flatten(), qs_mean_rows_GQ0.flatten(), 1)
p = np.poly1d(z)
plt.plot(qs_y_vals, p(qs_y_vals), "r--")
plt.suptitle("y={:.6f}x+{:.6f}".format(z[0], z[1]), y=0.85)
plt.show()

# Plot mean daily sediment discharge over 125 year period
qs_GQ0_daily = np.mean(qs_reshape_GQ0, axis=0)
plt.plot(qs_GQ0_daily, color="tab:brown")
plt.xlabel("Day of Year")
plt.ylabel("Discharge (kg/s)")
plt.title("Simulated Mean Daily Sediment Discharge, Ganges River, 1975-2100")
plt.show()

"""
Explanation: Explore mean annual sediment discharge trends through time
End of explanation
"""


print(
    "Mean Water Discharge = {} {}".format(
        q_GQ0.mean(),
        hydrotrend_GQ0.get_var_units("channel_exit_water__volume_flow_rate"),
    )
)
print(
    "Mean Bedload Discharge = {} {}".format(
        qb_GQ0.mean(),
        hydrotrend_GQ0.get_var_units(
            "channel_exit_water_sediment~bedload__mass_flow_rate"
        ),
    )
)
print(
    "Mean Suspended Sediment Discharge = {} {}".format(
        qs_GQ0.mean(),
        hydrotrend_GQ0.get_var_units(
            "channel_exit_water_sediment~suspended__mass_flow_rate"
        ),
    )
)
print(
    "Mean Suspended Sediment Concentration = {} {}".format(
        cs_GQ0.mean(),
        hydrotrend_GQ0.get_var_units(
            "channel_exit_water_sediment~suspended__mass_concentration"
        ),
    )
)

# Convert qs to MT/year
AnnualQs_GQ0 = (qs_GQ0.mean() * 1e-9) / 3.17098e-8
print(f"Mean Annual Suspended Sediment Discharge = {AnnualQs_GQ0} MT / year")

"""
Explanation: Get important mass balance information about your model run
End of explanation
"""


# Typically presented as a loglog pot
plt.scatter(np.log(q_GQ0), np.log(cs_GQ0), s=5, color="0.7")
plt.title("HydroTrend simulation of 25 year water discharge, Ganges River (1975-2100)")
plt.xlabel("Log River Discharge in m3/s")
plt.ylabel("Log Sediment concentration in kg/m3")
plt.show()

"""
Explanation: Create discharge-sedimentload relationship for this simulation
End of explanation
"""


# work for Q1b:

"""
Explanation: Now let's answer some questions about the above Ganges River simulation in HydroTrend

Exercise 1: How water and sediment discharge is predicted to change over the next century in the Ganges Basin.
Q1a: How does water discharge and suspended sediment discharge in the Ganges River change over the 125 year time period? Describe the general trend.
A1a:
Q1b:  What is the percent change in average water discharge from the beginning to the end of the simulation? You will need to do some calculations in the cell below.
A1b:
End of explanation
"""


# work for Q1c:

"""
Explanation: Q1c:  What about the percent change in sediment discharge over the simulation time period?
A1c:
End of explanation
"""


