# ============================================================
# File: _solve_ode_numerical_rk4_ivp.py
# Topic: Numerical ODE Solver — Runge-Kutta 4 (RK4) for IVPs
#
# IVP:
#   y' = f(x,y),   y(x0)=y0
#
# RK4 step:
#   k1 = f(x_n, y_n)
#   k2 = f(x_n + h/2, y_n + h*k1/2)
#   k3 = f(x_n + h/2, y_n + h*k2/2)
#   k4 = f(x_n + h,   y_n + h*k3)
#   y_{n+1} = y_n + (h/6)(k1 + 2k2 + 2k3 + k4)
#
# Returns:
#   {problem_s, soln_s, answer_s} in the same HTML/LaTeX style.
# ============================================================

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

def solve_ode_numerical_rk4_ivp(
    f_xy: Union[str, sp.Expr],
    *,
    x0: Union[int, float, str, sp.Expr] = 0,
    y0: Union[int, float, str, sp.Expr] = 1,
    h: Union[int, float, str, sp.Expr] = 0.1,
    n_steps: int = 5,
    x_symbol: str = "x",
    y_symbol: str = "y",
    decimals: int = 6,
    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    """
    RK4 approximate solution with detailed step-by-step explanation.
    """

    x = sp.Symbol(x_symbol, real=True)
    y = sp.Symbol(y_symbol, real=True)

    base_locals = {x_symbol: x, y_symbol: y, "pi": sp.pi, "E": sp.E, "exp": sp.exp, "sqrt": sp.sqrt, "sin": sp.sin, "cos": sp.cos, "tan": sp.tan}
    if locals_map:
        base_locals.update(locals_map)

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

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

    def _fmt_num(v: sp.Expr) -> str:
        """Format number for display (rounded float)."""
        vv = sp.N(v)
        try:
            return f"{float(vv):.{decimals}f}"
        except Exception:
            return str(vv)

    # 1. Parse Inputs
    f_raw = _parse(f_xy)
    f = sp.simplify(f_raw) if simplify else f_raw
    x0v = _parse(x0)
    y0v = _parse(y0)
    hv = _parse(h)

    # 2. Setup Loop
    f_sub = lambda xx, yy: f.subs({x: xx, y: yy})

    rows = []
    xn = x0v
    yn = y0v

    # Text buffer for the detailed Step 1 explanation
    step_one_explanation = ""

    # 3. Iteration
    for i in range(1, n_steps + 1):
        # Current state
        x_curr = xn
        y_curr = yn

        # RK4 Constants Calculation
        # k1 = f(x, y)
        k1 = f_sub(x_curr, y_curr)

        # k2 = f(x + h/2, y + h*k1/2)
        x_k2 = x_curr + hv/2
        y_k2 = y_curr + (hv * k1)/2
        k2 = f_sub(x_k2, y_k2)

        # k3 = f(x + h/2, y + h*k2/2)
        x_k3 = x_curr + hv/2
        y_k3 = y_curr + (hv * k2)/2
        k3 = f_sub(x_k3, y_k3)

        # k4 = f(x + h, y + h*k3)
        x_k4 = x_curr + hv
        y_k4 = y_curr + (hv * k3)
        k4 = f_sub(x_k4, y_k4)

        # Compute next y
        delta = (hv/6) * (k1 + 2*k2 + 2*k3 + k4)
        yn_next = y_curr + delta
        xn_next = x_curr + hv

        # Simplify results if requested (internal precision)
        if simplify:
            k1, k2, k3, k4 = sp.simplify(k1), sp.simplify(k2), sp.simplify(k3), sp.simplify(k4)
            yn_next = sp.simplify(yn_next)

        # Store row data for table
        rows.append({
            "i": i,
            "xn": x_curr,
            "yn": y_curr,
            "k1": k1, "k2": k2, "k3": k3, "k4": k4,
            "y_next": yn_next
        })

        # --- Generate Detailed Explanation for Step 1 Only ---
        if i == 1:
            step_one_explanation = (
                f"<h4>Detailed Calculation for Step 1</h4>"
                f"<p>We start at \\(x_0 = {_fmt_num(x0v)}\\) and \\(y_0 = {_fmt_num(y0v)}\\). The step size is \\(h = {_fmt_num(hv)}\\).</p>"
                f"<p><b>1. Calculate slopes (k values):</b></p>"
                f"$$ \\begin{{aligned}} "
                f"k_1 &= f(x_0, y_0) = f({_fmt_num(x_curr)}, {_fmt_num(y_curr)}) = {_fmt_num(k1)} \\\\[0.5em] "
                f"k_2 &= f(x_0 + \\tfrac{{h}}{{2}}, y_0 + \\tfrac{{h k_1}}{{2}}) \\\\"
                f"    &= f({_fmt_num(x_k2)}, {_fmt_num(y_k2)}) = {_fmt_num(k2)} \\\\[0.5em] "
                f"k_3 &= f(x_0 + \\tfrac{{h}}{{2}}, y_0 + \\tfrac{{h k_2}}{{2}}) \\\\"
                f"    &= f({_fmt_num(x_k3)}, {_fmt_num(y_k3)}) = {_fmt_num(k3)} \\\\[0.5em] "
                f"k_4 &= f(x_0 + h, y_0 + h k_3) \\\\"
                f"    &= f({_fmt_num(x_k4)}, {_fmt_num(y_k4)}) = {_fmt_num(k4)} "
                f"\\end{{aligned}} $$"
                f"<p><b>2. Calculate weighted average slope:</b></p>"
                f"$$ \\text{{slope}}_{{avg}} = \\frac{{1}}{{6}}(k_1 + 2k_2 + 2k_3 + k_4) = \\frac{{1}}{{6}}({_fmt_num(k1)} + {_fmt_num(2*k2)} + {_fmt_num(2*k3)} + {_fmt_num(k4)}) $$"
                f"<p><b>3. Update y:</b></p>"
                f"$$ y_1 = y_0 + h \\cdot (\\text{{slope}}_{{avg}}) = {_fmt_num(y_curr)} + {_fmt_num(delta)} = {_fmt_num(yn_next)} $$"
                f"<p>So, at \\(x_1 = {_fmt_num(xn_next)}\\), the approximate value is \\(y_1 \\approx {_fmt_num(yn_next)}\\).</p>"
                "<hr/>"
            )

        # Update for next iteration
        xn = xn_next
        yn = yn_next

    x_end = xn
    y_end = yn

    # 4. Build HTML Summary Table
    table_header = (
        "<table class='table table-sm table-striped'>"
        "<thead><tr>"
        "<th>Step (n)</th><th>\\(x_n\\)</th><th>\\(y_n\\)</th>"
        "<th>\\(k_1\\)</th><th>\\(k_2\\)</th><th>\\(k_3\\)</th><th>\\(k_4\\)</th>"
        "<th>\\(y_{n+1}\\)</th>"
        "</tr></thead><tbody>"
    )

    table_rows = ""
    for r in rows:
        table_rows += (
            "<tr>"
            f"<td>{r['i']}</td>"
            f"<td>{_fmt_num(r['xn'])}</td>"
            f"<td>{_fmt_num(r['yn'])}</td>"
            f"<td>{_fmt_num(r['k1'])}</td>"
            f"<td>{_fmt_num(r['k2'])}</td>"
            f"<td>{_fmt_num(r['k3'])}</td>"
            f"<td>{_fmt_num(r['k4'])}</td>"
            f"<td><b>{_fmt_num(r['y_next'])}</b></td>"
            "</tr>"
        )

    table_html = table_header + table_rows + "</tbody></table>"

    # 5. Output Construction
    problem_s = (
        f"Solve the Initial Value Problem (IVP) numerically using the <b>Runge-Kutta 4th Order (RK4)</b> method.<br>"
        f"<b>ODE:</b> \\( \\frac{{dy}}{{dx}} = {L(f)} \\)<br>"
        f"<b>Initial Condition:</b> \\( y({_fmt_num(x0v)}) = {_fmt_num(y0v)} \\)<br>"
        f"<b>Parameters:</b> Step size \\( h = {_fmt_num(hv)} \\) for \\( {n_steps} \\) steps."
    )

    soln_s = (
        "<div class='wrap'><div class='card'>"
        "<div class='badge'>Numerical Method: RK4</div>"
        "<h3>1. The RK4 Method</h3>"
        "<p>The Runge-Kutta 4 method uses a weighted average of four slopes (\\(k\\) values) to estimate the next point. <br>The formulas for each step are:</p>"
        "<div class='note'>"
        "$$ \\begin{aligned} "
        "k_1 &= f(x_n, y_n) \\\\"
        "k_2 &= f(x_n + \\tfrac{h}{2}, y_n + \\tfrac{h}{2}k_1) \\\\"
        "k_3 &= f(x_n + \\tfrac{h}{2}, y_n + \\tfrac{h}{2}k_2) \\\\"
        "k_4 &= f(x_n + h, y_n + h k_3) \\\\"
        "y_{n+1} &= y_n + \\frac{h}{6}(k_1 + 2k_2 + 2k_3 + k_4)"
        "\\end{aligned} $$"
        "</div>"
        f"{step_one_explanation}"
        "<h4>Summary Table (All Steps)</h4>"
        f"{table_html}"
        "</div></div>"
        f"{explain_rk4_derivation()}"
    )

    answer_s = f"After {n_steps} steps, \\( y({_fmt_num(x_end)}) \\approx {L(_fmt_num(y_end))} \\)."

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




def explain_rk4_derivation() -> str:
    """
    Returns an HTML/LaTeX explanation of the derivation/intuition behind RK4,
    connecting it to Simpson's Rule.
    """

    # We don't need dynamic variables for the derivation text,
    # but we use LaTeX for clear mathematical formatting.

    derivation_html = (
        "<div class='wrap'><div class='card'>"
        "<div class='badge'>Derivation & Intuition</div>"
        "<h3>Where does RK4 come from?</h3>"

        "<p>The Runge-Kutta 4th Order method is essentially an application of <b>Simpson's Rule</b> for numerical integration to differential equations.</p>"

        "<h4>1. The Integral Form</h4>"
        "<p>From the Fundamental Theorem of Calculus, we can write the solution as an integral:</p>"
        "$$ y(x_{n+1}) - y(x_n) = \\int_{x_n}^{x_{n+1}} f(x, y(x)) \\, dx $$"
        "$$ y_{n+1} = y_n + \\int_{x_n}^{x_n + h} f(x, y(x)) \\, dx $$"

        "<h4>2. Simpson's Rule Analogy</h4>"
        "<p>Simpson's Rule approximates an integral \\(\\int_a^b g(x)dx\\) using the values at the endpoints and the midpoint:</p>"
        "$$ \\text{Integral} \\approx \\frac{b-a}{6} \\left[ g(a) + 4g\\left(\\frac{a+b}{2}\\right) + g(b) \\right] $$"
        "<p>For our ODE step of size \\(h\\), this translates to:</p>"
        "$$ y_{n+1} \\approx y_n + \\frac{h}{6} \\left[ \\text{slope}_{start} + 4(\\text{slope}_{mid}) + \\text{slope}_{end} \\right] $$"

        "<h4>3. Estimating the Slopes (The 4 Steps)</h4>"
        "<p>Since we don't know \\(y(x)\\) inside the interval, we must estimate these slopes. RK4 does this in four stages:</p>"

        "<ul>"
        "<li><b>\\(k_1\\) (Start Slope):</b> The slope at the beginning of the interval.<br>"
        "\\( k_1 = f(x_n, y_n) \\)</li>"

        "<li><b>\\(k_2\\) (First Midpoint Estimate):</b> We use \\(k_1\\) to take a half-step (Euler step) to the midpoint and estimate the slope there.<br>"
        "\\( k_2 = f(x_n + \\frac{h}{2}, y_n + \\frac{h}{2}k_1) \\)</li>"

        "<li><b>\\(k_3\\) (Second Midpoint Estimate):</b> We use the slope \\(k_2\\) (which is likely more accurate than \\(k_1\\)) to redo the half-step and get a better midpoint slope.<br>"
        "\\( k_3 = f(x_n + \\frac{h}{2}, y_n + \\frac{h}{2}k_2) \\)</li>"

        "<li><b>\\(k_4\\) (End Slope):</b> We use \\(k_3\\) to take a full step to the end of the interval.<br>"
        "\\( k_4 = f(x_n + h, y_n + h k_3) \\)</li>"
        "</ul>"

        "<h4>4. The Final Weighted Average</h4>"
        "<p>Finally, RK4 combines these slopes. It splits the 'middle' weight of 4 from Simpson's Rule into two parts (\\(2k_2 + 2k_3\\)) to account for both midpoint estimates:</p>"

        "$$ \\text{Average Slope} = \\frac{1}{6}(k_1 + 2k_2 + 2k_3 + k_4) $$"

        "<p>This matches the Simpson's coefficients \\(1 : 4 : 1\\) (where the middle 4 is split \\(2+2\\)).</p>"

        "</div></div>"
    )

    return derivation_html
