"""TopEdu project 8: Why a ruler twists so easily

When you twist a bar, its stiffness comes from a "stress function" phi over the cross-section, shaped like a soap
film blown up over a frame the shape of the section:  -laplacian(phi) = 2 inside, phi = 0 on the edge.
The torsion constant is J = 2 x (volume under the film). For an a by b rectangle (b the short side) write
J = k a b^3. A square has k = 0.1406; a long thin strip approaches k = 1/3.
Twist per length = torque / (G J), with G the shear modulus.

Stage 1: run this file. It solves for phi on three rectangles and checks k against the exact series.
The CHECK line should say PASS.
Stage 2: twist real bars, measure their stiffness, and test the rule.

Run:  python 08_beam_torsion.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)


def soap_film(a, b, nx=65, ny=33, tolerance=1e-10):
    """Solves -laplacian(phi) = 2 on an a by b rectangle with phi = 0 on the edge.

    Red-black over-relaxation (SOR): update every other point like a chessboard, then the rest, each time pushing
    a little past the value the equation asks for (omega). Plain Jacobi works too, but needs thousands more sweeps.
    """
    x = np.linspace(0, a, nx); y = np.linspace(0, b, ny)
    hx, hy = x[1] - x[0], y[1] - y[0]
    wx, wy = 1 / hx ** 2, 1 / hy ** 2
    phi = np.zeros((nx, ny))
    I, Jj = np.meshgrid(np.arange(nx), np.arange(ny), indexing="ij")
    colours = [((I + Jj) % 2 == c)[1:-1, 1:-1] for c in (0, 1)]
    omega = 2 / (1 + np.sin(np.pi / max(nx, ny)))       # a good over-relaxation factor
    for sweep in range(20_000):
        biggest = 0.0
        for colour in colours:
            target = (wx * (phi[2:, 1:-1] + phi[:-2, 1:-1]) + wy * (phi[1:-1, 2:] + phi[1:-1, :-2]) + 2) / (2 * wx + 2 * wy)
            inner = phi[1:-1, 1:-1]
            step = omega * (target - inner) * colour
            inner += step
            biggest = max(biggest, np.abs(step).max())
        if biggest < tolerance:
            return x, y, phi
    raise RuntimeError("did not settle")


def torsion_factor(a, b, **grid):
    """k = J / (a b^3), with J = 2 x the volume under phi (trapezoid rule)."""
    x, y, phi = soap_film(a, b, **grid)
    volume = np.trapezoid(np.trapezoid(phi, y, axis=1), x) if hasattr(np, "trapezoid") else np.trapz(np.trapz(phi, y, axis=1), x)
    return 2 * volume / (a * b ** 3), (x, y, phi)


def series_factor(a, b, terms=50):
    """The exact k for a rectangle (a the long side, b the short one)."""
    n = np.arange(1, 2 * terms, 2)
    return (1 - 192 / np.pi ** 5 * (b / a) * np.sum(np.tanh(n * np.pi * a / (2 * b)) / n ** 5)) / 3


def main():
    worst = 0.0
    for a, b in [(1, 1), (2, 1), (4, 1)]:
        k, field = torsion_factor(a, b)
        exact = series_factor(a, b)
        error = abs(k - exact) / exact
        worst = max(worst, error)
        print(f"{a}:{b} bar: k = {k:.4f}, series {exact:.4f}, error {100 * error:.2f}%")
        if (a, b) == (2, 1):
            saved = field
    print(f"CHECK torsion factor against the series -> {'PASS' if worst < 0.01 else 'FAIL'} (limit 1%)")

    x, y, phi = saved
    fig, ax = plt.subplots(figsize=(5.5, 3.2))
    image = ax.contourf(x, y, phi.T, 20, cmap="viridis"); fig.colorbar(image, label="phi (film height)")
    ax.set_aspect("equal"); ax.set_title("The soap film over a 2 by 1 section")
    fig.tight_layout(); fig.savefig(os.path.join(OUT, "08_film.png"), dpi=150)

    ratios = np.array([1, 1.5, 2, 3, 5, 10, 20])
    fig, ax = plt.subplots(figsize=(5, 3.4))
    ax.semilogx(ratios, [series_factor(r, 1) for r in ratios], "o-", label="exact k")
    ax.semilogx(ratios, (1 - 0.63 / ratios) / 3, "--", label="thin-strip rule (1 - 0.63 b/a)/3")
    ax.axhline(1 / 3, color="gray", lw=0.6); ax.set_xlabel("a / b"); ax.set_ylabel("k = J / (a b^3)"); ax.legend()
    fig.tight_layout(); fig.savefig(os.path.join(OUT, "08_rule.png"), dpi=150)
    print(f"plots saved in {OUT}")

    # ---- YOUR TURN ----
    # 1. A ruler is about 30 by 3 mm; a pencil about 7 by 7 mm. Predict which is stiffer in twist, and by how much,
    #    using J = k a b^3. Is the answer what your hands tell you?
    # 2. Where does the thin-strip rule (the dashed line) go wrong by more than 5%? Read it from the plot, then check.
    # 3. Experiment: clamp one end of a foam or wood bar, hang a weight from an arm at the other end, and measure the
    #    twist angle. Stiffness = torque x length / angle should be G J. Do bars of different shapes agree on G?
    # 4. Honest limit: only rectangles, and only small twists. Try a hollow square by setting phi to a constant on
    #    an inner hole, and say what changes.
    if "--show" in sys.argv:
        plt.show()


if __name__ == "__main__":
    main()