"""
Performance analysis for stochastic modeling of Lorenz 96 and chaotic cavity flow.

The analysis section of "Data-driven modeling of strongly nonlinear chaotic flows 
with non-Gaussian statistics" by H. Arbabi and T. Sapsis
May 2020, arbabiha@gmail.com
"""

import numpy as np
import scipy.io as sio
import timeit
import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt


# custom tools
from sys import path
path.append('./thehood/')
import plotting_tools as pt   
from Lorenz96 import Lorenz96_SDEmodel
from Cavity import Cavity_SDEmodeling



plt.rc('text', usetex=True)
plt.rc('font', family='serif',size=9)
tfs = 10
MyColors=['#377eb8','#d95f02']
SavePath='./analysis_data/'

def compute_var_diag(rho_truth, rho_model, dx, trim_zeros=False):
    """Compute the variance diagnostic of PDF modeling.

    Variance diagnostic is a measure of the difference 
    between the truth distribution and the one generated by the model.
    
    Args:
        rho_truth: the PDF values of data on a uniform grid
        rho_model: the PDF values of data generated by the model on the same grid
        dx: the PDF grid spacing
        trim_zeros: if True, the 0 values for the model data are thrown out.
            Use with caution: using this option could mislead if there are zeors in meaty areas.

    Return:
        the variance diagnostic: Var_{\rho_truth} log(rho_truth/rho_model)
    """
    if trim_zeros:
        rho_truth=rho_truth[rho_model>0]
        rho_model=rho_model[rho_model>0]
    
    log_term = np.log(rho_model/rho_truth)
    
    average = np.average(log_term, weights=rho_truth*dx)
    var_diag= np.average((log_term-average)**2, weights=rho_truth*dx)
    
    return var_diag


def Cavity_analysis_data():
    """Generate model data for various sample sizes and polynomial degrees in cavity problem."""

    Ts = [100,250,500,2500,5000]
    basetag='analysis'

    for t_max in Ts:
        Cavity_SDEmodeling(t_max=t_max,poly_order=2,\
                                basetag=basetag,SavePath=SavePath)

    poly_orders = [1,2,3,4,5]

    for poly_order in poly_orders:
        Cavity_SDEmodeling(t_max=2500,poly_order=poly_order,\
                                basetag=basetag,SavePath=SavePath)




def Cavity_analysis_plot():
    """Analyze convergence and generate the anlaysis figure [6(a)] in the paper. """
    nbins=65
    smoothing_sigma = None


    Ts = [100,250,500,2500,5000]
    errors_sample_size = []

    for tmax in Ts:
        filename = SavePath+'Cavity_pointwise_analysis_T'+str(tmax)+'_Dim=10_ns10_po2'
        Data=sio.loadmat(filename)
        uv_model,uv_truth=Data['uv_model'],Data['uv_truth']
        Sensors = np.squeeze(Data['Sensors'])

        S = list(range(Sensors.shape[0]))
        sum_sensor_error=0

        for sensor in S:
            grid,p_truth=pt.pdf_1d(uv_truth[sensor,:],nx=nbins,smoothing_sigma=smoothing_sigma)
            dx0 = grid[1]-grid[0]
            _,p_model=pt.pdf_1d(uv_model[sensor,:],nx=nbins,smoothing_sigma=smoothing_sigma,
                                MyRange=[np.amin(grid),np.amax(grid)])
            sum_sensor_error= sum_sensor_error+compute_var_diag(p_truth,p_model,dx0,trim_zeros=True)

        errors_sample_size.append(sum_sensor_error/len(S))
        
    poly_orders = [1,2,3,4,5]
    errors_porder=[]

    for poly_order in poly_orders:
        filename = SavePath+'Cavity_pointwise_analysis_T2500_Dim=5_ns10_po'+str(poly_order)
        Data=sio.loadmat(filename)
        uv_model,uv_truth=Data['uv_model'],Data['uv_truth']
        Sensors = np.squeeze(Data['Sensors'])

        S = list(range(Sensors.shape[0]))
        sum_sensor_error=0

        for sensor in S:
            grid,p_truth=pt.pdf_1d(uv_truth[sensor,:],nx=nbins,smoothing_sigma=smoothing_sigma)
            dx0 = grid[1]-grid[0]
            _,p_model=pt.pdf_1d(uv_model[sensor,:],nx=nbins,smoothing_sigma=smoothing_sigma,
                                MyRange=[np.amin(grid),np.amax(grid)])
            sum_sensor_error= sum_sensor_error+compute_var_diag(p_truth,p_model,dx0,trim_zeros=True)

        errors_porder.append(sum_sensor_error/len(S))


    myfig=plt.figure(figsize=[5.5,2.2])
    axw,axx0,axy0=.33,.15,.2

    ax1 = myfig.add_axes([axx0,axy0,axw,.7])
    no_samples=[t_train*10 for t_train in Ts]
    plt.plot(no_samples,errors_sample_size,'-s',color=MyColors[0])
    ax1.tick_params(direction='in')
    plt.yscale('log'),plt.xscale('log')
    plt.yticks([1e-2,3e-2,1e-1],[r'$10^{-2}$',r'$3 \times 10^{-2}$',r'$10^{-1}$'])
    plt.minorticks_off()
    plt.xlabel(r'sample size')
    plt.title(r'$\tilde{e}$',fontsize=tfs)

    ax2 = myfig.add_axes([axx0+.15+axw,axy0,axw,.7])
    ax2.tick_params(direction='in')
    plt.plot(poly_orders,errors_porder,'-s',color=MyColors[0])
    plt.yscale('log')
    plt.yticks([1e-2,5e-2],[r'$ 10^{-2}$',r'$5\times 10^{-2}$'])
    plt.xticks([1,2,3,4,5])
    plt.minorticks_off()
    plt.xlabel(r'polynomial degree')
    plt.title(r'$\tilde{e}$',fontsize=tfs)

    plt.subplots_adjust(bottom=0.35)
    plt.savefig('cavity_analysis.pdf')


def Lorenz_analysis_data():
    """Generate model data for various sample sizes and polynomial degrees in Lorenz problem."""
    ts = [1,10,100,1000,5000]
    
    for t_max in ts:
        Lorenz96_SDEmodel(t_max=t_max, poly_order=3, SavePath=SavePath)
    
    poly_orders = [1,2,3,4,5]

    for po in poly_orders:
        Lorenz96_SDEmodel(t_max=1000, poly_order=po, SavePath=SavePath)

def Lorenz_analysis_plot():
    """Analyze convergence and generate the anlaysis figure [6(b)] in the paper. """

    x_truth = sio.loadmat('./mat files/Lorenz96Data.mat')['y'][:,0]
    nbins = 66
    smoothing_sigma = None
    grid,p_truth=pt.pdf_1d(x_truth,nx=nbins,smoothing_sigma=smoothing_sigma)
    xl = [np.amin(grid),np.amax(grid)] 
    dx0 = grid[1]-grid[0]


    Ts=[1,10,100,1000,5000]
    Ns,error_ss=[],[]

    for T in Ts:
        filename=SavePath+'LorenzModel_r3_t{}_ns1.mat'.format(T)
        data = sio.loadmat(filename)
        x_model = data['x_model'].squeeze()
        _,p_model=pt.pdf_1d(x_model,nx=nbins,smoothing_sigma=smoothing_sigma,MyRange=xl)
        
        Ns.append(data['x_train'].shape[0])
        error_ss.append(compute_var_diag(p_truth,p_model,dx0,trim_zeros=True))

    rs=[1,2,3,4,5]
    error_r=[]

    for r in rs:
        filename=SavePath+'LorenzModel_r{}_t1000_ns1.mat'.format(r)
        data = sio.loadmat(filename)
        x_model = data['x_model'].squeeze()
        _,p_model=pt.pdf_1d(x_model,nx=nbins,smoothing_sigma=smoothing_sigma,MyRange=xl)
        error_r.append(compute_var_diag(p_truth,p_model,dx0,trim_zeros=True))


    myfig=plt.figure(figsize=[5.5,2.2])
    axw,axx0,axy0=.33,.15,.2

    ax1 = myfig.add_axes([axx0,axy0,axw,.7])
    ax1.tick_params(direction='in')
    plt.plot(Ns,error_ss,'-s',color=MyColors[0])
    plt.yscale('log'), plt.xscale('log')
    plt.xlabel(r'sample size')
    plt.minorticks_off()
    plt.title(r'$e$',fontsize=tfs)

    ax2 = myfig.add_axes([axx0+.15+axw,axy0,axw,.7])
    ax2.tick_params(direction='in')
    plt.plot(rs,error_r,'-s',color=MyColors[0])
    plt.yscale('log')
    plt.xticks(rs)
    plt.xlabel(r'polynomial degree')
    plt.yticks([1e-2,5e-3],[r'$ 10^{-2}$',r'$5\times 10^{-3}$'])

    plt.xticks(rs)
    plt.minorticks_off()
    plt.title(r'$e$',fontsize=tfs)
    plt.subplots_adjust(bottom=0.35)
    plt.savefig('lorenz_analysis.pdf',dpi=450)




if __name__ == "__main__":
    Lorenz_analysis_data()
    Lorenz_analysis_plot()
    # Cavity_analysis_data() # this one takes a while
    # Cavity_analysis_plot()