
"""PKTron Algorithms Science — Monte Carlo Methods."""
import numpy as np
_R = np.float32

def estimate_pi(n_samples=500_000, seed=0):
    """Estimate pi via circle-in-square MC. Returns (estimate, error)."""
    rng = np.random.default_rng(seed)
    xy  = rng.random((n_samples,2), dtype=_R if hasattr(np.random.default_rng(0).random((2,)), "dtype") else float).astype(_R)
    inside = (xy[:,0]**2+xy[:,1]**2) <= _R(1)
    est    = float(4*inside.mean())
    return {"pi_estimate": est, "error": abs(est-np.pi), "n_samples": n_samples}

def monte_carlo_integrate(f, a, b, n_samples=100_000, seed=0):
    """
    1D Monte Carlo integration of f over [a,b].
    Returns dict: integral, std_error.
    """
    rng    = np.random.default_rng(seed)
    x      = rng.uniform(a, b, n_samples).astype(_R)
    fx     = np.array([f(xi) for xi in x], dtype=_R)
    result = (b-a)*fx.mean()
    stderr = (b-a)*fx.std()/np.sqrt(n_samples)
    return {"integral": float(result), "std_error": float(stderr),
            "n_samples": n_samples}

def monte_carlo_nd(f, bounds, n_samples=50_000, seed=0):
    """
    N-D MC integration over hyper-box defined by bounds=[(a1,b1),...].
    """
    rng   = np.random.default_rng(seed)
    dim   = len(bounds)
    vol   = float(np.prod([b-a for a,b in bounds]))
    x     = np.column_stack([rng.uniform(a,b,n_samples)
                              for a,b in bounds]).astype(_R)
    vals  = np.array([f(xi) for xi in x], dtype=_R)
    return {"integral": float(vol*vals.mean()),
            "std_error": float(vol*vals.std()/np.sqrt(n_samples))}

def metropolis_hastings(target_log_prob, x0, n_samples=10_000,
                         step_size=0.5, seed=0):
    """
    Metropolis-Hastings MCMC sampler.
    Returns (samples array, acceptance_rate).
    """
    rng     = np.random.default_rng(seed)
    x       = np.asarray(x0, dtype=_R)
    samples = [x.copy()]
    accepted = 0

    for _ in range(n_samples):
        x_prop = x + _R(step_size)*rng.standard_normal(x.shape).astype(_R)
        log_r  = target_log_prob(x_prop) - target_log_prob(x)
        if np.log(rng.random()) < float(log_r):
            x = x_prop
            accepted += 1
        samples.append(x.copy())

    return {"samples": np.array(samples, dtype=_R),
            "acceptance_rate": accepted/n_samples}

def importance_sampling(f, proposal_sample_fn, proposal_log_pdf_fn,
                         target_log_pdf_fn, n_samples=10_000, seed=0):
    """
    Importance sampling estimator of E_target[f].
    """
    rng     = np.random.default_rng(seed)
    samples = np.array([proposal_sample_fn(rng) for _ in range(n_samples)], dtype=_R)
    log_w   = np.array([target_log_pdf_fn(s)-proposal_log_pdf_fn(s)
                         for s in samples], dtype=_R)
    log_w  -= log_w.max()   # numerical stability
    w       = np.exp(log_w)
    w      /= w.sum()
    fvals   = np.array([f(s) for s in samples], dtype=_R)
    return {"estimate": float((w*fvals).sum()),
            "effective_sample_size": float(1/(w**2).sum())}
