#!/usr/bin/env python
# coding: utf-8

# # Quickstart of SuPy
# 
# This quickstart demonstrates the essential and simplest workflow of `supy` in SUEWS simulation:
# 
# 1. [load input files](#Load-input-files)
# 2. [run simulation](#Run-simulations)
# 3. [examine results](#Examine-results)
# 
# More advanced use of `supy` are available in the [tutorials](./tutorial.rst#tutorials)
# 
# Before we start, we need to load the following necessary packages.

# In[1]:


import matplotlib.pyplot as plt
import supy as sp
import pandas as pd
import numpy as np
from pathlib import Path

get_ipython().run_line_magic('matplotlib', 'inline')


# In[2]:


sp.show_version()


# ## Load input files

# ### For existing SUEWS users:

# First, a path to SUEWS `RunControl.nml` should be specified, which will direct `supy` to locate input files.

# In[3]:


path_runcontrol = Path('../sample_run') / 'RunControl.nml'


# In[4]:


df_state_init = sp.init_supy(path_runcontrol)


# A sample `df_state_init` looks below (note that `.T` is used here to produce a nicer tableform view):

# In[5]:


df_state_init.filter(like='method').T


# Following the convention of SUEWS, `supy` loads meteorological forcing (met-forcing) files at the grid level.

# In[6]:


grid = df_state_init.index[0]
df_forcing = sp.load_forcing_grid(path_runcontrol, grid)
# by default, two years of forcing data are included;
# to save running time for demonstration, we only use one year in this demo
df_forcing=df_forcing.loc['2012'].iloc[1:]


# ### For new users to SUEWS/SuPy:
# 

# To ease the input file preparation, a helper function `load_SampleData` is provided to get the sample input for SuPy simulations 

# In[7]:


df_state_init, df_forcing = sp.load_SampleData()
grid = df_state_init.index[0]
# by default, two years of forcing data are included;
# to save running time for demonstration, we only use one year in this demo
df_forcing=df_forcing.loc['2012'].iloc[1:]


# ### Overview of SuPy input 

# #### `df_state_init`
# 
# `df_state_init` includes model Initial state consisting of:
# 
# * surface characteristics (e.g., albedo, emissivity, land cover fractions, etc.; full details refer to [SUEWS documentation](https://suews-docs.readthedocs.io/en/latest/input_files/SUEWS_SiteInfo/SUEWS_SiteInfo.html))
# 
# * model configurations (e.g., stability; full details refer to [SUEWS documentation](https://suews-docs.readthedocs.io/en/latest/input_files/RunControl/RunControl.html))
# 
# Detailed description of variables in `df_state_init` refers to [SuPy input](../data-structure/supy-io.ipynb#df_state_init:-model-initial-states)

# Surface land cover fraction information in the sample input dataset:

# In[8]:


df_state_init.loc[:,['bldgh','evetreeh','dectreeh']]


# In[9]:


df_state_init.filter(like='sfr')


# #### `df_forcing`
# 
# `df_forcing` includes meteorological and other external forcing information.
# 
# Detailed description of variables in `df_forcing` refers to [SuPy input](../data-structure/supy-io.ipynb#df_forcing:-forcing-data).
# 
# Below is an overview of forcing variables of the sample data set used in the following simulations.

# In[10]:


list_var_forcing = [
    "kdown",
    "Tair",
    "RH",
    "pres",
    "U",
    "rain",
]
dict_var_label = {
    "kdown": "Incoming Solar\n Radiation ($ \mathrm{W \ m^{-2}}$)",
    "Tair": "Air Temperature ($^{\circ}}$C)",
    "RH": r"Relative Humidity (%)",
    "pres": "Air Pressure (hPa)",
    "rain": "Rainfall (mm)",
    "U": "Wind Speed (m $\mathrm{s^{-1}}$)",
}
df_plot_forcing_x = (
    df_forcing.loc[:, list_var_forcing].copy().shift(-1).dropna(how="any")
)
df_plot_forcing = df_plot_forcing_x.resample("1h").mean()
df_plot_forcing["rain"] = df_plot_forcing_x["rain"].resample("1h").sum()

axes = df_plot_forcing.plot(subplots=True, figsize=(8, 12), legend=False,)
fig = axes[0].figure
fig.tight_layout()
fig.autofmt_xdate(bottom=0.2, rotation=0, ha="center")
for ax, var in zip(axes, list_var_forcing):
    _ = ax.set_ylabel(dict_var_label[var])


# ### Modification of SuPy input

# Given `pandas.DataFrame` is the core data structure of SuPy, all operations, including modification, output, demonstration, etc., on SuPy inputs (`df_state_init` and `df_forcing`) can be done using `pandas`-based functions/methods.
# 
# Specifically, for modification, the following operations are essential:

# #### locating data

# Data can be located in two ways, namely:
# 1. by name via [`.loc`](http://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#selection-by-label);
# 2. by position via [`.iloc`](http://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#selection-by-position).

# In[11]:


# view the surface fraction variable: `sfr`
df_state_init.loc[:,'sfr']


# In[12]:


# view the second row of `df_forcing`, which is a pandas Series
df_forcing.iloc[1]


# In[13]:


# view a particular position of `df_forcing`, which is a value
df_forcing.iloc[8,9]


# #### setting new values

# Setting new values is very straightforward: after locating the variables/data to modify, just set the new values accordingly:

# In[14]:


# modify surface fractions
df_state_init.loc[:,'sfr']=[.1,.1,.2,.3,.25,.05,0]
# check the updated values
df_state_init.loc[:,'sfr']


# ## Run simulations

# Once met-forcing (via `df_forcing`) and initial conditions (via `df_state_init`) are loaded in, we call `sp.run_supy` to conduct a SUEWS simulation, which will return two `pandas` `DataFrame`s: `df_output` and `df_state`.

# In[15]:


df_output, df_state_final = sp.run_supy(df_forcing, df_state_init)


# ### `df_output`
# 
# `df_output` is an ensemble output collection of major SUEWS output groups, including:
# 
# * SUEWS: the essential SUEWS output variables
# * DailyState: variables of daily state information
# * snow: snow output variables (effective when `snowuse = 1` set in `df_state_init`)
# 
# Detailed description of variables in `df_output` refers to [SuPy output](../data-structure/supy-io.html#df_output:-model-output-results)

# In[16]:


df_output.columns.levels[0]


# ### `df_state_final`
# 
# `df_state_final` is a `DataFrame` for holding:
# 
# 1. all model states if `save_state` is set to `True` when calling `sp.run_supy` (`supy` may run significantly slower for a large simulation);
# 2. or, only the final state if `save_state` is set to `False` (the default setting), in which mode `supy` has a similar performance as the standalone compiled SUEWS executable.
# 
# Entries in `df_state_final` have the same data structure as `df_state_init` and can thus be used for other SUEWS simulations starting at the timestamp as in `df_state_final`.
# 
# Detailed description of variables in `df_state_final` refers to [SuPy output](../data-structure/supy-io.html#df_state_final:-model-final-states)

# In[17]:


df_state_final.T.head()


# ## Examine results

# Thanks to the functionality inherited from `pandas` and other packages under the [PyData](https://pydata.org) stack, compared with the standard SUEWS simulation workflow, `supy` enables more convenient examination of SUEWS results by statistics calculation, resampling, plotting (and many more).

# ### Ouptut structure
# 
# `df_output` is organised with `MultiIndex` `(grid,timestamp)` and `(group,varaible)` as `index` and `columns`, respectively.

# In[18]:


df_output.head()


# Here we demonstrate several typical scenarios for SUEWS results examination.
# 
# The essential `SUEWS` output collection is extracted as a separate variable for easier processing in the following sections. More [advanced slicing techniques](http://pandas.pydata.org/pandas-docs/stable/advanced.html#multiindex-advanced-indexing) are available in `pandas` documentation.

# In[19]:


df_output_suews = df_output['SUEWS']


# ### Statistics Calculation
# 
# We can use the `.describe()` method for a quick overview of the key surface energy balance budgets.

# In[20]:


df_output_suews.loc[:, ['QN', 'QS', 'QH', 'QE', 'QF']].describe()


# ### Plotting

# #### Basic example
# Plotting is very straightforward via the `.plot` method bounded with `pandas.DataFrame`.
# Note the usage of `loc` for two slices of the output `DataFrame`.

# In[21]:


# a dict for better display variable names
dict_var_disp = {
    'QN': '$Q^*$',
    'QS': r'$\Delta Q_S$',
    'QE': '$Q_E$',
    'QH': '$Q_H$',
    'QF': '$Q_F$',
    'Kdown': r'$K_{\downarrow}$',
    'Kup': r'$K_{\uparrow}$',
    'Ldown': r'$L_{\downarrow}$',
    'Lup': r'$L_{\uparrow}$',
    'Rain': '$P$',
    'Irr': '$I$',
    'Evap': '$E$',
    'RO': '$R$',
    'TotCh': '$\Delta S$',
}


# Quick look at the simulation results:

# In[22]:


ax_output = df_output_suews\
    .loc[grid]\
    .loc['2012 6 1':'2012 6 7',
         ['QN', 'QS', 'QE', 'QH', 'QF']]\
    .rename(columns=dict_var_disp)\
    .plot()
_ = ax_output.set_xlabel('Date')
_ = ax_output.set_ylabel('Flux ($ \mathrm{W \ m^{-2}}$)')
_ = ax_output.legend()


# #### More examples
# Below is a more complete example for examination of urban energy balance over the whole summer (June to August).

# In[23]:


# energy balance
ax_output = (
    df_output_suews.loc[grid]
    .loc["2012 6":"2012 8", ["QN", "QS", "QE", "QH", "QF"]]
    .rename(columns=dict_var_disp)
    .plot(figsize=(10, 3), title="Surface Energy Balance",)
)
_ = ax_output.set_xlabel("Date")
_ = ax_output.set_ylabel("Flux ($ \mathrm{W \ m^{-2}}$)")
_ = ax_output.legend()


# ### Resampling
# 
# The suggested runtime/simulation frequency of SUEWS is `300 s`, which usually results in a large output and may be over-weighted for storage and analysis. 
# Also, you may feel an apparent slowdown in producing the above figure as a large amount of data were used for the plotting.
# To slim down the result size for analysis and output, we can `resample` the default output very easily.

# In[24]:


rsmp_1d = df_output_suews.loc[grid].resample("1d")
# daily mean values
df_1d_mean = rsmp_1d.mean()
# daily sum values
df_1d_sum = rsmp_1d.sum()


# We can then re-examine the above energy balance at hourly scale and plotting will be significantly faster.

# In[25]:


# energy balance
ax_output = (
    df_1d_mean.loc[:, ["QN", "QS", "QE", "QH", "QF"]]
    .rename(columns=dict_var_disp)
    .plot(figsize=(10, 3), title="Surface Energy Balance",)
)
_ = ax_output.set_xlabel("Date")
_ = ax_output.set_ylabel("Flux ($ \mathrm{W \ m^{-2}}$)")
_ = ax_output.legend()


# Then we use the hourly results for other analyses.

# In[26]:


# radiation balance
ax_output = (
    df_1d_mean.loc[:, ["QN", "Kdown", "Kup", "Ldown", "Lup"]]
    .rename(columns=dict_var_disp)
    .plot(figsize=(10, 3), title="Radiation Balance",)
)
_ = ax_output.set_xlabel("Date")
_ = ax_output.set_ylabel("Flux ($ \mathrm{W \ m^{-2}}$)")
_ = ax_output.legend()


# In[27]:


# water balance
ax_output = (
    df_1d_sum.loc[:, ["Rain", "Irr", "Evap", "RO", "TotCh"]]
    .rename(columns=dict_var_disp)
    .plot(figsize=(10, 3), title="Surface Water Balance",)
)
_ = ax_output.set_xlabel("Date")
_ = ax_output.set_ylabel("Water amount (mm)")
_ = ax_output.legend()


# Get an overview of partitioning in energy and water balance at monthly scales:

# In[28]:


# get a monthly Resampler
df_plot = df_output_suews.loc[grid].copy()
df_plot.index = df_plot.index.set_names("Month")
rsmp_1M = df_plot.shift(-1).dropna(how="all").resample("1M", kind="period")
# mean values
df_1M_mean = rsmp_1M.mean()
# sum values
df_1M_sum = rsmp_1M.sum()


# In[29]:


# month names
name_mon = [x.strftime("%b") for x in rsmp_1M.groups]
# create subplots showing two panels together
fig, axes = plt.subplots(2, 1, sharex=True)
# surface energy balance
df_1M_mean.loc[:, ["QN", "QS", "QE", "QH", "QF"]].rename(columns=dict_var_disp).plot(
    ax=axes[0],  # specify the axis for plotting
    figsize=(10, 6),  # specify figure size
    title="Surface Energy Balance",
    kind="bar",
)
# surface water balance
df_1M_sum.loc[:, ["Rain", "Irr", "Evap", "RO", "TotCh"]].rename(
    columns=dict_var_disp
).plot(
    ax=axes[1],  # specify the axis for plotting
    title="Surface Water Balance",
    kind="bar",
)

# annotations
_ = axes[0].set_ylabel("Mean Flux ($ \mathrm{W \ m^{-2}}$)")
_ = axes[0].legend()
_ = axes[1].set_xlabel("Month")
_ = axes[1].set_ylabel("Total Water Amount (mm)")
_ = axes[1].xaxis.set_ticklabels(name_mon, rotation=0)
_ = axes[1].legend()


# ### Output

# The supy output can be saved as `txt` files for further analysis using supy function `save_supy`.

# In[30]:


df_output


# In[33]:


list_path_save = sp.save_supy(df_output, df_state_final,)


# In[32]:


for file_out in list_path_save:
    print(file_out.name)