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

# Define the parameters of the problem
np.random.seed(42)


def sample_erlang(k, mean, size):
    scale = mean / k
    return gamma.rvs(a=k, scale=scale, size=size)


def pdf(x, k, mean):
    scale = mean / k
    return gamma.pdf(x, a=k, scale=scale)


def dpdf(x, k, mean):
    beta = k / mean
    coeff = (beta**k) / scipy.special.gamma(k)
    term1 = (k - 1) * x ** (k - 2) * np.exp(-beta * x)
    term2 = beta * x ** (k - 1) * np.exp(-beta * x)
    return coeff * (term1 - term2)


def der_score(x_1, x_3, theta):
    dlogf1 = 2 * x_1 / (theta**2) - 2 / theta
    dlogf3 = -(2 * x_3 / (3 * (1 - theta) ** 2) - 2 / (1 - theta))
    return dlogf1 + dlogf3

# Number of samples for estimation
n_samples = int(1e5)

# Define the grid of \theta values
theta_grid = np.linspace(0.1, 0.9, 9)

# Placeholder to store results
ipa_estimates = []
lr_estimates = []
std_ipa = []
std_lr = []

# Simulation for IPA and LR estimators
for theta in theta_grid:
    X1 = sample_erlang(2, 1, n_samples)
    X2 = sample_erlang(2, 2, n_samples)
    X3 = sample_erlang(2, 3, n_samples)

    # Evaluate the max functions
    path1 = theta * X1 + X2
    path2 = (1 - theta) * X3
    max_path = np.maximum(path1, path2)

    # IPA estimator
    d_max_path = np.where(path1 > path2, X1, -X3)
    ipa_estimates.append(np.mean(d_max_path))
    std_ipa.append(np.std(d_max_path) / np.sqrt(n_samples))

    # LR estimator
    tX1 = theta * X1
    tX3 = (1 - theta) * X3
    lr_grad = max_path * der_score(tX1, tX3, theta)
    lr_estimates.append(np.mean(lr_grad))
    std_lr.append(np.std(lr_grad) / np.sqrt(n_samples))

# Plot estimates
plt.figure(figsize=(10, 6))
plt.plot(theta_grid, ipa_estimates, label="IPA Estimate", marker="o")
plt.plot(theta_grid, lr_estimates, label="LR Estimate", marker="s")
plt.xlabel("Theta (\u03B8)")
plt.ylabel("Estimate")
plt.title("IPA and LR Estimators")
plt.legend()
plt.grid()
plt.show()

# Plot standard deviations
plt.figure(figsize=(10, 6))
plt.plot(theta_grid, std_ipa, label="IPA Standard Deviation", marker="o")
plt.plot(theta_grid, std_lr, label="LR Standard Deviation", marker="s")
plt.xlabel("Theta (\u03B8)")
plt.ylabel("Standard Deviation")
plt.title("Standard Deviations of IPA and LR Estimators")
plt.legend()
plt.grid()
plt.show()


# Stochastic Gradient Descent (SGD)
true_theta_star = 0.6253

# SGD with decreasing step size
sgd_steps_decreasing = []
theta_sgd_decreasing = 0.5  # Initial theta
errors_decreasing = []

max_iter_decreasing = 50000
n_samples = int(1.0e0)
for k in range(max_iter_decreasing):
    X1 = sample_erlang(2, 1, n_samples)
    X2 = sample_erlang(2, 2, n_samples)
    X3 = sample_erlang(2, 3, n_samples)

    path1 = theta_sgd_decreasing * X1 + X2
    path2 = (1 - theta_sgd_decreasing) * X3

    d_max_path = np.where(path1 > path2, X1, -X3)
    grad_estimate = np.mean(d_max_path)

    step_size = 1 / (k + 1)  # Decreasing step size
    theta_sgd_decreasing -= step_size * grad_estimate
    theta_sgd_decreasing = max(
        0, min(1, theta_sgd_decreasing)
    )  # Projection to keep theta in [0, 1]

    sgd_steps_decreasing.append(theta_sgd_decreasing)
    errors_decreasing.append(abs(theta_sgd_decreasing - true_theta_star))

# SGD with fixed step size and geometrically increasing sample size
sgd_steps_fixed = []
theta_sgd_fixed = 0.5  # Initial theta
errors_fixed = []
sample_size = 1000  # Initial sample size
step_size_fixed = 0.01

delta = 1.1  # Multiplication factor at each iteration
max_iter_geometric = int(
    1 + np.log(1 + (delta - 1) * max_iter_decreasing) / np.log(delta)
)  # To keep the total number of samples equal
num_grad_evals_geometric = np.zeros(max_iter_geometric)
for k in range(max_iter_geometric):
    sample_size = int(n_samples * delta**k)  # Geometrically increasing sample size
    num_grad_evals_geometric[k] = sample_size
    X1 = sample_erlang(2, 1, sample_size)
    X2 = sample_erlang(2, 2, sample_size)
    X3 = sample_erlang(2, 3, sample_size)

    path1 = theta_sgd_fixed * X1 + X2
    path2 = (1 - theta_sgd_fixed) * X3

    d_max_path = np.where(path1 > path2, X1, -X3)
    grad_estimate = np.mean(d_max_path)

    theta_sgd_fixed -= step_size_fixed * grad_estimate
    theta_sgd_fixed = max(0, min(1, theta_sgd_fixed))  # Keep theta in [0, 1]

    sgd_steps_fixed.append(theta_sgd_fixed)
    errors_fixed.append(abs(theta_sgd_fixed - true_theta_star))

# Plot SGD progress for fixed step size and increasing sample size
num_grad_evals_decreasing = np.arange(1, max_iter_decreasing + 1) * n_samples
num_grad_evals_geometric = np.cumsum(num_grad_evals_geometric)
plt.figure(figsize=(10, 6))
plt.plot(
    num_grad_evals_decreasing,
    errors_decreasing,
    label="Decreasing Step Size",
    color="b",
)
plt.plot(
    num_grad_evals_geometric,
    errors_fixed,
    label="Fixed Step Size & Increasing Sample Size",
    color="g",
)
plt.xlabel("Number of Gradient Evaluations")
plt.ylabel("Error |\u03B8_k - \u03B8*|")
plt.title("Convergence of SGD")
plt.yscale("log")
plt.xscale("log")
plt.grid()
plt.legend()
plt.show()