"""TopEdu project 4: How thick can lab-grown tissue be?

Oxygen soaks into tissue from its surface while the cells use it up. With no blood vessels, the centre of a
large piece starves. In scaled units the oxygen c obeys  -laplacian(c) + c = 0  inside, with c = 1 on the surface.
Lengths are measured in "penetration lengths" L = sqrt(D / k): D is how fast oxygen spreads, k how fast it is used.

Stage 1: run this file. It solves a flat slab (exact answer: cosh) and a cube (exact answer: a series), and checks
both. The CHECK line should say PASS.
Stage 2: find the largest cube, slab and bar whose centre stays above an oxygen threshold you choose.

Run:  python 04_tissue_oxygen.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 slab(thickness, n=81):
    """A flat sheet, open on both faces: solve -c'' + c = 0 with c = 1 at both faces, as one small linear system."""
    x = np.linspace(0, thickness, n)
    h = x[1] - x[0]
    inner = n - 2
    A = np.diag(np.full(inner, 2 / h ** 2 + 1)) + np.diag(np.full(inner - 1, -1 / h ** 2), 1) + np.diag(np.full(inner - 1, -1 / h ** 2), -1)
    rhs = np.zeros(inner); rhs[0] = rhs[-1] = 1 / h ** 2      # the surface value c = 1 enters here
    return x, np.concatenate(([1.0], np.linalg.solve(A, rhs), [1.0]))


def box(sides, n=17, tolerance=1e-9):
    """A box with sides (a, b, c), open on every face. Solved by Jacobi sweeps: each point becomes the value
    the equation wants, given its six neighbours, until nothing changes. Returns the oxygen at every point."""
    h = [s / (n - 1) for s in sides]
    w = [1 / hh ** 2 for hh in h]
    c = np.ones((n, n, n))
    for sweep in range(100_000):
        new = c.copy()
        new[1:-1, 1:-1, 1:-1] = (w[0] * (c[2:, 1:-1, 1:-1] + c[:-2, 1:-1, 1:-1])
                                 + w[1] * (c[1:-1, 2:, 1:-1] + c[1:-1, :-2, 1:-1])
                                 + w[2] * (c[1:-1, 1:-1, 2:] + c[1:-1, 1:-1, :-2])) / (2 * sum(w) + 1)
        if np.abs(new - c).max() < tolerance:
            return new
        c = new
    raise RuntimeError("did not settle")


def box_centre_series(sides, terms=60):
    """Exact oxygen at the centre of a box: 1 minus a triple sine series (only odd terms survive)."""
    a, b, c = sides
    total = 0.0
    for l in range(1, terms, 2):
        for m in range(1, terms, 2):
            for k in range(1, terms, 2):
                sign = (-1) ** ((l - 1) // 2 + (m - 1) // 2 + (k - 1) // 2)
                total += sign * 64 / (np.pi ** 3 * l * m * k) / (1 + np.pi ** 2 * (l ** 2 / a ** 2 + m ** 2 / b ** 2 + k ** 2 / c ** 2))
    return 1 - total


def main():
    x, c = slab(4.0)
    exact = np.cosh(x - 2.0) / np.cosh(2.0)
    slab_error = np.abs(c - exact).max()
    print(f"slab 4 thick: centre {c[len(c) // 2]:.4f}, exact {exact[len(c) // 2]:.4f}, worst error {slab_error:.1e}")

    worst = slab_error
    for side in (1.0, 2.0, 4.0):
        field = box((side, side, side))
        centre = field[8, 8, 8]
        series = box_centre_series((side, side, side))
        error = abs(centre - series) / series
        worst = max(worst, error)
        print(f"cube side {side:g}: centre oxygen {centre:.4f}, series {series:.4f}, error {100 * error:.2f}%")
    print(f"CHECK slab and cubes against exact answers -> {'PASS' if worst < 0.01 else 'FAIL'} (limit 1%)")

    fig, ax = plt.subplots(figsize=(5.5, 3.8))
    ax.plot(x, c, label="computed"); ax.plot(x, exact, "--", label="exact cosh")
    ax.set_xlabel("depth across the slab (penetration lengths)"); ax.set_ylabel("oxygen (surface = 1)"); ax.legend()
    ax.set_title("Oxygen across a slab 4 penetration lengths thick")
    fig.tight_layout(); fig.savefig(os.path.join(OUT, "04_slab.png"), dpi=150)
    print(f"plots saved in {OUT}")

    # ---- YOUR TURN ----
    # 1. Pick a threshold, say the centre must keep 20% of the surface oxygen. Predict the largest cube side,
    #    then find it by trying sides between 2 and 8 with box_centre_series (fast) and checking one with box().
    # 2. Same volume, other shapes: a slab 0.5 thick and a long square bar (make one side 10). Which keeps its
    #    centre alive best? Is it the one with the most surface?
    # 3. Turn units back into micrometres: look up D for oxygen in tissue and a cell consumption rate, compute
    #    L = sqrt(D / k), and compare your limit with published organoid sizes (often a few hundred micrometres).
    # 4. Honest limit: real cells use oxygen at a nearly fixed rate until it runs low. Would that make your
    #    limit larger or smaller? Argue it before you test it.
    if "--show" in sys.argv:
        plt.show()


if __name__ == "__main__":
    main()