"""
FTCS (Forward-Time Centered-Space) explicit finite difference solver
for the 1D heat equation:

    dT/dt = alpha * d²T/dx²

Boundary conditions:
  - x=0  (Neumann): heat flux step Q(t) applied via ghost cell
  - x=L  (Robin):   convective boundary condition

Returns T_matrix of shape (NX, NT).
"""

import sys
import os
import numpy as np

# Allow importing config from the project root
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import config


def solve():
    """Run the FTCS solver and return (T_matrix, x_vals, t_vals).

    Returns
    -------
    T_matrix : np.ndarray, shape (NX, NT)
        Temperature at each spatial node (row) and time step (column).
        T_matrix[i, n] = T(x_i, t_n).
    x_vals : np.ndarray, shape (NX,)
        Spatial node positions from 0 to L.
    t_vals : np.ndarray, shape (NT,)
        Time values from 0 to T_END.
    """
    # -----------------------------------------------------------------
    # Parameters from config
    # -----------------------------------------------------------------
    alpha = config.ALPHA
    k     = config.K
    L     = config.L
    T0    = config.T0
    Q_val = config.Q_VAL
    t_step = config.T_STEP
    h     = config.H_CONV
    T_amb = config.T_AMB
    T_end = config.T_END
    NX    = config.NX
    NT    = config.NT

    # -----------------------------------------------------------------
    # Grid
    # -----------------------------------------------------------------
    x_vals = np.linspace(0.0, L, NX)
    t_vals = np.linspace(0.0, T_end, NT)
    dx = x_vals[1] - x_vals[0]
    dt = t_vals[1] - t_vals[0]

    # -----------------------------------------------------------------
    # Stability check (CFL for explicit heat equation: r <= 0.5)
    # -----------------------------------------------------------------
    r = alpha * dt / dx**2
    if dt > dx**2 / (2.0 * alpha):
        print(
            f"[FDM WARNING] Stability condition violated: "
            f"dt={dt:.6g} > dx²/(2*alpha)={dx**2/(2.0*alpha):.6g}  (r={r:.4f} > 0.5). "
            "Solution may diverge."
        )

    # -----------------------------------------------------------------
    # Allocate output matrix and set initial condition
    # -----------------------------------------------------------------
    T_matrix = np.zeros((NX, NT), dtype=np.float64)
    T_matrix[:, 0] = T0  # uniform IC

    # Working array for the current time level
    T_cur = np.full(NX, T0, dtype=np.float64)

    # -----------------------------------------------------------------
    # Time integration
    # -----------------------------------------------------------------
    for n in range(NT - 1):
        t_now = t_vals[n]

        # --- Neumann BC at x=0: ghost cell ---
        # Q(t) = Q_val if t >= t_step else 0
        Q = Q_val if t_now >= t_step else 0.0
        T_ghost = T_cur[1] + 2.0 * dx * Q / k  # ghost node T[-1]

        # --- Interior FTCS update (indices 1 .. NX-2) ---
        T_new = T_cur.copy()
        T_new[1:-1] = T_cur[1:-1] + r * (T_cur[2:] - 2.0 * T_cur[1:-1] + T_cur[:-2])

        # --- Apply Neumann BC at i=0 using ghost cell ---
        T_new[0] = T_cur[0] + r * (T_cur[1] - 2.0 * T_cur[0] + T_ghost)

        # --- Apply Robin BC at x=L (explicit) ---
        T_new[-1] = (T_cur[-2] + dx * h / k * T_amb) / (1.0 + dx * h / k)

        T_cur = T_new
        T_matrix[:, n + 1] = T_cur

    return T_matrix, x_vals, t_vals