import numpy as np 
import scipy.stats as st 
import matplotlib.pyplot as plt 
import scipy.optimize as so

np.random.seed(42)

alpha = 0.05
eps = 0.1


# Distribution definitions 

### U[-1,1]
a = -1
b = 1
mu_unif = (a+b)/2
var_unif = (b-a)**2/(12)
sigma_unif = np.sqrt(var_unif)
beta_unif = 3*np.sqrt(3)/4  # Provided in problem, tested numerically using MC.

### Pareto (xm=1, gamma=3.1)
xm = 1
gamma = 3.1
mu_pareto = gamma * xm / (gamma - 1)
var_pareto = xm**2 * gamma / ((gamma - 1)**2 * (gamma - 2))
sigma_pareto = np.sqrt(var_pareto)
beta_pareto = 18.79 # Provided in problem, tested numerically using MC.

### Lognormal (mu_l=0,sigma_l=1)
mu_l = 0
sigma_l = 1
mu_logn = np.exp(mu_l + sigma_l**2/2)
var_logn = ( np.exp(sigma_l**2) - 1 ) * np.exp(2*mu_l + sigma_l**2)
sigma_logn = np.sqrt(var_logn)
beta_logn = 6.35 # Provided in problem, tested numerically using MC, matches theoretical value.

# Functions for N bounds
# Chebyshev: N >= var / (alpha * eps^2)
def N_chebyshev(var, alpha, eps):
    return np.ceil(var / (alpha * eps**2))

# CLT normal approximation: N >= (z_{1-alpha/2}^2 * var) / eps^2
def N_clt(var, alpha, eps):
    z = st.norm.ppf(1 - alpha/2)
    return np.ceil((z**2 * var) / (eps**2))

def N_BE(sigma, beta, alpha, eps):
    K_BE = 0.4748
    # We need to solve the equation (see slides):
    # 1 - alpha/2 - Phi(sqrt(N)*eps/sigma) + K_BE * beta / sqrt(N) = 0
    f = lambda N: 1 - alpha/2 - st.norm.cdf(np.sqrt(N)*eps/sigma) + K_BE * beta / np.sqrt(N)
    # Use a root-finding method to find N
    N_solution = so.bisect(f, 10, 1e9)
    return int(np.ceil(N_solution))

vals = {
    'Pareto (xm=1, gamma=3.1)': (mu_pareto, var_pareto, beta_pareto),
    'Lognormal (mu=0,sigma=1)': (mu_logn, var_logn, beta_logn),
    'Uniform [-1,1]': (mu_unif, var_unif, beta_unif)
}

# compute for each distribution
results = {}
for name, (mu, var, beta) in vals.items():
    sigma = np.sqrt(var)
    n_cheb = N_chebyshev(var, alpha, eps)
    n_clt = N_clt(var, alpha, eps)
    n_be = N_BE(sigma, beta, alpha, eps)
    results[name] = (mu, var, sigma, beta, n_cheb, n_clt, n_be)

print(f"{'Distribution':<30} {'mu':<10} {'var':<15} {'sigma':<15} {'beta':<10} {'N_Chebyshev':<15} {'N_CLT':<15} {'N_BE':<10}")
for name in results.keys():
    mu, var, sigma, beta, n_cheb, n_clt, n_be = results[name]
    print(f"{name:<30} {mu:<10.4f} {var:<15.4f} {sigma:<15.4f} {beta:<10.4f} {n_cheb:<15} {n_clt:<15} {n_be}")

print("-" * 80) 
# -- Monte Carlo simulation --
M = 10000  # Number of independent MC approximations
for name in ["Uniform [-1,1]", "Pareto (xm=1, gamma=3.1)", "Lognormal (mu=0,sigma=1)"]:
    mu, var, sigma, beta, n_cheb, n_clt, n_be = results[name]

    # sampler using scipy distributions
    if name == 'Uniform [-1,1]':
       sampler = lambda size: st.uniform.rvs(loc=a, scale=b-a, size=size) 
    elif name == 'Pareto (xm=1, gamma=3.1)':
       sampler = lambda size: st.pareto.rvs(b=gamma, scale=xm, size=size)
    elif name == 'Lognormal (mu=0,sigma=1)':
        sampler = lambda size: st.lognorm.rvs(s=sigma_l, size=size)
    else:
       raise ValueError("Unknown distribution") 

    for technique in ['Chebyshev', 'CLT', 'BE']:
        N_key = f'N_{technique}'
        N = int(results[name][{'Chebyshev':4, 'CLT':5, 'BE':6}[technique]])
        samples = sampler((M, N))
        I_estimates = samples.mean(axis=1)
        prob_empirical = np.mean(np.abs(I_estimates - mu) > eps)
        print(f"{name} using {technique} N={N}: Empirical probability = {prob_empirical:.4f}, Target alpha = {alpha:.4f}")

    print("-" * 80)