"""TopEdu project 6: Why earthquakes pick certain buildings

A valley of soft sediment sits in rock. Shear waves bounce between the rock walls and floor and the free ground
surface, so the valley rings at a natural period. Buildings whose own sway period is close to it shake hardest.

This model is a rectangle W wide and H deep. Ground motion u obeys the wave equation with u = 0 at the rock (the
two sides and the floor) and no stress at the surface. Its slowest ring has period
    T = 2 pi / omega,   omega = c pi sqrt(1/W^2 + 1/(2H)^2).
Engineers' quick estimate ignores the width:  T = 4H / c.

Stage 1: run this file. It shakes the basin, listens at the middle of the surface, and checks the period it hears
against the exact one. The CHECK line should say PASS.
Stage 2: map where the quick estimate is good enough, and compare with a real, documented basin.

This is a first explanation, not a hazard estimate: never present it as a prediction for a real place.
Run:  python 06_basin_resonance.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)

W, H, C = 400.0, 40.0, 80.0    # basin width (m), sediment depth (m), shear wave speed (m/s)


def exact_period(w, h, c):
    return 2 * np.pi / (c * np.pi * np.sqrt(1 / w ** 2 + 1 / (2 * h) ** 2))


def simulate(w=W, h=H, c=C, nx=49, ny=17, periods=24):
    """Leapfrog steps of u_tt = c^2 laplacian(u). Rows y = 0 (rock floor) and the side columns stay at 0; the top
    row y = H is a free surface, done with a mirror point. Returns the time step and motion at the surface centre."""
    x = np.linspace(0, w, nx); y = np.linspace(0, h, ny)
    hx, hy = x[1] - x[0], y[1] - y[0]
    X, Y = np.meshgrid(x, y, indexing="ij")
    dt = 0.9 / (c * np.sqrt(1 / hx ** 2 + 1 / hy ** 2))
    steps = int(periods * 4 * h / c / dt)

    def laplacian(u):
        p = np.pad(u, ((0, 0), (0, 1)))
        p[:, -1] = u[:, -2]                               # mirror above the surface: no stress there
        out = np.zeros_like(u)
        out[1:-1, 1:] = ((p[2:, 1:-1] - 2 * p[1:-1, 1:-1] + p[:-2, 1:-1]) / hx ** 2
                         + (p[1:-1, 2:] - 2 * p[1:-1, 1:-1] + p[1:-1, :-2]) / hy ** 2)
        return out

    u_old = (np.sin(np.pi * X / w) + 0.4 * np.sin(3 * np.pi * X / w)) * np.sin(np.pi * Y / (2 * h))
    u = u_old + 0.5 * (c * dt) ** 2 * laplacian(u_old)
    heard = [u_old[nx // 2, -1], u[nx // 2, -1]]
    for _ in range(steps):
        u, u_old = 2 * u - u_old + (c * dt) ** 2 * laplacian(u), u
        heard.append(u[nx // 2, -1])
    return dt, np.array(heard)


def strongest_period(signal, dt):
    """Period of the strongest frequency, refined between FFT bins with a parabola."""
    signal = (signal - signal.mean()) * np.hanning(len(signal))
    size = 16 * len(signal)
    spectrum = np.abs(np.fft.rfft(signal, size)); freqs = np.fft.rfftfreq(size, dt)
    i = np.argmax(spectrum[1:]) + 1
    a, b, c = np.log(spectrum[i - 1:i + 2])
    f = freqs[i] + 0.5 * (a - c) / (a - 2 * b + c) * (freqs[1] - freqs[0])
    return 1 / f


def main():
    dt, heard = simulate()
    period = strongest_period(heard, dt)
    exact = exact_period(W, H, C)
    error = abs(period - exact) / exact
    print(f"basin {W:g} m wide, {H:g} m deep, c = {C:g} m/s")
    print(f"simulated period {period:.3f} s, exact {exact:.3f} s, quick estimate 4H/c = {4 * H / C:.3f} s")
    print(f"CHECK period against the exact formula: error {100 * error:.2f}% -> {'PASS' if error < 0.01 else 'FAIL'} (limit 1%)")

    fig, ax = plt.subplots(figsize=(6, 3.4))
    ax.plot(np.arange(len(heard)) * dt, heard, lw=0.8)
    ax.set_xlabel("time (s)"); ax.set_ylabel("surface motion"); ax.set_title("Middle of the basin surface")
    fig.tight_layout(); fig.savefig(os.path.join(OUT, "06_surface.png"), dpi=150)

    # How wrong is 4H/c? Its error depends only on the shape, W/H.
    shapes = np.array([1, 2, 4, 6, 10, 20, 50])
    errors = [abs(4 * H / C - exact_period(k * H, H, C)) / exact_period(k * H, H, C) for k in shapes]
    for k, e in zip(shapes, errors):
        print(f"  width/depth {k:3g}: 4H/c is off by {100 * e:5.1f}%")
    fig, ax = plt.subplots(figsize=(5, 3.4))
    ax.semilogx(shapes, 100 * np.array(errors), "o-")
    ax.set_xlabel("width / depth"); ax.set_ylabel("error of 4H/c (%)"); ax.set_title("When is the quick estimate enough?")
    fig.tight_layout(); fig.savefig(os.path.join(OUT, "06_estimate.png"), dpi=150)
    print(f"plots saved in {OUT}")

    # ---- YOUR TURN ----
    # 1. Before running: at what width/depth does 4H/c reach 10% error? Read it off the plot and check with simulate().
    # 2. The formula assumes a rectangle. Make the floor slope (set u = 0 below a slanted line each step) and see
    #    how the period moves. Explain why.
    # 3. Mexico City, 1985: look up published lakebed depth and shear wave speed. Does the model's period come near
    #    the roughly 2 s that damaged mid-height buildings? Say clearly that this explains, and does not predict.
    # 4. Rule of thumb: a building's sway period is about 0.1 s per storey. Which heights match your basin?
    if "--show" in sys.argv:
        plt.show()


if __name__ == "__main__":
    main()