#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import numpy as np
import math
import matplotlib.pyplot as plt
from scipy.special import erf
import sobol_new as sn
from matplotlib import rc
from numpy import matlib
import scipy.stats as st
####################################################################
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)

def mc(typess, d,M):
    x = st.uniform.rvs( size = (d,M) )
    data, exact = evaluate(types, x)
    est = np.mean(data)
    err_est = np.std(data)/np.sqrt(M)
    err = np.abs(est-exact)
    return est, err_est, err

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, exact = 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
    err = np.abs(est-exact)
    return est, err_est, err

def evaluate(types, x):
    d = x.shape[0]
    if (types == 1):
        w = 0.5
        c = 9./d
        val = np.cos(2*np.pi*w + c*np.sum(x,0))
        exact = np.real( np.exp(2*np.pi*1j*w) * ((np.exp(1j*c)-1)/(1j*c))**d )
    elif(types == 2):
        w = 0.5
        c = 7.25/d
        val = np.prod(1/(c**(-2) + (x - w)**2), 0)
        exact = (c * (np.arctan(c*(1-w)) + np.arctan(c*w)))**d
    elif(types == 3):
        w = 0.5
        c = 7.03/d
        val = np.exp(-np.sum(c**2 * (x - w)**2,0))
        exact = (np.sqrt(np.pi)*(erf(c*(1-w))+erf(c*w))/(2*c))**d
    elif(types == 4):
        w = 0.5
        c = 2.04/d
        val = np.exp(- np.sum(c*np.abs(x-w),0))
        exact = ((1/c)*(2-np.exp(-c*w)-np.exp(-c*(1-w))))**d
    elif(types == 5):
        w1 = np.pi/4
        w2 = np.pi/5
        c = 4.3/d
        val = np.exp(np.sum(c*x,0))
        val[(x[0,:]>w1) + (x[1,:]>w2)] = 0
        exact = (1/c**d) * (np.exp(c*w1)-1)*(np.exp(c*w2)-1) * \
        (np.exp(c)-1)**(d-2)
    elif(types == 6):
        val = (np.sum(x,0)<=1)
        exact = 1./math.factorial(d)
    
    return val, exact

for types in range(1,7):
    print(types) 
    fig, axes = plt.subplots(nrows = 1, ncols = 2, figsize = (12,4))
    
    Mlist = np.array([50*2**i for i in range(1,11)])
    nM = np.size(Mlist)
    ds = [2,20]
    for ax in axes:
        d=ds[0]
        K = 20

        # Run the Monte Carlo
        mc_err_est = np.zeros(nM)
        mc_err = np.zeros(nM)
        for i in range(nM):
            val, mc_err_est[i], mc_err[i] = mc(types, d, Mlist[i])
        
        # Run the QMC
        qmc_err_est = np.zeros(nM)
        qmc_err = np.zeros(nM)
        for i in range(nM):
            dummy, qmc_err_est[i], qmc_err[i] = qmc(types, d, Mlist[i]/K, K)
        
        ax.loglog(Mlist, mc_err_est, color='red', label = 'MC error estimate')
        ax.loglog(Mlist, mc_err,'--', color='red', label = 'MC error')
        ax.loglog(Mlist, qmc_err_est, color=[0.0,0.8,0.0], label = 'QMC error estimate')
        ax.loglog(Mlist, qmc_err, '--', color=[0.0,0.8,0.0], label = 'QMC error')
        ax.loglog(Mlist, Mlist**-0.5, '--', label = r'$M^{-1/2}$',color='gray')
        ax.loglog(Mlist, Mlist**-1.0, ':', label = r'$M^{-1}$',color='gray')

        ax.grid(True,which='both')
        plt.suptitle(r'Function ' +str(types))   

        d=ds[1]
    ax.legend(fontsize = 11,loc='upper center', 
              bbox_to_anchor=(0.0000001, -0.15), shadow=False, ncol=4)  

    plt.show()