import numpy as np
import matplotlib.pyplot as plt
from numpy.polynomial import legendre
import pandas as pd

plt.rcParams.update(
    {
        "savefig.dpi": 300,
        "figure.autolayout": True,
        "text.usetex": True,
        "font.family": "serif",
        "font.serif": ["Times New Roman"],
        "pdf.fonttype": 42,
        "ps.fonttype": 42,
        "font.size": 14,
        "axes.labelsize": 14,
        "axes.titlesize": 14,
        "xtick.labelsize": 14,
        "ytick.labelsize": 14,
        "legend.fontsize": 14,
        "lines.linewidth": 1.0,
        "lines.markersize": 5,
        "axes.linewidth": 1.0,
        "lines.markerfacecolor": "none",
        "lines.markeredgecolor": "auto",
        "xtick.direction": "in",
        "ytick.direction": "in",
        "legend.facecolor": "white",
        "legend.edgecolor": "black",
        "legend.framealpha": 1.0,
    }
)

def burgers_f(u):
    """Flux function f(u) = u^2/2"""
    return u**2 / 2


def burgers_df(u):
    """Flux derivative f'(u) = u"""
    return u

def burgers_sin_solver_kernel(x, t):
    """
    Solve the sine-wave initial value problem for Burgers equation.

    Parameters
    ----------
    x : array
        Spatial coordinates in [-pi, pi]
    t : float
        Time

    Returns
    -------
    u : array
        Solution u(x,t) = sin(x - u*t)
    """
    u = x / (np.pi / 2 + t)  # initial guess

    # Newton iteration
    for _ in range(100000):
        du = (u - np.sin(x - u * t)) / (1 + np.cos(x - u * t) * t)
        u = u - du

        if np.max(np.abs(du)) < 1e-10:
            break

    return u


def burgers_sin_solver(x, t, a, b, w, phi):
    """
    Solve Burgers equation with sinusoidal initial condition.

    u(x,0) = a + b*sin(w*x + phi)

    Parameters
    ----------
    x : array
        Spatial coordinates
    t : float
        Time
    a : float
        Mean value
    b : float
        Amplitude
    w : float
        Angular frequency
    phi : float
        Phase shift

    Returns
    -------
    u : array
        Solution at (x,t)
    """
    x_tilde = w * x + phi - a * w * t
    x_tilde = np.mod(x_tilde + np.pi, 2 * np.pi) - np.pi  # map to [-pi, pi]
    t_tilde = b * w * t

    return a + b * burgers_sin_solver_kernel(x_tilde, t_tilde)

fig, ax = plt.subplots()

for tend in [0, 0.5, 1.0, 2.0, 3.0]:
    x, dx = np.linspace(0, 2 * np.pi, 100, endpoint=False, retstep=True)
    x += dx / 2

    u = burgers_sin_solver(x, tend, a=0.0, b=1.0, w=1.0, phi=0)

    if tend >= 1.0:
        u0 = lambda s: np.sin(s)
        ax.plot(x + u0(x) * tend, u0(x), "--")

    ax.plot(x, u, label=f"t={tend}")

ax.legend()

def SSP_RK3_with_post_process(u, t, dt, op_L, post_process, **kwargs):
    """
    SSP-RK3 time stepping with optional post-processing.

    Parameters
    ----------
    u : array
        Current solution
    t : float
        Current time
    dt : float
        Time step
    op_L : callable
        Spatial operator L(u, t)
    post_process : callable
        Post-processing operator (e.g., limiter)

    Returns
    -------
    unew : array
        Solution at next time level
    """
    u1_pre = u + dt * op_L(u, t, **kwargs)
    u1 = post_process(u1_pre)
    u2_pre = 3 / 4 * u + 1 / 4 * (u1 + dt * op_L(u1, t + dt, **kwargs))
    u2 = post_process(u2_pre)
    unew_pre = 1 / 3 * u + 2 / 3 * (u2 + dt * op_L(u2, t + dt / 2, **kwargs))
    unew = post_process(unew_pre)

    return unew

def get_LF_flux(f, get_alpha):
    """
    Create Lax-Friedrichs numerical flux function.

    Parameters
    ----------
    f : callable
        Physical flux function
    get_alpha : callable
        Function to compute maximum wave speed

    Returns
    -------
    fhat_lf : callable
        Numerical flux function
    """
    def fhat_lf(ul, ur):
        return (f(ul) + f(ur)) / 2 - get_alpha(ul, ur) / 2 * (ur - ul)

    return fhat_lf


def burgers_get_alpha(ul, ur):
    """Maximum wave speed for Burgers equation."""
    return np.maximum(np.abs(burgers_df(ul)), np.abs(burgers_df(ur)))

def gauss_quadrature(f, xl, xr, gk):
    """
    Compute integral using Gaussian quadrature.

    Parameters
    ----------
    f : callable
        Function to integrate
    xl, xr : float
        Integration interval
    gk : int
        Number of quadrature points

    Returns
    -------
    result : float or array
        Integral value
    """
    x = (xl + xr) / 2
    dx = xr - xl

    scalar_flag = np.isscalar(x)
    x = np.atleast_1d(x)  # shape = (n,)

    points, weights = legendre.leggauss(gk)  # shape = (gk,), (gk,)

    # shape = (n, np.newaxis) + (gk,) = (n, gk)
    x_nodes = x[:, np.newaxis] + dx / 2 * points
    # shape = (n, gk) * (gk,) = (n, gk)
    integral_values = f(x_nodes) * weights
    result = dx / 2 * np.sum(integral_values, axis=1)  # shape = (n,)

    return result if not scalar_flag else result[0]  # shape = (n,) or scalar

def basis_eval(k, points):
    """
    Evaluate Legendre basis functions.

    Parameters
    ----------
    k : int
        Polynomial degree
    points : array
        Points in [-1, 1]

    Returns
    -------
    result : array
        Shape (k+1, len(points)) or (k+1,) if scalar
    """
    scalar_flag = np.isscalar(points)
    points = np.atleast_1d(points)

    coeff = np.identity(k + 1)
    result = np.zeros((k + 1, len(points)))

    for i in range(k + 1):
        result[i] = legendre.Legendre(coeff[i])(points)

    # shape = (k+1, len(points)) or (k+1,)
    return result[:, 0] if scalar_flag else result

def basis_dx_eval(k, points):
    """
    Evaluate derivatives of Legendre basis functions.

    Parameters
    ----------
    k : int
        Polynomial degree
    points : array
        Points in [-1, 1]

    Returns
    -------
    result : array
        Shape (k+1, len(points)) or (k+1,) if scalar
    """
    scalar_flag = np.isscalar(points)
    points = np.atleast_1d(points)

    coeff = np.identity(k + 1)
    result = np.zeros((k + 1, len(points)))

    for i in range(1, k + 1):
        f1 = legendre.Legendre(coeff[i])
        f2 = legendre.Legendre(coeff[i - 1])
        result[i] = i / (points**2 - 1) * (points * f1(points) - f2(points))

    # shape = (k+1, len(points)) or (k+1,)
    return result[:, 0] if scalar_flag else result

def basis_inner_vec(k, gk):
    """
    Compute L2 inner products of basis functions.

    Parameters
    ----------
    k : int
        Polynomial degree
    gk : int
        Quadrature order

    Returns
    -------
    result : array
        Vector of inner products (k+1,)
    """
    coeff = np.identity(k + 1)

    result = np.zeros(k + 1)
    for i in range(k + 1):
        f = lambda t: legendre.Legendre(coeff[i])(t) ** 2
        result[i] = gauss_quadrature(f, -1, 1, gk)  # = 2 / (2 * i + 1)

    return result

def l2_projection(f, x, dx, k, gk):
    """
    L2 projection of function f onto piecewise polynomial space.

    Parameters
    ----------
    f : callable
        Function to project
    x : array
        Cell centers
    dx : float
        Cell size
    k : int
        Polynomial degree
    gk : int
        Quadrature order

    Returns
    -------
    result : array
        Projection coefficients (n, k+1) or (k+1,) if scalar
    """
    scalar_flag = np.isscalar(x)
    x = np.atleast_1d(x)

    n = len(x)
    coeff = np.identity(k + 1)

    # Denominator: integral of basis function squared
    inners = np.zeros(k + 1)
    for i in range(k + 1):
        f2 = lambda t: legendre.Legendre(coeff[i])(t) ** 2
        inners[i] = dx / 2 * gauss_quadrature(f2, -1, 1, gk)

    # Numerator: integral of f * basis
    result = np.zeros([n, k + 1])
    for j in range(n):
        for i in range(k + 1):
            f1 = lambda t: f(t) * legendre.Legendre(coeff[i])(2 * (t - x[j]) / dx)
            result[j][i] = (
                gauss_quadrature(f1, x[j] - dx / 2, x[j] + dx / 2, gk) / inners[i]
            )

    # shape = (n, k+1) or (k+1,)
    return result if not scalar_flag else result[0]

def error4DG(uexact, uh, x, dx, t, gk):
    """
    Compute L1, L2, and Linf errors for DG solution.

    Parameters
    ----------
    uexact : callable
        Exact solution function u(x, t)
    uh : array
        DG coefficients (n, k+1)
    x : array
        Cell centers
    dx : float
        Cell size
    t : float
        Time
    gk : int
        Quadrature order

    Returns
    -------
    err_inf, err_l1, err_l2 : float
        L∞, L1, and L2 errors
    """
    points, weights = legendre.leggauss(gk)
    err_inf, err_l1, err_l2 = 0, 0, 0

    for xi, ui in zip(x, uh):
        u = uexact(xi + dx / 2 * points, t)
        uh_val = legendre.Legendre(ui)(points)

        err_inf = np.max([err_inf, np.max(np.abs(u - uh_val))])
        err_l1 += np.abs(u - uh_val).dot(weights)
        err_l2 += ((u - uh_val) ** 2).dot(weights)

    err_l1 *= dx / 2
    err_l2 = np.sqrt(dx / 2 * err_l2)
    return err_inf, err_l1, err_l2

def minmod(a1, a2, a3, h):
    """
    Minmod TVD limiter.

    Parameters
    ----------
    a1, a2, a3 : array
        Input values
    h : float
        Cell size (unused)

    Returns
    -------
    result : array
        Minmod of inputs
    """
    pos_mask = (a1 > 0) & (a2 > 0) & (a3 > 0)
    neg_mask = (a1 < 0) & (a2 < 0) & (a3 < 0)

    result = np.where(pos_mask, np.minimum(np.minimum(a1, a2), a3), 0)
    result = np.where(neg_mask, np.maximum(np.maximum(a1, a2), a3), result)

    return result


def tvd():
    """Return minmod as TVD limiter."""
    return minmod


def tvb(m):
    """
    Create TVB minmod limiter.

    Parameters
    ----------
    m : float
        TVB constant

    Returns
    -------
    minmod_tilde : callable
        TVB limiter function
    """

    def minmod_tilde(a1, a2, a3, h):
        condition = np.abs(a1) < m * h**2
        return np.where(condition, a1, minmod(a1, a2, a3, h))

    return minmod_tilde

def dg_lim(lim, k, gk, h):
    """
    Create DG limiter post-processor.

    Parameters
    ----------
    lim : callable
        TVD/TVB limiter function
    k : int
        Polynomial degree
    gk : int
        Quadrature order
    h : float
        Cell size

    Returns
    -------
    processer : callable
        Post-processing function
    """
    assert k >= 1

    points, weights = legendre.leggauss(gk)  # shape = # (gk,), (gk,)
    M = basis_eval(k, points)  # shape = (k+1, gk)
    vm = 1 / 2 * M.dot(weights)  # cell averages, shape = (k+1,)

    vl = basis_eval(k, -1)  # left boundary, shape = (k+1,)
    vr = basis_eval(k, 1)  # right boundary

    if k == 1:
        Mat2 = np.array([[vm[0], vm[1]], [vl[0], vl[1]]])
    else:
        Mat3 = np.array(
            [
                [vm[0], vm[1], vm[2]],
                [vl[0], vl[1], vl[2]],
                [vr[0], vr[1], vr[2]],
            ]
        )

    def processer(uh):
        """Apply limiter to DG coefficients."""
        ul = uh.dot(vl)  # shape = (n,)
        um = uh.dot(vm)
        ur = uh.dot(vr)

        um_delta_r = np.roll(um, -1) - um
        um_delta_l = um - np.roll(um, 1)

        ul_mod = um - lim(um - ul, um_delta_l, um_delta_r, h)  # shape = (n,)
        ur_mod = um + lim(ur - um, um_delta_l, um_delta_r, h)

        uh_mod = np.zeros_like(uh)  # shape = (n, k+1)
        if k == 1:
            B = np.stack([um, ul_mod], axis=0)  # shape = (2, n)
            uh_mod = np.linalg.solve(Mat2, B).T  # shape = (n, 2)
        else:
            B = np.stack([um, ul_mod, ur_mod], axis=0)  # shape = (3, n)
            uh_mod[:, :3] = np.linalg.solve(Mat3, B).T  # shape = (n, 3)

        return uh_mod

    return processer

def RKDG(uh0, dx, tend, k, gk, f, fhat, df, **kwargs):
    """
    Runge-Kutta Discontinuous Galerkin solver.

    Parameters
    ----------
    uh0 : array
        Initial DG coefficients (n, k+1)
    dx : float
        Cell size
    tend : float
        Final time
    k : int
        Polynomial degree
    gk : int
        Quadrature order
    f : callable
        Physical flux function
    fhat : callable
        Numerical flux function
    df : callable
        Flux derivative for CFL
    **kwargs : dict
        Additional options (e.g., lim for limiter)

    Returns
    -------
    uh : array
        DG coefficients at final time
    """
    points, weights = legendre.leggauss(gk)  # shape = (gk,), (gk,)

    # Basis matrices
    M = basis_eval(k, points)  # shape = (k+1, gk)
    D = (2 / dx) * basis_dx_eval(k, points)  # shape = (k+1, gk)
    M_inv = (2 / dx) * np.diag(1.0 / basis_inner_vec(k, gk))  # shape =(k+1, k+1)

    # Boundary evaluation matrices
    vl = basis_eval(k, -1)  # shape = (k+1,)
    vr = basis_eval(k, 1)
    vc = basis_eval(k, 0)

    def op_L(uh, t):
        """Spatial operator L(u, t)."""
        # Neighbor cells
        uh_l = np.roll(uh, shift=1, axis=0)  # shape = (n, k+1)
        uh_r = np.roll(uh, shift=-1, axis=0)

        # Boundary values
        ul_m = uh_l.dot(vr)  # left value from left neighbor,shape = (n,)
        ul_p = uh.dot(vl)  # left value from current cell
        ur_m = uh.dot(vr)  # right value from current cell
        ur_p = uh_r.dot(vl)  # right value from right neighbor

        # Numerical fluxes
        fhat_l = fhat(ul_m, ul_p)  # shape = (n,)
        fhat_r = fhat(ur_m, ur_p)

        # Volume term: integral of f(u) * v_x
        tmp_u = uh @ M  # (n, k+1) @ (k+1, gk) = (n, gk)
        tmp_fu = f(tmp_u) * weights  # (n, gk) * (gk,) = (n, gk)
        int_fu_vx = (dx / 2) * (tmp_fu @ D.T)  # (n, gk) @ (gk, k+1) = (n, k+1)

        # Flux terms at boundaries
        fhatl_vl = fhat_l[:, np.newaxis] * vl  # (n, np.newaxis) * (k+1,) = (n, k+1)
        fhatr_vr = fhat_r[:, np.newaxis] * vr

        # Residual
        # (n, k+1) @ (k+1, k+1) = (n, k+1)
        delta = (int_fu_vx - fhatr_vr + fhatl_vl) @ M_inv.T
        return delta

    # Post-processing (limiter)
    if "lim" in kwargs and k > 0:
        post_processer = dg_lim(kwargs["lim"], k, gk, dx)
    else:
        post_processer = lambda u: u

    # Time stepping
    cfl = 1 / (2 * k + 1)
    dt0 = cfl * (dx ** ((k + 1) / 3)) if k > 2 else cfl * dx

    uh = uh0.copy()
    t = 0
    while t < tend:
        # Adaptive time step based on CFL
        alpha_global = np.max(np.abs(df(uh.dot(vc))))
        dt = min(dt0 / alpha_global, tend - t)
        uh = SSP_RK3_with_post_process(uh, t, dt, op_L, post_processer)
        t += dt

    return uh

def burgers_init(x):
    """Initial condition: u(x,0) = 1/2 + sin(x)"""
    return burgers_sin_solver(x, t=0, a=0.5, b=1, w=1, phi=0)


def burgers_exact(x, t):
    """Exact solution using method of characteristics."""
    return burgers_sin_solver(x, t=t, a=0.5, b=1, w=1, phi=0)

def order(xlist, ylist):
    """Compute convergence order from error data."""
    return -np.log(ylist[1:] / ylist[:-1]) / np.log(xlist[1:] / xlist[:-1])


def err_table(nums, err_l1, err_l2, err_inf):
    """Create error table with convergence orders."""
    ord_l1 = order(nums, err_l1)
    ord_l2 = order(nums, err_l2)
    ord_linf = order(nums, err_inf)

    data = {
        "m": nums,
        "L1 error": err_l1,
        "L1 order": np.r_[np.nan, ord_l1],
        "L2 error": err_l2,
        "L2 order": np.r_[np.nan, ord_l2],
        "Linf error": err_inf,
        "Linf order": np.r_[np.nan, ord_linf],
    }

    return pd.DataFrame(data)


def order_test(*, init, exact, xleft, xright, nlist, tend, k, gk=7, **kwargs):
    """
    Convergence order test.

    Parameters
    ----------
    init : callable
        Initial condition
    exact : callable
        Exact solution
    xleft, xright : float
        Domain bounds
    nlist : array
        List of mesh sizes
    tend : float
        Final time
    k : int
        Polynomial degree
    gk : int
        Quadrature order
    **kwargs : dict
        Additional DG parameters
    """
    err_inf = np.zeros_like(nlist, dtype=np.float64)
    err_l1 = np.zeros_like(nlist, dtype=np.float64)
    err_l2 = np.zeros_like(nlist, dtype=np.float64)

    for cnt, n in enumerate(nlist):
        x, dx = np.linspace(xleft, xright, n, endpoint=False, retstep=True)
        x += dx / 2

        u0 = l2_projection(init, x, dx, k, gk)
        uh = RKDG(u0, dx, tend, k, gk, **kwargs)
        err_inf[cnt], err_l1[cnt], err_l2[cnt] = error4DG(exact, uh, x, dx, tend, gk)

    df = err_table(nlist, err_l1, err_l2, err_inf)
    print(df.to_string(index=False, float_format="%.2e"))

def shock_test(*, init, exact, xleft, xright, nlist, n_ref, tend, k, gk=7, **kwargs):
    """
    Shock capturing test: compare DG solution with reference exact solution.
    """
    x_ref, dx_ref = np.linspace(xleft, xright, n_ref, endpoint=False, retstep=True)
    x_ref += dx_ref / 2

    u_exact_ref = exact(x_ref, tend)

    for n in nlist:
        x, dx = np.linspace(xleft, xright, n, endpoint=False, retstep=True)
        x += dx / 2

        u0 = l2_projection(init, x, dx, k, gk)
        uh = RKDG(u0, dx, tend, k, gk, **kwargs)

        # Evaluate at cell centers
        uh_point_value = np.zeros_like(x)
        for i in range(len(x)):
            uh_point_value[i] = legendre.Legendre(uh[i])(0)

        fig, ax = plt.subplots(figsize=(6, 5))
        ax.plot(x_ref, u_exact_ref, label=r"$u_{exact}$")
        ax.plot(
            x,
            uh_point_value,
            linestyle="None",
            marker="o",
            markersize=5,
            markerfacecolor="none",
            label=r"$u_h$",
        )
        ax.legend()
        ax.grid()

    plt.show()

# DG P0
config = {
    "k": 0,
    "init": burgers_init,
    "exact": burgers_exact,
    "xleft": -np.pi,
    "xright": np.pi,
    "f": burgers_f,
    "df": burgers_df,
    "fhat": get_LF_flux(burgers_f, burgers_get_alpha),
}

print("=== DG P0: Convergence Test (t=0.5) ===")
order_test(tend=0.5, nlist=np.array([20, 40, 80, 160, 320]), **config)
print("\n=== DG P0: Shock Test (t=1.5) ===")
shock_test(tend=1.5, nlist=np.array([40, 160]), n_ref=320, **config)

# DG P1
config = {
    "k": 1,
    "init": burgers_init,
    "exact": burgers_exact,
    "xleft": -np.pi,
    "xright": np.pi,
    "f": burgers_f,
    "df": burgers_df,
    "fhat": get_LF_flux(burgers_f, burgers_get_alpha),
}

print("=== DG P1: Convergence Test (t=0.5) ===")
order_test(tend=0.5, nlist=np.array([20, 40, 80, 160, 320]), **config)
print("\n=== DG P1: Shock Test (t=1.5) ===")
shock_test(tend=1.5, nlist=np.array([40, 160]), n_ref=320, **config)

# DG P2
config = {
    "k": 2,
    "init": burgers_init,
    "exact": burgers_exact,
    "xleft": -np.pi,
    "xright": np.pi,
    "f": burgers_f,
    "df": burgers_df,
    "fhat": get_LF_flux(burgers_f, burgers_get_alpha),
}

print("=== DG P2: Convergence Test (t=0.5) ===")
order_test(tend=0.5, nlist=np.array([20, 40, 80, 160, 320]), **config)
print("\n=== DG P2: Shock Test (t=1.5) ===")
shock_test(tend=1.5, nlist=np.array([40, 160]), n_ref=320, **config)

# DG P2 + TVD limiter
config = {
    "k": 2,
    "init": burgers_init,
    "exact": burgers_exact,
    "xleft": -np.pi,
    "xright": np.pi,
    "f": burgers_f,
    "df": burgers_df,
    "fhat": get_LF_flux(burgers_f, burgers_get_alpha),
    "lim": tvd(),
}

print("=== DG P2 + TVD: Convergence Test (t=0.5) ===")
order_test(tend=0.5, nlist=np.array([20, 40, 80, 160, 320]), **config)
print("\n=== DG P2 + TVD: Shock Test (t=1.5) ===")
shock_test(tend=1.5, nlist=np.array([40, 160]), n_ref=320, **config)

# DG P2 + TVB limiter (M=5)
config = {
    "k": 2,
    "init": burgers_init,
    "exact": burgers_exact,
    "xleft": -np.pi,
    "xright": np.pi,
    "f": burgers_f,
    "df": burgers_df,
    "fhat": get_LF_flux(burgers_f, burgers_get_alpha),
    "lim": tvb(5),
}

print("=== DG P2 + TVB (M=5): Convergence Test (t=0.5) ===")
order_test(tend=0.5, nlist=np.array([20, 40, 80, 160, 320]), **config)
print("\n=== DG P2 + TVB (M=5): Shock Test (t=1.5) ===")
shock_test(tend=1.5, nlist=np.array([40, 160]), n_ref=320, **config)