"""Part A. Random walkers: a full numpy solution (no Python loops over walkers or steps).

Run:  python3 walk.py      -> plots in ./figures, numerical answers in the console.
"""
from pathlib import Path

import numpy as np
import matplotlib.pyplot as plt

FIG = Path(__file__).parent / "figures"
FIG.mkdir(exist_ok=True)


# ---------------------------------------------------------------- 1. walk()
def walk(M, N, L=1.0, rng=None):
    """M independent walkers, N steps of length L in random directions.

    Returns x, y of shape (M, N+1); column 0 is the origin.
    """
    rng = np.random.default_rng(rng)
    theta = rng.uniform(0.0, 2.0 * np.pi, size=(M, N))   # all angles at once
    x = np.zeros((M, N + 1))
    y = np.zeros((M, N + 1))
    x[:, 1:] = np.cumsum(L * np.cos(theta), axis=1)       # position = sum of steps
    y[:, 1:] = np.cumsum(L * np.sin(theta), axis=1)
    return x, y


# ---------------------------------------------------------------- 5. variants
def _positions(steps):
    """steps: (M, N, d) -> positions (M, N+1, d), starting at zero."""
    M, _, d = steps.shape
    return np.concatenate([np.zeros((M, 1, d)), np.cumsum(steps, axis=1)], axis=1)


def walk_line(M, N, L=1.0, rng=None):
    rng = np.random.default_rng(rng)
    return _positions(L * rng.choice([-1.0, 1.0], size=(M, N, 1)))


def walk_grid(M, N, L=1.0, rng=None):
    rng = np.random.default_rng(rng)
    dirs = L * np.array([[0, 1], [0, -1], [1, 0], [-1, 0]], dtype=float)  # N S E W
    return _positions(dirs[rng.integers(0, 4, size=(M, N))])


def walk_3d(M, N, L=1.0, rng=None):
    rng = np.random.default_rng(rng)
    v = rng.standard_normal((M, N, 3))                   # isotropic distribution
    v /= np.linalg.norm(v, axis=2, keepdims=True)        # -> uniform on the sphere
    return _positions(L * v)


def walk_lazy(M, N, L=1.0, rng=None):
    rng = np.random.default_rng(rng)
    theta = rng.uniform(0.0, 2.0 * np.pi, size=(M, N))
    move = rng.random((M, N)) < 0.5                      # stays put with probability 1/2
    steps = L * np.stack([np.cos(theta), np.sin(theta)], axis=2) * move[..., None]
    return _positions(steps)


def msd(pos):
    """Ensemble <r^2> as a function of step number: shape (N+1,)."""
    return np.mean(np.sum(pos**2, axis=-1), axis=0)


def main():
    L = 1.0

    # ------------------------------------------------------------ 2. 5 trajectories
    x, y = walk(5, 1000, L, rng=7)
    fig, ax = plt.subplots(figsize=(7, 7))
    for i in range(5):
        ax.plot(x[i], y[i], lw=0.8, label=f"walker {i + 1}")
        ax.plot(x[i, -1], y[i, -1], "o", ms=7, mec="k")
    ax.plot(0, 0, "k*", ms=14, label="start")
    ax.set_aspect("equal")
    ax.set_title("5 walkers, N = 1000")
    ax.legend()
    fig.savefig(FIG / "a2_trajectories.png", dpi=150)

    # ------------------------------------------------------------ 3. <r^2>(N)
    M, N = 1000, 1000
    x, y = walk(M, N, L, rng=1)
    r2 = np.mean(x**2 + y**2, axis=0)
    n = np.arange(N + 1)
    slope, intercept = np.polyfit(n, r2, 1)
    slope0 = np.sum(n * r2) / np.sum(n * n)              # fit through the origin
    print(f"[A3] linear fit: <r^2> = {slope:.4f} N + {intercept:.2f};  "
          f"through origin: slope = {slope0:.4f}  (theory L^2 = {L**2})")

    fig, ax = plt.subplots()
    ax.plot(n, r2, label="simulation")
    ax.plot(n, n * L**2, "--", label="theory $NL^2$")
    ax.plot(n, slope * n + intercept, ":", label=f"fit, slope {slope:.3f}")
    ax.set_xlabel("N"); ax.set_ylabel(r"$\langle r^2\rangle$"); ax.legend()
    fig.savefig(FIG / "a3_msd_linear.png", dpi=150)

    # ------------------------------------------------------------ 4. log-log
    alpha, logc = np.polyfit(np.log(n[1:]), np.log(r2[1:]), 1)
    print(f"[A4] log-log slope alpha = {alpha:.4f}, prefactor = {np.exp(logc):.3f}")
    fig, ax = plt.subplots()
    ax.loglog(n[1:], r2[1:], label="simulation")
    ax.loglog(n[1:], n[1:] * L**2, "--", label="$NL^2$")
    ax.set_xlabel("N"); ax.set_ylabel(r"$\langle r^2\rangle$"); ax.legend()
    fig.savefig(FIG / "a4_msd_loglog.png", dpi=150)

    # ------------------------------------------------------------ 5. variants
    fig, ax = plt.subplots()
    for name, f in [("line ±L", walk_line), ("grid NSEW", walk_grid),
                    ("3D sphere", walk_3d), ("lazy (p=1/2)", walk_lazy)]:
        m = msd(f(M, N, L, rng=2))
        k = np.sum(n * m) / np.sum(n * n)
        a, _ = np.polyfit(np.log(n[1:]), np.log(m[1:]), 1)
        print(f"[A5] {name:13s}: <r^2>/N = {k:.4f} L^2,  alpha = {a:.4f}")
        ax.plot(n, m, label=f"{name}: slope {k:.3f}")
    ax.plot(n, n * L**2, "k--", label="$NL^2$")
    ax.plot(n, n * L**2 / 2, "k:", label="$NL^2/2$")
    ax.set_xlabel("N"); ax.set_ylabel(r"$\langle r^2\rangle$"); ax.legend()
    fig.savefig(FIG / "a5_variants.png", dpi=150)


if __name__ == "__main__":
    main()