"""Simulate 2D polymer random walks and compare coil shapes with
end-to-end distance statistics."""

import random
import math
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from equi_gel.plot_binary import STYLE_PATH

plt.style.use(STYLE_PATH)

# =========================
# Parameters
# =========================

NUM_WALKS = 50_000
PLOT_SUBSET = 3000

CHOOOSEN_WALK = 177

N_STEPS_LONG = 1000
N_STEPS_SHORT = 500

STEP_SIZE = 1

random.seed(42)


# =========================
# Random walk utilities
# =========================


def generate_walk_path(n_steps, step_size=1):
    """Perform random walk"""
    _x = [0]
    _y = [0]

    moves = [
        (0, step_size),
        (0, -step_size),
        (-step_size, 0),
        (step_size, 0),
    ]

    for _ in range(n_steps):
        dx, dy = random.choice(moves)
        _x.append(_x[-1] + dx)
        _y.append(_y[-1] + dy)

    return _x, _y


def end_to_end_distance(_x, _y):
    """Calculates end to end distance"""
    return math.hypot(_x[-1], _y[-1])


def rayleigh_pdf(r, n_steps, step_size=1):
    """Calculates rayleigh distribution"""
    sigma_sq = n_steps * step_size**2 / 2.0
    return (r / sigma_sq) * np.exp(-(r**2) / (2.0 * sigma_sq))


# =========================
# Simulation
# =========================

paths_short = []
paths_long = []

dist_short = []
dist_long = []

for _ in range(NUM_WALKS):
    xs, ys = generate_walk_path(N_STEPS_SHORT, STEP_SIZE)
    xl, yl = generate_walk_path(N_STEPS_LONG, STEP_SIZE)

    paths_short.append((xs, ys))
    paths_long.append((xl, yl))

    dist_short.append(end_to_end_distance(xs, ys))
    dist_long.append(end_to_end_distance(xl, yl))

dist_short = np.array(dist_short)
dist_long = np.array(dist_long)

# %% Plot
# =========================
# Figure layout
# =========================

fig, (ax_coil, ax_dist) = plt.subplots(1, 2)


# pull colors from style cycle
cycle_colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]

COLOR_LONG_BG = cycle_colors[1]
COLOR_LONG = "darkorange"
COLOR_SHORT = "deepskyblue"
COLOR_SHORT_BG = cycle_colors[0]


# =========================
# Left panel — polymer coils
# =========================

ax_coil.grid(False)

for idx, (x, y) in enumerate(paths_long[:PLOT_SUBSET]):
    if idx == CHOOOSEN_WALK:
        ax_coil.plot(
            x,
            y,
            alpha=0.8,
            color=COLOR_LONG,
            linewidth=0.8,
            label=f"N={N_STEPS_LONG}",
            zorder=11,
        )
    else:
        ax_coil.plot(
            x,
            y,
            alpha=0.8,
            color=COLOR_LONG_BG,
            linewidth=0.4,
        )

for idx, (x, y) in enumerate(paths_short[:PLOT_SUBSET]):
    if idx == CHOOOSEN_WALK:
        ax_coil.plot(
            x,
            y,
            alpha=0.8,
            color=COLOR_SHORT,
            linewidth=0.8,
            label=f"N={N_STEPS_SHORT}",
            zorder=11,
        )
    else:
        ax_coil.plot(
            x,
            y,
            alpha=0.8,
            color=COLOR_SHORT_BG,
            linewidth=0.4,
        )


# slimmer / smaller origin cross
ax_coil.scatter(
    0,
    0,
    marker="x",
    s=40,  # smaller size
    linewidths=0.8,
    color="black",
    zorder=20,
)

ax_coil.set_aspect("equal")
ax_coil.set_xticks([])
ax_coil.set_yticks([])


# =========================
# Right panel — distributions
# =========================

BINS = 20

ax_dist.hist(
    dist_long,
    bins=BINS,
    density=True,
    alpha=0.5,
    color=COLOR_LONG_BG,
    label=rf"$N={N_STEPS_LONG}$",
)

ax_dist.hist(
    dist_short,
    bins=BINS,
    density=True,
    alpha=0.5,
    color=COLOR_SHORT_BG,
    label=rf"$N={N_STEPS_SHORT}$",
)


# Theory curves — solid for both
r_grid = np.linspace(
    0,
    max(dist_long.max(), dist_short.max()) * 1.05,
    500,
)

ax_dist.plot(
    r_grid,
    rayleigh_pdf(r_grid, N_STEPS_LONG, STEP_SIZE),
    linewidth=1.2,
    color=COLOR_LONG_BG,
    label="Theory",
)

ax_dist.plot(
    r_grid,
    rayleigh_pdf(r_grid, N_STEPS_SHORT, STEP_SIZE),
    linewidth=1.2,
    color=COLOR_SHORT_BG,
)

ax_dist.set_xlim(left=0)
ax_dist.set_yticks([])
ax_dist.set_xlabel(r"$\langle r\rangle \;[b]$")
ax_dist.set_ylabel("Probability density [-]")
ax_dist.grid(False)
ax_dist.legend(loc="upper right")

plt.tight_layout()
script_path = Path(__file__).resolve()
figures_dir = script_path.parent / "figures"
figures_dir.mkdir(parents=True, exist_ok=True)

plt.savefig(
    figures_dir / f"{script_path.stem}.png", dpi=600, bbox_inches="tight"
)
# %%
