"""TopEdu project 1: Where should the hot chip go?

A chip heats a square metal plate. The bottom edge (y = 0) is held cold, the other three edges are insulated.
How hot does the plate get, and how much does the chip's position change that?

Stage 1 (weeks 1 to 3): run this file. It solves the heat equation step by step until the plate settles, then
checks the peak temperature against an exact series answer. You should see a CHECK line that says PASS.
Stage 2 (weeks 4 to 9): change the chip, then answer your own question in the YOUR TURN section at the bottom.

Units are scaled: the plate is 1 by 1, and temperature is "degrees above the cold edge".
Run:  python 01_chip_cooling.py          (add --show to open the plots in a window)
Needs: Python 3 with numpy and matplotlib  (pip install numpy 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)

# ---- The setup (the same defaults as the online lab) ----
X0, Y0 = 0.5, 0.5   # chip centre; Y0 is the distance from the cold edge
SIZE = 0.12         # chip size (the width of its heat spot)
POWER = 100.0       # peak heating rate
N = 33              # grid points along each side


def heat_source(x, y, x0=X0, y0=Y0, size=SIZE):
    """Heat put in per unit area: a smooth spot centred on the chip."""
    return POWER * np.exp(-((x - x0) ** 2 + (y - y0) ** 2) / (2 * size ** 2))


def solve_plate(f, n=N, tolerance=1e-7):
    """Steps the heat equation  du/dt = laplacian(u) + f  until the temperature stops changing.

    u[i, j] is the temperature at x = i*h, y = j*h. Insulated edges use a mirror ("ghost") point, so no heat
    crosses them; the cold edge y = 0 is held at 0. Returns the settled temperature and the number of steps.
    """
    h = 1.0 / (n - 1)
    dt = 0.24 * h * h          # explicit steps blow up above 0.25 h^2: try 0.26 and watch
    u = np.zeros((n, n))
    for step in range(1, 2_000_000):
        p = np.pad(u, 1, mode="reflect")   # mirror points make the edges insulated
        laplacian = (p[2:, 1:-1] + p[:-2, 1:-1] + p[1:-1, 2:] + p[1:-1, :-2] - 4 * u) / h ** 2
        change = dt * (laplacian + f)
        change[:, 0] = 0.0                # the cold edge stays at 0
        u = u + change
        if np.abs(change).max() < tolerance * dt:
            return u, step
    raise RuntimeError("did not settle")


def series_answer(f_of, x, y, terms=40, quad=400):
    """The exact settled temperature, written as a sum of cosine-times-sine shapes (a Fourier series).

    Each shape cos(m pi x) sin((k + 1/2) pi y) already obeys the edge rules, so the answer is a sum of them, each
    weighted by how much of the heat source looks like that shape divided by how fast that shape cools.
    """
    s = (np.arange(quad) + 0.5) / quad            # midpoint rule on a fine grid
    X, Y = np.meshgrid(s, s, indexing="ij")
    F = f_of(X, Y)
    total = np.zeros_like(x, dtype=float)
    for m in range(terms):
        cx = np.cos(m * np.pi * X)
        for k in range(terms):
            sy = np.sin((k + 0.5) * np.pi * Y)
            weight = (F * cx * sy).mean() / ((cx * sy) ** 2).mean()
            rate = (m * np.pi) ** 2 + ((k + 0.5) * np.pi) ** 2
            total += weight / rate * np.cos(m * np.pi * x) * np.sin((k + 0.5) * np.pi * y)
    return total


def main():
    grid = np.linspace(0, 1, N)
    x, y = np.meshgrid(grid, grid, indexing="ij")
    u, steps = solve_plate(heat_source(x, y))
    i, j = np.unravel_index(np.argmax(u), u.shape)
    exact = series_answer(heat_source, np.array(grid[i]), np.array(grid[j]))
    error = abs(u[i, j] - exact) / exact
    print(f"settled after {steps} steps; peak {u[i, j]:.4f} at x = {grid[i]:.2f}, y = {grid[j]:.2f}")
    print(f"series answer at that point: {float(exact):.4f}")
    print(f"CHECK peak temperature against the series: error {100 * error:.2f}% -> {'PASS' if error < 0.01 else 'FAIL'} (limit 1%)")

    fig, ax = plt.subplots(figsize=(5, 4.2))
    image = ax.imshow(u.T, origin="lower", extent=[0, 1, 0, 1], cmap="inferno")
    fig.colorbar(image, label="temperature above the cold edge")
    ax.set_xlabel("x"); ax.set_ylabel("y (distance from the cold edge)")
    ax.set_title("Settled plate temperature")
    fig.tight_layout(); fig.savefig(os.path.join(OUT, "01_plate.png"), dpi=150)

    # A first sweep: move the chip away from the cold edge and watch the peak.
    distances = np.linspace(0.15, 0.85, 8)
    peaks = [solve_plate(heat_source(x, y, y0=d), tolerance=1e-5)[0].max() for d in distances]
    for d, p in zip(distances, peaks):
        print(f"  chip at distance {d:.2f} from the cold edge: peak {p:.3f}")
    fig, ax = plt.subplots(figsize=(5, 3.6))
    ax.plot(distances, peaks, "o-")
    ax.set_xlabel("chip distance from the cold edge"); ax.set_ylabel("peak temperature")
    ax.set_title("Peak temperature against chip position")
    fig.tight_layout(); fig.savefig(os.path.join(OUT, "01_sweep.png"), dpi=150)
    print(f"plots saved in {OUT}")

    # ---- YOUR TURN ----
    # 1. Write your prediction first: if the chip size doubles with the same total power, does the peak halve?
    #    (Hint: to keep the total power fixed, scale POWER by (old size / new size)^2.) Then run it.
    # 2. Add a second chip: heat_source(x, y, 0.3, 0.4) + heat_source(x, y, 0.7, 0.6). Search for the two
    #    positions with the lowest peak. Is the best layout symmetric? Is it what you guessed?
    # 3. Grid check: run with N = 17, 33 and 65. How much does the peak change? Report that in your paper.
    # 4. Real world: heat an aluminium plate with a power resistor and cool one edge in ice water. Which of your
    #    predictions survive? The gap is the interesting part.
    if "--show" in sys.argv:
        plt.show()


if __name__ == "__main__":
    main()