#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
from scipy.stats import norm
import matplotlib.pyplot as plt
from matplotlib import rc
from pdb import set_trace
import sobol_new as sn
from numpy import matlib
# ----------------------------------------------
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)

#Defines which function to evaluate 

def alpha_1(x,w,mu=1.0,sigma=2.0):
    N=len(w)
    a=0
    #transforms to the region of interest [-1,1]
    w=-1+2*w
    for n in range(1,N+1):
       a=a+np.cos(n*np.pi*x)*w[n-1]/(n**2.0)
    a=mu+(sigma**2.0/np.pi**2.0)*a
    return a
    
def alpha_2(x,w,mu=1,sigma=4):
    N=len(w)
    a=0
    #transforms to uniform rvs
    w=norm.ppf(w)
    for n in range(1,N+1):
       a=a+np.sin((n-0.5)*np.pi*x)*w[n-1]/((n-0.5)*np.pi)
    a=x+(2**0.5)*a
    return np.exp(a)

def evaluate(types,w,L=1):
    I=w.shape[0]
    x=np.linspace(0,L,I)
    h=x[1]-x[0]
    z=np.zeros(w.shape[1])
    for j in range(w.shape[1]):
        if types==1:
            alpha=alpha_1(x,w[:,j])
        elif types==2:
            alpha=alpha_2(x,w[:,j])
        else:
            raise "type not defined. must be 1 or 2"
        z[j]=2*h*np.sum(1./(1+alpha))
        if np.isnan(z[j])==1:
            set_trace()
    return z        

#------------------------------------------------------------------
#
#defines wheather to use Monte Carlo or QMC
#
#------------------------------------------------------------------
def mc(typess, d,M):
    x = np.random.random( size = (d,M) );
    data= evaluate(types, x);
    est = np.mean(data);
    err_est = np.std(data)/np.sqrt(M);
    return est, err_est

def  qmc(types,d,N,K):
    x = sn.generate_points(int(N),int(d),0)
    x=x.T
    data = np.zeros(K)
    for i in range(K):
        dat = evaluate(types, np.mod(x+matlib.repmat(
                np.random.random(size =(int(d),1)),1,int(N)),1))
        data[i] = np.mean(dat);
    
    est = np.mean(data);
    err_est = 3*np.std(data)/np.sqrt(K); # use 3 in the error estimate : P(sqrt(M)*err/std<3) approx 0
    return est, err_est

# Computes integrals via MC and QMC for an N dimensional random field
NN=[2,4,16,32]
nN=len(NN)

final_est_mc=np.zeros((nN,2))

final_er_mc=np.zeros((nN,2))

for N in NN:
    print(N)
    fig, axes = plt.subplots(nrows = 1, ncols = 2, figsize = (12,4))

    types=1
    for ax in axes:
        Mlist = 50*2**np.arange(1,10);
        nM = np.size(Mlist);    
        mc_err_est = np.zeros(nM);
        qmc_err_est = np.zeros(nM);

        est_mc = np.zeros(nM);
        est_qmc = np.zeros(nM);
        K = 20;
        for i in range(nM):
            est_mc[i], mc_err_est[i] = mc(types, N, Mlist[i]);
            est_qmc[i], qmc_err_est[i] = qmc(types, N, Mlist[i]/K, K);
        
        print('N='+str(N))
        print('type ' +str(types))
        print('mean MC ' ,est_mc[i])
        print('mean QMC ' ,est_qmc[i])

        print('error MC ' ,mc_err_est[i])
        print('error QMC ' ,qmc_err_est[i])
        print('------------------')
        
        ax.loglog(Mlist, mc_err_est, '-', label = 'MC error estimate')
        ax.loglog(Mlist, qmc_err_est, '-', label = 'QMC error estimate')
            
        #fits line to QMC
        
        xx=np.log(Mlist)
        yy=np.log(qmc_err_est)
        
        
        aa,bb=np.polyfit(xx,yy,deg=1)
        
        ax.loglog(Mlist, Mlist**-0.5, '--', label = r'$N^{-1/2}$',color='gray')
        ax.loglog(np.exp(xx), np.exp(aa*xx+bb), ':', label = r'$N^{'+str(round(aa,2))+'}$',color='gray')

        ax.grid(True,which='both')
        ax.legend(fontsize = 11,loc='upper right')    

        plt.suptitle(r'$N=$' +str(N))
        types+=1
    plt.show()