#!/usr/bin/env python3
"""
example.py — Usage demonstrations for hv_fluctuation_os.

NumPy only.
"""

import numpy as np

from hv_fluctuation_os import (
    MovingHarmonicTrap, TunableQuadratic, DoubleWellTilted,
    LinearSchedule, sample_equilibrium, sample_forward,
    sample_forward_both, jarzynski_estimate, crooks_test,
)


def demo_jarzynski_quadratic():
    print("=" * 72)
    print("Demo 1 — Jarzynski on the tunable quadratic")
    print("=" * 72)
    print()

    system = TunableQuadratic()
    schedule = LinearSchedule(lam_start=1.0, lam_end=4.0, T_total=5.0)
    beta = 1.0
    dt = 1e-3
    n_steps = int(schedule.T_total / dt)

    exact_df = (system.free_energy(schedule.lam_end, beta)
                - system.free_energy(schedule.lam_start, beta))

    print(f"  System:      {system.name}")
    print(f"  λ:           {schedule.lam_start} → {schedule.lam_end}")
    print(f"  T_total:     {schedule.T_total}")
    print(f"  β:           {beta}")
    print(f"  Exact ΔF:    {exact_df:.6f}")
    print()

    rng = np.random.default_rng(0)
    x0 = sample_equilibrium(system, schedule.lam_start, beta,
                             n_traj=10000, rng=rng)
    W = sample_forward(system, schedule, beta, x0, dt, n_steps, rng)
    j = jarzynski_estimate(W, beta)

    print(f"  n_traj:              {j['n']}")
    print(f"  <W>:                 {W.mean():.4f}")
    print(f"  σ_W:                 {W.std():.4f}")
    print(f"  <exp(-βW)>:          {j['mean_exp']:.6f}")
    print(f"  exp(-βΔF):           {np.exp(-beta * exact_df):.6f}")
    print(f"  βΔF_naive:           {beta * j['delta_F_naive']:.4f}")
    print(f"  βΔF_bc:              {beta * j['delta_F_bc']:.4f}")
    print(f"  Exact βΔF:           {beta * exact_df:.4f}")
    print(f"  Error (naive):       "
          f"{abs(j['delta_F_naive'] - exact_df) / exact_df:.2%}")
    print()


def demo_crooks():
    print("=" * 72)
    print("Demo 2 — Crooks fluctuation theorem")
    print("=" * 72)
    print()

    system = TunableQuadratic()
    schedule = LinearSchedule(lam_start=1.0, lam_end=4.0, T_total=5.0)
    beta = 1.0
    dt = 1e-3
    n_steps = int(schedule.T_total / dt)

    rng = np.random.default_rng(0)
    W_f, W_r = sample_forward_both(system, schedule, beta, n_traj=5000,
                                     dt=dt, n_steps=n_steps,
                                     n_equil=0, rng=rng)

    exact_df = (system.free_energy(schedule.lam_end, beta)
                - system.free_energy(schedule.lam_start, beta))

    crooks = crooks_test(W_f, W_r, beta, n_bins=25)

    print(f"  <W_F>:           {W_f.mean():+.4f}   σ: {W_f.std():.4f}")
    print(f"  <W_R>:           {W_r.mean():+.4f}   σ: {W_r.std():.4f}")
    print()
    print(f"  Crooks slope:    {crooks['slope']:+.4f}   (expected β={beta})")
    print(f"  Crooks intercept:{crooks['intercept']:+.4f}")
    print(f"  ΔF from fit:     {crooks['delta_F']:+.4f}   "
          f"(exact {exact_df:+.4f})")
    print(f"  Abs error:       {abs(crooks['delta_F'] - exact_df):.4f}")
    print(f"  Bins used:       {crooks['bins_used']}")
    print()


def demo_convergence():
    print("=" * 72)
    print("Demo 3 — 1/√N convergence of the Jarzynski estimator")
    print("=" * 72)
    print()

    system = TunableQuadratic()
    schedule = LinearSchedule(1.0, 4.0, 5.0)
    beta = 1.0
    dt = 1e-3
    n_steps = int(schedule.T_total / dt)
    exact_df = (system.free_energy(schedule.lam_end, beta)
                - system.free_energy(schedule.lam_start, beta))

    print(f"  Exact βΔF = {beta * exact_df:.4f}")
    print()
    print(f"  {'n_traj':>7s}  {'mean βΔF':>10s}  {'std over seeds':>15s}")
    print("  " + "-" * 36)

    for n_traj in [100, 500, 2000, 10000]:
        dfs = []
        for seed in range(10):
            rng = np.random.default_rng(seed * 1000 + n_traj)
            x0 = sample_equilibrium(system, schedule.lam_start, beta,
                                     n_traj, rng)
            W = sample_forward(system, schedule, beta, x0, dt, n_steps, rng)
            j = jarzynski_estimate(W, beta)
            dfs.append(j['delta_F_naive'])
        dfs = np.array(dfs)
        print(f"  {n_traj:>7d}  {dfs.mean():>+10.4f}  {dfs.std():>15.4f}")
    print()


def demo_double_well():
    print("=" * 72)
    print("Demo 4 — double-well tilt (barrier crossing)")
    print("=" * 72)
    print()

    system = DoubleWellTilted(a=4.0)
    schedule = LinearSchedule(lam_start=0.0, lam_end=1.5, T_total=20.0)
    beta = 1.0
    dt = 5e-4
    n_steps = int(schedule.T_total / dt)

    exact_df = (system.free_energy(schedule.lam_end, beta)
                - system.free_energy(schedule.lam_start, beta))

    print(f"  System:      {system.name}")
    print(f"  λ:           {schedule.lam_start} → {schedule.lam_end}")
    print(f"  Exact ΔF:    {exact_df:.6f}")
    print()

    rng = np.random.default_rng(0)
    x0 = sample_equilibrium(system, schedule.lam_start, beta,
                             n_traj=2000, rng=rng)
    W = sample_forward(system, schedule, beta, x0, dt, n_steps, rng)
    j = jarzynski_estimate(W, beta)

    print(f"  <W>:             {W.mean():+.4f}")
    print(f"  σ_W:             {W.std():.4f}")
    print(f"  βΔF_naive:       {beta * j['delta_F_naive']:+.4f}")
    print(f"  Exact βΔF:       {beta * exact_df:+.4f}")
    print(f"  2nd law:         <W> = {W.mean():.4f} ≥ ΔF = "
          f"{exact_df:.4f}: {W.mean() >= exact_df}")
    print()


def demo_quasi_static():
    print("=" * 72)
    print("Demo 5 — quasi-static limit")
    print("=" * 72)
    print()

    system = TunableQuadratic()
    beta = 1.0
    dt = 1e-3
    exact_df = (system.free_energy(4.0, beta)
                - system.free_energy(1.0, beta))

    print(f"  Exact ΔF = {exact_df:.4f}")
    print()
    print(f"  {'T_total':>8s}  {'mean W':>10s}  {'std W':>10s}  "
          f"{'|W - ΔF|':>10s}")
    print("  " + "-" * 44)

    for T in [1.0, 2.0, 5.0, 20.0, 50.0]:
        schedule = LinearSchedule(1.0, 4.0, T)
        n_steps = int(T / dt)
        rng = np.random.default_rng(42)
        x0 = sample_equilibrium(system, 1.0, beta, 500, rng)
        W = sample_forward(system, schedule, beta, x0, dt, n_steps, rng)
        print(f"  {T:>8.1f}  {W.mean():>+10.4f}  {W.std():>10.4f}  "
              f"{abs(W.mean() - exact_df):>10.4f}")
    print()


def demo_custom_system():
    print("=" * 72)
    print("Demo 6 — custom system (quartic confinement)")
    print("=" * 72)
    print()

    class QuarticTrap:
        name = "quartic_trap"

        def potential(self, x, lam):
            return 0.25 * lam * x ** 4

        def force(self, x, lam):
            return -lam * x ** 3

        def dU_dlam(self, x, lam):
            return 0.25 * x ** 4

        def free_energy(self, lam, beta):
            # Z = ∫ exp(-β λ x^4/4) dx
            # Sub u = (β λ / 4)^{1/4} x, Z = (4/(β λ))^{1/4} · Γ(5/4) · 2
            # F = -(1/β) ln Z
            Z = (4.0 / (beta * lam)) ** 0.25 * 2 * 0.906402477055477  # Γ(5/4)
            return -np.log(Z) / beta

        def equilibrium_std(self, lam, beta):
            # Numerical: compute std via quadrature
            xs = np.linspace(-10, 10, 8001)
            U = self.potential(xs, lam)
            p = np.exp(-beta * U)
            p /= p.sum()
            mean = float(np.sum(xs * p))
            var = float(np.sum((xs - mean) ** 2 * p))
            return var ** 0.5

    system = QuarticTrap()
    schedule = LinearSchedule(lam_start=0.5, lam_end=2.0, T_total=10.0)
    beta = 1.0
    dt = 5e-4
    n_steps = int(schedule.T_total / dt)
    exact_df = (system.free_energy(schedule.lam_end, beta)
                - system.free_energy(schedule.lam_start, beta))

    rng = np.random.default_rng(0)
    x0 = sample_equilibrium(system, schedule.lam_start, beta,
                             n_traj=2000, rng=rng)
    W = sample_forward(system, schedule, beta, x0, dt, n_steps, rng)
    j = jarzynski_estimate(W, beta)

    print(f"  System: {system.name}")
    print(f"  U(x, λ) = (λ/4) x^4")
    print(f"  λ: {schedule.lam_start} → {schedule.lam_end}")
    print()
    print(f"  Exact ΔF:      {exact_df:+.4f}")
    print(f"  <W>:           {W.mean():+.4f}")
    print(f"  βΔF_naive:     {beta * j['delta_F_naive']:+.4f}")
    print(f"  Error:         {abs(beta * j['delta_F_naive'] - beta * exact_df):.4f}")
    print()
    print("  Any system implementing potential / force / dU_dlam /")
    print("  free_energy can be plugged into the engine.")
    print()


def main():
    demo_jarzynski_quadratic()
    demo_crooks()
    demo_convergence()
    demo_double_well()
    demo_quasi_static()
    demo_custom_system()


if __name__ == "__main__":
    main()