# ============================================================
# File: _solve_first_order_nonlinear_pde.py
# Topic: First-Order Nonlinear PDE (Method of Characteristics)
#
# We generate a clean quasi-linear PDE with a closed-form solution:
#
#   u_t + (u + alpha) u_x = 0,     u(x,0) = m x + c
#
# Characteristics:
#   dt/ds = 1,  dx/ds = u + alpha,  du/ds = 0
#
# So u = f(xi) along characteristics, and
#   x = xi + (f(xi) + alpha) t
#
# With f(xi)=m xi + c:
#   x = xi(1 + m t) + t(c+alpha)
#   xi = (x - t(c+alpha))/(1 + m t)
#   u(x,t) = m xi + c = (m x + c - m alpha t)/(1 + m t)
#
# Note: t != -1/m (characteristics intersect there for this IC).
# ============================================================

# ============================================================
# File: _solve_first_order_nonlinear_pde.py
# Topic: First-Order Nonlinear PDE (Method of Characteristics)
# ============================================================

from __future__ import annotations
from typing import Any, Dict, Optional, Union
import sympy as sp

def solve_first_order_nonlinear_pde(
    m: Union[int, str, sp.Expr],
    c: Union[int, str, sp.Expr],
    alpha: Union[int, str, sp.Expr],
    *,
    x_symbol: str = "x",
    t_symbol: str = "t",
    u_symbol: str = "u",
    simplify: bool = True,
    locals_map: Optional[Dict[str, Any]] = None,
) -> dict:
    """
    Solve PDE by characteristics:
        u_t + (u + alpha) u_x = 0,   u(x,0)=m x + c
    Returns detailed HTML/LaTeX solution steps.
    """

    x = sp.Symbol(x_symbol, real=True)
    t = sp.Symbol(t_symbol, real=True)
    # Parameter for characteristic curve
    s = sp.Symbol('s', real=True)
    # Initial position parameter
    xi = sp.Symbol(r'\xi', real=True)

    u = sp.Function(u_symbol)
    uxt = u(x, t)

    base_locals = {x_symbol: x, t_symbol: t, "pi": sp.pi, "E": sp.E, "exp": sp.exp}
    if locals_map:
        base_locals.update(locals_map)

    def L(expr: sp.Expr) -> str:
        return sp.latex(expr)

    def _parse(expr):
        if isinstance(expr, str):
            return sp.sympify(expr, locals=base_locals)
        return sp.sympify(expr)

    # 1. Parse Inputs
    mv = _parse(m)
    cv = _parse(c)
    alphav = _parse(alpha)

    if simplify:
        mv = sp.simplify(mv)
        cv = sp.simplify(cv)
        alphav = sp.simplify(alphav)

    # 2. Prepare Math Objects
    # Speed term coefficient: (u + alpha)
    # We display it carefully for LaTeX
    if alphav == 0:
        speed_disp = "u"
    else:
        # Format nice string for u + alpha
        alpha_s = L(alphav)
        if alphav.is_number and alphav < 0:
             speed_disp = f"u - {L(-alphav)}"
        else:
             speed_disp = f"u + {alpha_s}"

    # Initial Condition function f(xi) = m*xi + c
    f_xi = mv * xi + cv

    # 3. Solve Process (Algebraic Steps for Display)
    # Characteristic equations
    # dt/ds = 1 => t = s
    # du/ds = 0 => u = f(xi)
    # dx/ds = u + alpha => x = s*(f(xi) + alpha) + xi

    # Since t=s, we substitute s->t directly in display logic
    u_char = f_xi # u is constant along characteristic
    slope_char = u_char + alphav

    # x = xi + (slope)*t
    x_eq_rhs = xi + slope_char * t

    # Solve for xi in terms of x, t
    # x = xi + (m*xi + c + alpha)*t
    # x = xi(1 + m*t) + t(c + alpha)
    # xi = (x - t(c+alpha)) / (1 + m*t)

    denom = 1 + mv * t
    num = x - t * (cv + alphav)
    xi_solved = num / denom

    # Final u(x,t) = m*xi_solved + c
    u_final_raw = mv * xi_solved + cv
    u_final = sp.simplify(u_final_raw) if simplify else u_final_raw

    # 4. Construct Output Strings

    # Problem Statement
    problem_s = (
        f"Solve the first-order nonlinear PDE using the <b>Method of Characteristics</b>.<br>"
        f"<b>PDE:</b> \\( u_t + ({speed_disp})u_x = 0 \\)<br>"
        f"<b>Initial Condition:</b> \\( u(x, 0) = {L(mv*x + cv)} \\)"
    )

    # Solution Steps
    soln_s = (
        "<div class='wrap'><div class='card'>"
        "<div class='badge'>Method of Characteristics</div>"

        "<h3>1. Set up Characteristic Equations</h3>"
        "<p>We convert the PDE into a system of Ordinary Differential Equations (ODEs) parametrized by \\(s\\). <br>"
        "Matching coefficients with the chain rule form \\(\\frac{du}{ds} = \\frac{\\partial u}{\\partial t}\\frac{dt}{ds} + \\frac{\\partial u}{\\partial x}\\frac{dx}{ds}\\):</p>"
        "$$ \\begin{aligned} "
        "\\frac{dt}{ds} &= 1 & \\implies t &= s \\quad (\\text{assuming } t_0=0) \\\\[0.5em] "
        f"\\frac{{dx}}{{ds}} &= {speed_disp} & & \\\\[0.5em] "
        "\\frac{du}{ds} &= 0 & \\implies u &= \\text{constant along characteristics} "
        "\\end{aligned} $$"

        "<h3>2. Integrate along Characteristics</h3>"
        "<p>Since \\(u\\) is constant along a characteristic curve starting at \\(x = \\xi\\) when \\(t=0\\):</p>"
        f"$$ u(x,t) = u(\\xi, 0) = f(\\xi) = {L(f_xi)} $$"
        "<p>Now substitute \\(u = f(\\xi)\\) into the \\(dx/ds\\) equation. Since \\(t=s\\):</p>"
        f"$$ \\frac{{dx}}{{dt}} = {speed_disp} = f(\\xi) + {L(alphav)} $$"
        "<p>Integrating with respect to \\(t\\) (noting that \\(\\xi\\) is fixed for a specific characteristic):</p>"
        f"$$ x = \\xi + (f(\\xi) + {L(alphav)})t $$"

        "<h3>3. Solve for \\(\\xi\\)</h3>"
        "<p>We need to map the current point \\((x,t)\\) back to its starting point \\(\\xi\\). Substitute \\(f(\\xi) = {L(f_xi)}\\) into the equation for \\(x\\):</p>"
        f"$$ x = \\xi + ({L(f_xi)} + {L(alphav)})t $$"
        "$$ x = \\xi + m\\xi t + ct + \\alpha t $$"
        f"$$ x = \\xi(1 + {L(mv)}t) + t({L(cv + alphav)}) $$"
        "<p>Rearranging to isolate \\(\\xi\\):</p>"
        f"$$ \\xi = \\frac{{x - t({L(cv + alphav)})}}{{1 + {L(mv)}t}} $$"

        "<h3>4. Final Substitution</h3>"
        "<p>Finally, substitute \\(\\xi\\) back into the solution \\(u(x,t) = {L(f_xi)}\\):</p>"
        f"$$ u(x,t) = {L(mv)} \\left( \\frac{{x - t({L(cv + alphav)})}}{{1 + {L(mv)}t}} \\right) + {L(cv)} $$"
        f"$$ u(x,t) = {L(u_final)} $$"
        "</div></div>"
        f"{explain_method_of_characteristics()}"
    )

    answer_s = f"The solution is \\( u(x,t) = {L(u_final)} \\)."

    return {
        "problem_s": problem_s,
        "soln_s": soln_s,
        "answer_s": answer_s
    }

def explain_method_of_characteristics() -> str:
    """
    Returns an HTML/LaTeX explanation of the intuition behind
    Method of Characteristics.
    """

    explanation = (
        "<div class='wrap'><div class='card'>"
        "<div class='badge'>Concept & Intuition</div>"
        "<h3>What is the Method of Characteristics?</h3>"

        "<p>The Method of Characteristics solves First-Order PDEs by reducing them to ODEs along specific curves called <b>characteristic curves</b>.</p>"

        "<h4>1. Geometric Interpretation</h4>"
        "<p>Imagine the solution \\(u(x,t)\\) as a surface. The PDE tells us that at every point, the directional derivative of \\(u\\) in a specific direction is zero.</p>"
        "<ul>"
        "<li>This means \\(u\\) is <b>constant</b> along these specific curves.</li>"
        "<li>If you know the value of \\(u\\) at the start of the curve (the initial condition), that value simply 'transports' along the curve.</li>"
        "</ul>"

        "<h4>2. The Traffic Analogy</h4>"
        "<p>Consider traffic flow where density is \\(u\\). The equation \\(u_t + u u_x = 0\\) says that the density propagates at speed \\(u\\).</p>"
        "<ul>"
        "<li>Regions with higher density (higher \\(u\\)) move faster.</li>"
        "<li>Regions with lower density move slower.</li>"
        "<li>If faster traffic catches up to slower traffic, the characteristic curves intersect, forming a <b>shock wave</b> (mathematical discontinuity).</li>"
        "</ul>"

        "</div></div>"
    )
    return explanation