"""TopEdu project 3: How small can a nature reserve be?

Animals in a rectangular reserve multiply at rate r and wander at rate D. Any that cross the edge are lost.
Write the density as u = e^(r t) v. Then v obeys the plain heat equation  dv/dt = D laplacian(v)  with v = 0 on the
edge, and v dies away at some rate lambda. The species survives exactly when r is bigger than lambda.

Stage 1: run this file. It measures lambda for a 1 x 1 and a 2 x 0.5 reserve (the same area) and checks both
against the exact rate  D pi^2 (1/a^2 + 1/b^2). The CHECK line should say PASS.
Stage 2: use real dispersal and growth rates and ask which shapes of a real reserve would keep a species.

Run:  python 03_critical_patch.py  [--show]
Needs: numpy and matplotlib
"""
import os
import sys

import numpy as np
import matplotlib

if "--show" not in sys.argv:
    matplotlib.use("Agg")
import matplotlib.pyplot as plt

OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "output")
os.makedirs(OUT, exist_ok=True)

D = 1.0     # roaming rate (area per unit time)
N = 33      # grid points along each side


def decay_rate(a, b, d=D, n=N):
    """Steps dv/dt = D laplacian(v) in an a by b rectangle, v = 0 on the edge, and measures how fast v dies.

    Returns (measured rate, times, log of the total amount of v). Late on, log(total) falls in a straight line,
    and its slope is minus the decay rate.
    """
    x = np.linspace(0, a, n); y = np.linspace(0, b, n)
    hx, hy = x[1] - x[0], y[1] - y[0]
    X, Y = np.meshgrid(x, y, indexing="ij")
    v = X * (a - X) * Y * (b - Y)                         # a starting population, zero on the edge
    dt = 0.4 / (d * (2 / hx ** 2 + 2 / hy ** 2))         # explicit steps blow up above 0.5 / (...)
    exact_rate = d * np.pi ** 2 * (1 / a ** 2 + 1 / b ** 2)
    final = 2.0 / exact_rate                              # long enough for the slowest shape to take over
    steps = int(final / dt)
    times, logs = [], []
    for step in range(steps + 1):
        if step % max(1, steps // 60) == 0:
            times.append(step * dt); logs.append(np.log(v.sum()))
        lap = np.zeros_like(v)
        lap[1:-1, 1:-1] = ((v[2:, 1:-1] - 2 * v[1:-1, 1:-1] + v[:-2, 1:-1]) / hx ** 2
                           + (v[1:-1, 2:] - 2 * v[1:-1, 1:-1] + v[1:-1, :-2]) / hy ** 2)
        v = v + d * dt * lap
    times, logs = np.array(times), np.array(logs)
    late = times > final / 2
    rate = -np.polyfit(times[late], logs[late], 1)[0]
    return rate, times, logs


def main():
    worst = 0.0
    fig, ax = plt.subplots(figsize=(5.5, 3.8))
    for a, b in [(1.0, 1.0), (2.0, 0.5)]:
        rate, times, logs = decay_rate(a, b)
        exact = D * np.pi ** 2 * (1 / a ** 2 + 1 / b ** 2)
        error = abs(rate - exact) / exact
        worst = max(worst, error)
        print(f"{a:g} x {b:g} reserve: measured decay rate {rate:.3f}, exact {exact:.3f}, error {100 * error:.2f}%")
        print(f"   so the species needs a growth rate r > {rate:.2f} to survive here")
        ax.plot(times, logs - logs[0], label=f"{a:g} x {b:g}")
    print(f"CHECK decay rates against D pi^2 (1/a^2 + 1/b^2) -> {'PASS' if worst < 0.01 else 'FAIL'} (limit 1%)")
    ax.set_xlabel("time"); ax.set_ylabel("log(total), from its start"); ax.legend(title="reserve")
    ax.set_title("Same area, different shapes: who dies faster?")
    fig.tight_layout(); fig.savefig(os.path.join(OUT, "03_decay.png"), dpi=150)

    # Fixed area 1, changing shape: how much harder does survival get as the reserve stretches?
    ratios = np.array([1, 1.5, 2, 3, 4, 6])
    needed = [D * np.pi ** 2 * (1 / (np.sqrt(k)) ** 2 + 1 / (1 / np.sqrt(k)) ** 2) for k in ratios]
    for k, r in zip(ratios, needed):
        print(f"  area 1, length/width {k:g}: needs r > {r:.2f}  ({r / needed[0]:.2f} times the square)")
    print(f"plots saved in {OUT}")

    # ---- YOUR TURN ----
    # 1. The list above uses the formula. Check two of its shapes with decay_rate() before you trust it.
    # 2. A real species: find published values of D (km^2 per year) and r (per year) for one animal or plant.
    #    What is the smallest square reserve that keeps it? Compare with the size of a real reserve.
    # 3. Predict, then test: does an L-shaped reserve do better or worse than a square of the same area?
    #    (Hint: set v to zero on a corner block every step to cut the L out of a rectangle.)
    # 4. Say what the model leaves out: here nothing limits growth, so it tells you "survives or not", not "how many".
    if "--show" in sys.argv:
        plt.show()


if __name__ == "__main__":
    main()