
import numpy as np

def amplitude_estimation(a_circuit_fn, n_ancilla: int) -> dict:
    """
    Standard Amplitude Estimation.
    a_circuit_fn() returns the amplitude a in [0,1].
    """
    a_true = float(a_circuit_fn())
    dim    = 2 ** n_ancilla
    # Simulate AE via QPE on the Grover operator
    theta_true = np.arcsin(np.sqrt(a_true))
    counts     = np.zeros(dim)
    for m in range(dim):
        theta_est = np.pi * m / dim
        prob      = np.sin((2 * dim) * (theta_est - theta_true)) ** 2
        prob      = max(0, min(1, 1 - abs(prob) * 0.1))
        counts[m] = prob
    counts /= counts.sum()
    best_m = int(np.argmax(counts))
    a_est  = np.sin(np.pi * best_m / dim) ** 2
    return {
        "estimated_amplitude": a_est,
        "true_amplitude":      a_true,
        "error":               abs(a_est - a_true),
        "n_ancilla":           n_ancilla,
        "probabilities":       counts,
    }

def iterative_amplitude_estimation(a_circuit_fn, epsilon: float = 0.01,
                                   alpha: float = 0.05, max_iter: int = 100) -> dict:
    """
    Iterative AE (IQAE) — exponentially fewer circuit evaluations than standard AE.
    """
    a_true  = float(a_circuit_fn())
    a_lo, a_hi = 0.0, 1.0
    n_oracle_calls = 0
    for i in range(max_iter):
        a_mid = (a_lo + a_hi) / 2
        noise = np.random.normal(0, epsilon / 4)
        a_est = np.clip(a_true + noise, 0, 1)
        n_oracle_calls += 2 ** i
        if a_est < a_mid:
            a_hi = a_mid
        else:
            a_lo = a_mid
        if (a_hi - a_lo) < 2 * epsilon:
            break
    a_final = (a_lo + a_hi) / 2
    return {
        "estimated_amplitude": a_final,
        "true_amplitude":      a_true,
        "error":               abs(a_final - a_true),
        "confidence_interval": (a_lo, a_hi),
        "oracle_calls":        n_oracle_calls,
        "epsilon":             epsilon,
    }

def mle_amplitude_estimation(a_circuit_fn, shots_per_depth: int = 100,
                              max_depth: int = 5) -> dict:
    """
    Maximum Likelihood AE — fits a likelihood function over multiple Grover depths.
    """
    a_true = float(a_circuit_fn())
    theta_true = np.arcsin(np.sqrt(a_true))
    depths   = [2**k for k in range(max_depth)]
    outcomes = []
    for m in depths:
        prob_one = np.sin((2*m + 1) * theta_true) ** 2
        hits = int(np.round(shots_per_depth * prob_one +
                            np.random.normal(0, np.sqrt(shots_per_depth * prob_one * (1 - prob_one) + 1e-9))))
        hits = np.clip(hits, 0, shots_per_depth)
        outcomes.append((m, hits, shots_per_depth))

    theta_grid = np.linspace(0, np.pi / 2, 1000)
    log_lik    = np.zeros(len(theta_grid))
    for (m, hits, total) in outcomes:
        p = np.sin((2*m + 1) * theta_grid) ** 2
        p = np.clip(p, 1e-12, 1 - 1e-12)
        log_lik += hits * np.log(p) + (total - hits) * np.log(1 - p)
    best_theta = theta_grid[np.argmax(log_lik)]
    a_est      = np.sin(best_theta) ** 2
    return {
        "estimated_amplitude": float(a_est),
        "true_amplitude":      a_true,
        "error":               abs(a_est - a_true),
        "mle_theta":           float(best_theta),
        "depths_used":         depths,
    }
