###############################################################################
# SCRIPT TO RUN AN IDEALIZED TEST OF ENKF ON RESCUE MODEL
# AUTHORS: Man-Yau Chan (1,2) , Biing Han Chia (1,3), Li Yang Tan (1,3)
# ORGANIZATIONS: 1. Team NUS_Singapore_Sci, iGEM 2018
#                2. The Pennsylvania State University, Department of Meteorology
#                   and Atmospheric Sciences
#                3. National University of Singapore, Special Programme in
#                   Science
###############################################################################
# PURPOSE: Running script for an idealized test of EnKF
###############################################################################
# DOCUMENTATION: See RESCUE_EnKF_wiki.pdf in the GitHub repository.
#                It should be available in this directory if you downloaded the
#                repository.
###############################################################################
# SECTIONS:
# 1) Simulated truth model run and simulated observation generation
# 2) Generate prior ensemble and plot stats
# 3) Run EnKF and plot stats
###############################################################################
# Date of last update: Oct 16, 2018
###############################################################################

import funclib_enkf as enkf
import funclib_stats as stats
import funclib_model as model
import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt
import numpy as np
import math as m

###############################################################################
# Section 1: Simulated truth model run and simulated obs generation
###############################################################################
# Concentration units: molecules / cell
# Parameters used for truth model run that produced realistic looking mCherry
# and GFP-mCherry concentrations
d_GC = m.log(2)/(5.0)                       # 1 / hr
k_GC = m.log(2)/(10.0/60.0)                 # 1 / hr
d_C  = m.log(2)/(5.0)                       # 1 / hr
k_C  = m.log(2)/(10.0/60.0)                 # 1 / hr
d_P  = m.log(2)/(120./60.0)                 # 1 / hr
d_s  = m.log(2)/(120./60.0)                 # 1 / hr
Keq  = 20                                   # molecules / cell
k_P  = 0.1                                  # 1 / hr
g_s  = 30                                   # molecules / hr

# Parameter list
param_truth = np.array([g_s, d_s, k_P, Keq, d_P, k_C, d_C, k_GC, d_GC])

# Initial reactant concentration list
reactant_truth = np.array( [0.00, 0.00, 100, 0.000, 0.000, 0.000])
dt = 0.1                 # In units of hours
max_t_steps = 5000

# Run the model
[eqbm_truth, flag] = model.run_model(reactant_truth, param_truth, dt, max_t_steps)
eqbm_truth = np.array(eqbm_truth)

# Plot outcome
colors = ['-k','-r','-g','-b','--k','--r']
names=[ 'S','ES','E','P','C','GC']
for i in range(6):
    plt.plot( np.arange( len(eqbm_truth[i])) * dt, eqbm_truth[i], colors[i], label=names[i] )
plt.yscale('log')
plt.ylim([1,5e3])
plt.ylabel('molecules / cell')
plt.xlabel('time (hrs)')
plt.title('Concentration time series of RESCUE components')
plt.legend()
plt.savefig("simulated_truth.png")

truth = eqbm_truth[:,-1]


# Generate observation (simulated laboratory measurements) with 1% error perturbation
# standard deviation
obs_sigma = truth[-2:]*0.01
obs = truth[-2:] + np.random.normal(loc=0, scale = obs_sigma)


###############################################################################
# Section 2: Generate prior ensemble with erroneous parameters
###############################################################################

# UNCOMMENT THESE LINES TO GENERATE A NEW PRIOR ENSEMBLE

# Set up of erroneous parameters
param_guess = param_truth*0.8

# Uncertainty of erroneous parameters
param_sigma = param_guess*0.3

# Initial concentrations (same as truth)
reactant_t0_guess = np.array( [0.00, 0.00, 100, 0.000, 0.000, 0.000])

# Just to shut off the part of the enkf that creates an ensemble fo initial
# concentrations
reactant_t0_sigma = reactant_t0_guess * 0.0

# Number of ensemble members
ens_size = 20000

## Run the prior ensemble to eqbm
#prior_ens = enkf.ensemble_eqbm_runs( param_guess, param_sigma, reactant_t0_guess,
#                                     reactant_t0_sigma, ens_size, dt,
#                                     max_t_steps)
#
## Save ensemble for future use
#np.savetxt( 'prior_ens.txt', prior_ens)

# Import prior ensemble data (If needed)

# COMMENT THIS LINE IF THE PRIOR ENSEMBLE IS BEING GENERATED BY THIS CODE
prior_ens = np.loadtxt('prior_ens.txt')

# Subset the prior ensemble concentration field
ens_size = (prior_ens.shape)[0]
prior_conc_ens = prior_ens[:,-6:]

# Plot prior ens histograms
fig_C = plt.figure( figsize=(12,12) )
stats.conc_hist2d_C( prior_conc_ens, fig_C, truth)
fig_C.subplots_adjust( wspace=0.2, hspace=0.2 )
plt.savefig( 'prior_C_conc_histogram.png', bbox_inches='tight', dpi=300 )
plt.close()

fig_GC = plt.figure( figsize=(12,12) )
stats.conc_hist2d_GC( prior_conc_ens, fig_GC, truth)
fig_GC.subplots_adjust( wspace=0.2, hspace=0.2 )
plt.savefig( 'prior_GC_conc_histogram.png', bbox_inches='tight', dpi=300 )
plt.close()


###############################################################################
# Section 3: Run the EnKF and plot statistics
###############################################################################

# Run EnKF
poste_ens = enkf.enkf_update( prior_ens, obs, obs_sigma )

# Subset only the concentration fields
poste_conc_ens = poste_ens[:,-6:]

# Plot posterior ens
fig_C = plt.figure( figsize=(12,12) )
stats.conc_hist2d_C( poste_conc_ens, fig_C, truth)
fig_C.subplots_adjust( wspace=0.2, hspace=0.2 )
plt.savefig( 'poste_C_conc_histogram.png', bbox_inches='tight', dpi=300 )
plt.close()

fig_GC = plt.figure( figsize=(12,12) )
stats.conc_hist2d_GC( poste_conc_ens, fig_GC, truth)
fig_GC.subplots_adjust( wspace=0.2, hspace=0.2 )
plt.savefig( 'poste_GC_conc_histogram.png', bbox_inches='tight', dpi=300 )
plt.close()

# Compute errors
prior_mean = np.mean(prior_conc_ens, axis=0)
poste_mean = np.mean(poste_conc_ens, axis=0)

print "Errors of prior ensemble avg concentration fields wrt truth"
print (prior_mean - truth)/truth

print ""

print "Errors of posterio ensemble avg concentration fields wrt truth"
print (poste_mean - truth)/truth
print ""

print "Ratio of prior to posterior ensemble average error"
print (prior_mean - truth)/(poste_mean - truth)


