import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import norm

np.random.seed(42)

# Define auxiliary functions
def phi(x):
    """Standard normal PDF."""
    return norm.pdf(x)

def I_theta_derivative(theta):
    """Analytical derivative of I(theta)."""
    return phi(1 / theta) / theta**2

def I_epsilon_analytical_derivative(theta, epsilon):
    """Analytical derivative of I_epsilon(theta)."""
    denominator = (
        np.exp(1 / (2 * (epsilon**2 + theta**2))) *
        epsilon *
        np.sqrt(2 * np.pi) *
        (epsilon**2 + theta**2) *
        np.sqrt(1 + theta**2 / epsilon**2)
    )
    return theta / denominator

def I_epsilon_approximation(theta, epsilon, n_samples=10000):
    """Monte Carlo approximation of I_epsilon(theta)."""
    X = np.random.normal(0, 1, n_samples)
    smoothed_term = norm.cdf((theta * X - 1) / epsilon)
    return np.mean(smoothed_term)

def I_epsilon_derivative_approximation(theta, epsilon, n_samples=10000):
    """IPA approximaiton of the derivative of I_epsilon(theta)."""
    X = np.random.normal(0, 1, n_samples)
    phi_term = phi((theta * X - 1) / epsilon)
    return np.mean(phi_term * X / epsilon)

def likelihood_ratio_method(theta, n_samples=10000):
    """Likelihood Ratio (LR) method for dI/dtheta."""
    X = np.random.normal(0, 1, n_samples)
    indicator = (theta * X > 1)
    term1 = -1 / theta
    term2 = X**2 / theta
    return np.mean(indicator * (term1 + term2))

# Parameters
theta = 2
epsilons = np.logspace(0, -2, 5)  # Log-spaced grid for epsilon from 1 to 10^-4
n_samples_grid = np.logspace(1, 7, 10, dtype=int)  # Log-spaced grid for sample sizes

# Compute analytical values
analytical_value = I_theta_derivative(theta)
smoothed_analytical_values = [I_epsilon_analytical_derivative(theta, epsilon) for epsilon in epsilons]
analytical_errors = [abs(smoothed - analytical_value) for smoothed in smoothed_analytical_values]

plt.figure(figsize=(12, 8))
colors = plt.cm.viridis(np.linspace(0, 1, len(epsilons)))

# Loop over epsilon values
for i, epsilon in enumerate(epsilons):
    approximation_errors = []
    # Compute IPA approximations
    for n_samples in n_samples_grid:
        approx = I_epsilon_derivative_approximation(theta, epsilon, n_samples)
        approximation_errors.append(abs(approx-analytical_value))

    # Plot IPA approximation errors
    plt.plot(n_samples_grid, approximation_errors, label=f"IPA ε = {epsilon:.1e}", color=colors[i])
    plt.axhline(abs(smoothed_analytical_values[i]-analytical_value), color=colors[i], linestyle='--', linewidth=1,
                label=f"IPA Analytical Smoothed (ε={epsilon:.1e})")

# Compute LR approximaitons
LR_approximation_errors = []
for n_samples in n_samples_grid:
    approx = likelihood_ratio_method(theta, n_samples)
    LR_approximation_errors.append(abs(approx-analytical_value))

# Plot LR approximaiton errors
plt.plot(n_samples_grid, LR_approximation_errors, label=f"LR")

# Customize the plot
plt.xscale('log')
plt.yscale('log')
plt.xlabel("Number of Samples (log scale)")
plt.ylabel("Derivative Approximation Error (log scale)")
plt.title(f"Comparison of Analytical and Approximate Derivatives\nfor θ = {theta}")
plt.legend(loc="best", fontsize='small')
plt.grid(True, which="both", linestyle="--", linewidth=0.5)
plt.tight_layout()
plt.savefig('../figures/smoothed-ipa-approximaitons-theta{}.pdf'.format(theta))
# plt.show()