"""TopEdu project 2: How deep does summer go?

The ground's surface warms and cools once a year. Underground the swing gets smaller and arrives later.
How much smaller, and how much later, tells you the soil's thermal diffusivity (kappa).

Stage 1: run this file. It steps the heat equation through one year in a 20 m column of soil, driven by a yearly
cycle at the surface, and compares the whole field with the exact answer. The CHECK line should say PASS.
Stage 2: fit kappa from temperatures at several depths, first from the model, then from a real station.

Units: depth in metres, time in years, temperature in degrees C away from the yearly average.
Run:  python 02_soil_temperature.py  [--show]   or   python 02_soil_temperature.py my_station.csv
Needs: numpy and matplotlib
"""
import csv
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)

KAPPA = 15.8       # m^2 per year (about 5e-7 m^2/s, a damp soil)
AMPLITUDE = 10.0   # degrees C: the surface swings from -10 to +10 around its average
DEPTH = 20.0       # metres of soil in the model
N = 161            # grid points


def damping_depth(kappa):
    """d = sqrt(2 kappa / omega). At depth d the swing has shrunk to 1/e (37%) of the surface swing."""
    omega = 2 * np.pi            # one cycle per year
    return np.sqrt(2 * kappa / omega)


def exact(z, t, kappa=KAPPA, amplitude=AMPLITUDE):
    d = damping_depth(kappa)
    return amplitude * np.exp(-z / d) * np.sin(2 * np.pi * t - z / d)


def run_year(kappa=KAPPA, amplitude=AMPLITUDE, n=N, samples=200):
    """Steps du/dt = kappa d2u/dz2 for one year. Returns times, depths and the temperature at each sampled time."""
    z = np.linspace(0, DEPTH, n)
    h = z[1] - z[0]
    steps = int(np.ceil(1.0 / (0.4 * h * h / kappa)))   # explicit steps need dt < 0.5 h^2 / kappa
    steps = samples * int(np.ceil(steps / samples))
    dt = 1.0 / steps
    u = exact(z, 0.0, kappa, amplitude)          # start from the settled yearly pattern
    times, frames = [0.0], [u.copy()]
    for step in range(1, steps + 1):
        t = step * dt
        inside = u[1:-1] + kappa * dt * (u[2:] - 2 * u[1:-1] + u[:-2]) / h ** 2
        bottom = u[-1] + kappa * dt * 2 * (u[-2] - u[-1]) / h ** 2   # insulated bottom (mirror point)
        u = np.concatenate(([amplitude * np.sin(2 * np.pi * t)], inside, [bottom]))
        if step % (steps // samples) == 0:
            times.append(t); frames.append(u.copy())
    return np.array(times), z, np.array(frames)


def fit_yearly_cycle(t, temperature):
    """Least-squares fit  T = mean + a sin(2 pi t) + b cos(2 pi t), with t in years.

    Returns (mean, amplitude, lag in years behind a pure sine that starts at t = 0). Works for model output and for real data with gaps.
    """
    t = np.asarray(t, float); temperature = np.asarray(temperature, float)
    keep = np.isfinite(t) & np.isfinite(temperature)
    A = np.column_stack([np.ones(keep.sum()), np.sin(2 * np.pi * t[keep]), np.cos(2 * np.pi * t[keep])])
    mean, a, b = np.linalg.lstsq(A, temperature[keep], rcond=None)[0]
    return mean, np.hypot(a, b), -np.arctan2(b, a) / (2 * np.pi)


def kappa_from_depths(depths, amplitudes, phases):
    """Two estimates of kappa: from how fast the swing shrinks, and from how fast it falls behind.

    ln(amplitude) falls by 1/d per metre; the delay grows by 1/(2 pi d) years per metre. If both give the same
    kappa, the simple model fits. If they disagree, that is a finding: layers or moisture may be the cause.
    """
    slope_amp = np.polyfit(depths, np.log(amplitudes), 1)[0]
    slope_lag = np.polyfit(depths, -np.unwrap(2 * np.pi * np.asarray(phases)) / (2 * np.pi), 1)[0]
    d_amp, d_lag = -1 / slope_amp, 1 / (2 * np.pi * slope_lag)
    return np.pi * d_amp ** 2, np.pi * d_lag ** 2      # kappa = omega d^2 / 2 = pi d^2 (per year)


def read_station(path):
    """Reads a CSV with a column 'days' (days since any start) and one column per depth named like 'cm_10'.

    Real stations come in many layouts. NOAA's U.S. Climate Reference Network files have soil temperatures at
    5, 10, 20, 50 and 100 cm; read their header file, then copy the columns you need into this simple layout.
    Missing readings (-9999 in USCRN files) are skipped.
    """
    with open(path, newline="") as handle:
        rows = list(csv.DictReader(handle))
    days = np.array([float(row["days"]) for row in rows])
    series = {}
    for key in rows[0]:
        if key.startswith("cm_"):
            values = np.array([float(row[key]) if row[key] not in ("", "-9999", "-9999.0") else np.nan for row in rows])
            series[float(key[3:]) / 100] = values
    return days / 365.25, series


def main():
    times, z, frames = run_year()
    worst = np.abs(frames - exact(z[None, :], times[:, None])).max()
    print(f"largest error over the year: {worst:.2e} degrees C ({100 * worst / AMPLITUDE:.3f}% of the swing)")
    print(f"CHECK whole year against the exact answer -> {'PASS' if worst < 0.01 * AMPLITUDE else 'FAIL'} (limit 1% of the swing)")

    # Pretend the model is a weather station: read it at the depths a real station uses.
    depths = np.array([0.05, 0.10, 0.20, 0.50, 1.00])
    amplitudes, phases = [], []
    for depth in depths:
        values = np.array([np.interp(depth, z, frame) for frame in frames])
        _, amp, phase = fit_yearly_cycle(times, values)
        amplitudes.append(amp); phases.append(phase)
        print(f"  depth {depth:4.2f} m: swing {amp:5.2f} C, {365 * (phase % 1):5.1f} days behind the surface")
    k_amp, k_lag = kappa_from_depths(depths, amplitudes, phases)
    print(f"kappa from shrinking: {k_amp:.2f}   kappa from lag: {k_lag:.2f}   (the model used {KAPPA})")
    print(f"damping depth d = {damping_depth(KAPPA):.2f} m: below about 3d = {3 * damping_depth(KAPPA):.1f} m the year is almost invisible")

    fig, ax = plt.subplots(figsize=(5.5, 4))
    for frame_index in range(0, len(times), 25):
        ax.plot(frames[frame_index], z, lw=1)
    ax.set_ylim(6, 0); ax.set_xlabel("temperature away from the average (C)"); ax.set_ylabel("depth (m)")
    ax.set_title("Temperature with depth, at eight times of year")
    fig.tight_layout(); fig.savefig(os.path.join(OUT, "02_profiles.png"), dpi=150)

    station = [arg for arg in sys.argv[1:] if arg.endswith(".csv")]
    if station:
        years, series = read_station(station[0])
        fits = {depth: fit_yearly_cycle(years, values) for depth, values in sorted(series.items())}
        for depth, (mean, amp, phase) in fits.items():
            print(f"  station {depth:.2f} m: mean {mean:.1f} C, swing {amp:.2f} C, phase {phase:.3f} years")
        print("station kappa (shrinking, lag): %.2f, %.2f m^2/year" % kappa_from_depths(
            list(fits), [f[1] for f in fits.values()], [f[2] for f in fits.values()]))
    print(f"plots saved in {OUT}")

    # ---- YOUR TURN ----
    # 1. Predict first: if the soil were twice as diffusive (wet), how much deeper would the 1 C swing reach?
    #    Then run with KAPPA = 31.6 and check.
    # 2. Download two years of soil temperatures from a USCRN station, put them in the simple CSV layout described
    #    in read_station, and run:  python 02_soil_temperature.py my_station.csv
    # 3. Compare two stations with different soils or rainfall. Do "kappa from shrinking" and "kappa from lag"
    #    agree? If not, why might a real soil break this one-layer model?
    if "--show" in sys.argv:
        plt.show()


if __name__ == "__main__":
    main()