#!/usr/bin/env python3
"""
Output:
    mp4 animation
    png final figure
"""

import os
import numpy as np
#import matplotlib
#matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.animation as animation


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

XMIN = -50.0
XMAX = 50.0
DX = 1.0

U = 0.0           # constant advection velocity
Kdiff = 0.1

# 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


def diffusion_tendency(T, x, Kdiff):

    Fdiff = np.zeros_like(T)
    NX = len(T)

    for i in range(1, NX-1):
        Fdiff[i] = Kdiff * ( T[i+1] - 2.0 * T[i] + T[i-1]) / DX**2

    # Cyclic boundary points
    Fdiff[0] = Kdiff * ( T[1] - 2.0 * T[0] + T[-1]) / DX**2
    Fdiff[-1] = Kdiff * ( T[0] - 2.0 * T[-1] + T[-2]) / DX**2


    return Fdiff

# ============================================================
# Advection tendency
# Students modify this section later
# ============================================================
def advection_tendency(T, x, U):

    #scheme = "upwind"
    #scheme = "central"
    scheme = "spectral"

    dTdt = np.zeros_like(T)
    NX = len(T)

    # ==================================================
    # Upwind Scheme
    # ==================================================
    if scheme == "upwind":

        dTdt[1:] = -U * np.diff(T) / DX

        # Cyclic boundary
        dTdt[0] = -U * (T[0] - T[-1]) / DX

    # ==================================================
    # Central-Difference Scheme
    # ==================================================
    elif scheme == "central":

        # Interior points
        dTdt[1:NX-1] = -U * (T[2:NX] - T[0:NX-2]) / (2.0 * DX)

        # Cyclic boundary points
        dTdt[0] = -U * (T[1] - T[-1]) / (2.0 * DX)
        dTdt[-1] = -U * (T[0] - T[-2]) / (2.0 * DX)

    # ==================================================
    # Spectral Scheme
    # ==================================================
    elif scheme == "spectral":

        L = x[-1] - x[0]

        T_hat = np.fft.fft(T)

        if NX % 2 == 0:
            k = np.asarray(
                list(range(0, NX // 2))
                + [0]
                + list(range(-NX // 2 + 1, 0)),
                dtype=np.float64
            )
        else:
            k = np.asarray(
                list(range(0, (NX - 1) // 2))
                + [0]
                + list(range(-(NX - 1) // 2, 0)),
                dtype=np.float64
            )

        k = k * 2.0 * np.pi / L

        dTdx = np.fft.ifft(
            1.0j * k * T_hat
        ).real

        dTdt = -U * dTdx

    else:
        raise ValueError(
            f"Unknown scheme '{scheme}'. "
            "Choose 'upwind', 'central', or 'spectral'."
        )

    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):
    # -- Euler --
    Fadv = advection_tendency(T, x, U)

    # -- RK4 --
    #k1 = advection_tendency(T, x, U)
    #k2 = advection_tendency( T + 0.5 * DT * k1, x, U)
    #k3 = advection_tendency( T + 0.5 * DT * k2, x, U)
    #k4 = advection_tendency( T + DT * k3, x, U)
    #Fadv =  ( k1 + 2.0 * k2 + 2.0 * k3 + k4) / 6.0

    Fdiff = diffusion_tendency(T, x, Kdiff)

    T = T + DT * (Fadv + Fdiff)

    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()

#outfile_gif = script_name + ".gif"
#os.system( f'ffmpeg -i {outfile_mp4} -vf "fps=10,scale=900:-1" {outfile_gif}')