import numpy as np
from numpy import matlib
from scipy.stats import t as ts
import matplotlib.pyplot as plt
from matplotlib import rc
from aux_plot import density_scatter #auxiliary library to plot scatterdensitites
#
#   imports graphical paramters
#
plt.style.use('ggplot')
rc('font',**{'family':'serif','serif':['Computer Modern Roman'],
     'size' : '12'})
rc('text', usetex=True)
rc('lines', linewidth=2)
plt.rcParams['axes.facecolor']='w'
import matplotlib
latex_preamble = r'\usepackage{amsmath} \usepackage{amssymb}'
matplotlib.rcParams.update({
    'text.usetex': True,
    'text.latex.preamble': latex_preamble
})

np.random.seed(42)

digits=6 #how many digits of accuracy to display
mu = np.array([0,0])
Sigma = np.array([[4,-1],[-1,4]]);
A = np.linalg.cholesky(Sigma);
alpha = 0.05; # 0.95 = 1-alpha
N = int(1E5);
cval = ts.ppf(1-alpha/2,N-1);
alist = np.array([1,3,5,8,10])
y_normal=A@np.random.standard_normal((2,N))
fig,ax= plt.subplots(nrows = 1, ncols = 1, figsize = (6,4))
aa=7
ab=12

ax=density_scatter(y_normal[0,:],y_normal[1,:],sort=True,
                   bins=20,cmap='Spectral')

xy=np.linspace(aa,ab)
ax.fill_between(xy, xy, np.max(xy), color='#539ecd')
ax.fill_between(xy, xy, aa,color='#539ecd')
plt.legend([r'Samples $\mathcal{N}(0,\Sigma)$',r'$A$'],fancybox=True, 
            framealpha=0.0,loc='upper center', bbox_to_anchor=(0.5, -0.05), 
            shadow=False, ncol=2)
plt.show()

#these values can acutally be computed. 
truep = np.array([0.06503208932, 0.001380140995, 0.000002979188004,\
                  2.905959243*10**(-12), 1.217860285*10**(-17)])
na  = np.size(alist);

def psi(x,a):
    return ((x[1,:]>a)*(x[0,:]>a))

def w(xstar,Sigma,x):
    return np.exp(-xstar.T@(np.linalg.solve(Sigma,x)) + 
                  0.5*xstar.T@np.linalg.solve(Sigma,xstar))
#for 2c
def wd(xstar,Sigma,x,d):
    return np.exp(-xstar.T@(np.linalg.solve(Sigma,x)) + 
                  0.5*xstar.T@np.linalg.solve(Sigma,xstar)/d)

for i in range(na):
    a = alist[i]
    print('a = ',a,', N= ',N);
    # set-up importance distribution and weight function
    xstar = np.array([a,a])
    
    # perform sampling   
    xi = np.random.standard_normal((2,int(N)))
    X = matlib.repmat(mu,2,int(N/2)) + A@xi;
    Y = matlib.repmat(xstar,2,int(N/2)) + A@xi;
    # compute mean and standard dev.
    Z = psi(X,a);
    mean_mc = np.mean(Z);
    std_mc  = np.std(Z);
    print('MC confidence interval (alpha=',alpha,' N= ',N,') P  ',
          mean_mc, ' +- ',cval*std_mc/np.sqrt(N));
    print('relative RMSE(MC) = ',round(std_mc/np.sqrt(N)/truep[i],digits));
    print("                                                   ")
    ZZ = psi(Y,a)*w(xstar,Sigma,Y);
    mean_is = np.mean(ZZ);
    std_is = np.std(ZZ);
    print('IS confidence interval (alpha=',alpha,' N= ',N,') P  ', mean_is,
          ' +- ',cval*std_is/np.sqrt(N));
    print('relative RMSE(IS) = ',round(std_is/np.sqrt(N)/truep[i],digits));
    print("                                                   ")
    print("-------------------------------------------------------------\
          ------------------------------------")

############ Plays with \delta
deltalist = np.linspace(0.5,5,100)
ndelta = np.size(deltalist);
for i in range(na):
    
    a = alist[i]
    # set-up importance distribution and weight function
    xstar = np.array([a,a]);
    #perform sampling    
    xi = np.random.standard_normal([2,int(N)])
    errorlist = np.zeros(ndelta)
    stdlist = np.zeros(ndelta)
    for j in range(ndelta):
        delta = deltalist[j];
        Y = np.matlib.repmat(xstar,2,int(N/2)) + A@xi;
        ZZ = psi(Y,a)*wd(xstar,Sigma,Y,delta);        
        mean_is = np.mean(ZZ);
        std_is = np.std(ZZ);
        errorlist[j] = mean_is - truep[i];
        stdlist[j] = std_is;
    
############# Plots        
    plt.clf()
    print('a= ',a)
    plt.subplot(131)
    
    plt.plot(deltalist,abs(errorlist));
    plt.title(r'$\left|\text{Emp. er}\right|$');
    plt.xlabel('$\delta$')
    
    plt.subplot(132)
    
    plt.plot(deltalist,stdlist);
    plt.title('Emp. st. dev.');
    plt.xlabel('$\delta$')
    
    plt.subplot(133)

    plt.plot(deltalist,np.sqrt( (errorlist)**2 + stdlist**2/N )/truep[i]);
    plt.title('Emp. rel. rmse');
    plt.xlabel('$\delta$')
    plt.tight_layout()
    plt.show()