# -*- coding: utf-8 -*-
"""
Created on Sun Oct 15 16:23:10 2023

@authors: Luise Wanner (TUD Dresden Technical University, luise.wanner@tu-dresden.de), Martin Jung (Max Planck Institute for Biogeochemistry)
 description:
 application of the model of dispersive heat fluxes (Wanner et al. Towards Energy-Balance Closure With a Model of Dispersive Fluxes,
 submitted to Boundary-Layer Meteorology) to CHEESEHEAD19 field measurements using ERA5 reanalyis data
"""

import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestRegressor
path = "insert your path here"

#-----------------------------------------------------------------------------#
#-- SET UP MODEL based on idealized large eddy simulations                  --#
#-----------------------------------------------------------------------------#
#-- read training data
dat = pd.read_csv(path+"Disp_flux_model_training.csv")

#-- extrct predicting variables for both H and LE
featuresInH=["u*/w*","HP","z/z_i","dTa"]
featuresInLE=["u*/w*","HP","z/z_i","dqa"]

X_IN_H=dat[featuresInH].to_numpy()
X_IN_LE=dat[featuresInLE].to_numpy()

#-- set up model for H    
YH=dat['H_d'].to_numpy()

rfH = RandomForestRegressor(n_estimators = 200,n_jobs=20)
rfH.fit(X_IN_H, np.squeeze(YH))

#-- set up model for LE
YLE=dat['LE_d'].to_numpy()

rfLE = RandomForestRegressor(n_estimators = 200,n_jobs=20)
rfLE.fit(X_IN_LE, np.squeeze(YLE))

#-----------------------------------------------------------------------------#
#-- APPLY MODEL to example dataset                                          --#
#-----------------------------------------------------------------------------#
ds = pd.read_csv(path+"example_dataset.csv")

ds['H_pred'] = rfH.predict(ds[featuresInH].to_numpy())
ds['LE_pred'] = LE_pred = rfLE.predict(ds[featuresInLE].to_numpy())

ds.to_csv(path+"example_dataset_with_predicted_dispersive_fluxes.csv",index=False)