#!/usr/bin/env python3
"""
PSC6125 template:
One-dimensional advection equation.

Equation:

    dT/dt = -u dT/dx

Students can modify:
    1. Initial condition
    2. Advection tendency
    3. Numerical scheme

Output:
    mp4 animation
    png final figure
"""

import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation


# ============================================================
# Parameters
# ============================================================
NT = 250          # number of time steps
DT = 1.0          # time step

XMIN = -50.0
XMAX = 50.0
DX = 1.0

U = 0.5           # constant advection velocity

# Initial bell-shape parameters
X0 = 0.0          # initial center location
K = 15            # half-width in grid points


# ============================================================
# Grid
# ============================================================
nx = int((XMAX - XMIN) / DX) + 1
x = np.linspace(XMIN, XMAX, nx)
time = np.arange(NT) * DT


# ============================================================
# Initial condition
# ============================================================
def initial_condition(x):
    """
    Create a bell-shaped initial tracer distribution.
    """

    sigma = float(K) * DX / 3.5

    T0 = np.exp(-((x - X0)**2) / (2.0 * sigma**2))

    # Make the field exactly zero far from the initial center
    T0[np.abs(x - X0) > K * DX] = 0.0

    return T0


# ============================================================
# Advection tendency
# Students modify this section later
# ============================================================
def advection_tendency(T, x, U):
    """
    Compute dT/dt for the advection equation.

    Template version:
    No advection is applied yet.
    """

    dTdt = np.zeros_like(T)

    return dTdt


# ============================================================
# Time integration
# ============================================================
T_save = np.zeros((NT, nx))

T = initial_condition(x)
T_initial = T.copy()

T_save[0, :] = T

for n in range(NT - 1):
    dTdt = advection_tendency(T, x, U)

    T = T + DT * dTdt

    T_save[n + 1, :] = T


# ============================================================
# Animation setup
# ============================================================
fig, ax = plt.subplots(figsize=(9, 4))

line0, = ax.plot(x, T_initial, color="k", linewidth=1.5, label="Initial")
lineT, = ax.plot([], [], "o-", color="C0", markersize=4, linewidth=2.5, label="Numerical")
time_text = ax.text(
    0.78, 0.85, "",
    transform=ax.transAxes,
    bbox=dict(facecolor="white", edgecolor="none", alpha=0.8)
)

ax.set_xlim(XMIN, XMAX)
ax.set_ylim(-0.5, 1.5)
ax.set_xlabel("x")
ax.set_ylabel("T")
ax.set_title("One-Dimensional Advection")
ax.grid(True)
ax.legend()


def init():
    lineT.set_data([], [])
    time_text.set_text("")
    return lineT, time_text


def update(frame):
    lineT.set_data(x, T_save[frame, :])
    time_text.set_text(f"t = {time[frame]:.1f}")
    return lineT, time_text


ani = animation.FuncAnimation(
    fig,
    update,
    frames=NT,
    init_func=init,
    interval=50,
    blit=True
)


# ============================================================
# Save animation and final figure
# ============================================================
plt.tight_layout()

script_name = os.path.splitext(os.path.basename(__file__))[0]

outfile_mp4 = script_name + ".mp4"
writer = animation.FFMpegWriter(fps=20, bitrate=1800)
ani.save(outfile_mp4, writer=writer)
print("Animation saved as:", outfile_mp4)

outfile_png = script_name + ".png"
plt.savefig(outfile_png, dpi=150)
print("Final figure saved as:", outfile_png)

plt.show()