"""
pde_solver.py — A PDE solver whose matmul goes through the VGPU cache.

The PDE is described in a small text language; the parser builds the
discretised grid and operator matrix, the solver then time-steps (or
solves directly) so every matrix multiplication reuses the compiled
VGPU kernel for that shape.

Run it from the shell:

    python pde_solver.py                                                            # help
    python pde_solver.py "heat 1D alpha=0.1 L=1 T=0.5 nx=100 dt=0.001 ic=sin(pi*x)"
    python pde_solver.py "heat 2D alpha=0.1 L=1 H=1 T=0.1 nx=40 ny=40 dt=0.0001 ic=sin(pi*x)*sin(pi*y)"
    python pde_solver.py "wave 1D c=1 L=2 T=2 nx=200 ic=sin(pi*x) dt=0.005"
    python pde_solver.py "poisson 1D L=1 nx=100 source=sin(pi*x) bc=dirichlet"
    python pde_solver.py "du/dt = 0.1*d2u/dx2 - 0.5*du/dx + 0.2*u"
    python pde_solver.py "du/dt = D*d2u/dx2  D=0.05"

Or import it:

    from pde_solver import solve, PDEParser, cache
    res = solve("heat 1D alpha=0.1 L=1 T=0.5 nx=100 dt=0.001 ic=sin(pi*x)")
"""
import re
import sys
import numpy as np

try:
    from vgpu_cache import VGPUCache
except ImportError as e:
    sys.stderr.write(
        "pde_solver requires vgpu_cache.py in the same directory.\n")
    raise

# A single shared cache instance — every matmul benefits from kernel reuse.
cache = VGPUCache()

# ================================================================
#  Section 1 — Expression compiler
# ================================================================
# Strings like 'sin(pi*x)*exp(-(x-0.5)^2)' compile into numpy-vectorised
# functions so the parser can accept any IC / source term.
_FN_NAMES = {
    'sin': 'np.sin',   'cos': 'np.cos',   'tan': 'np.tan',
    'sinh': 'np.sinh', 'cosh': 'np.cosh', 'tanh': 'np.tanh',
    'exp': 'np.exp',   'log': 'np.log',   'sqrt': 'np.sqrt',
    'abs': 'np.abs',
}

def _preprocess(expr: str) -> str:
    code = expr.strip().replace('^', '**')
    code = re.sub(r'\be\b(?!xp)', 'np.e', code)
    code = re.sub(r'\bpi\b', 'np.pi', code)
    for name, repl in _FN_NAMES.items():
        code = re.sub(r'\b' + name + r'\(', repl + '(', code)
    return code

def make_fn(expr: str, var_names=('x',)):
    """Compile a math expression into a vectorised callable."""
    code = _preprocess(expr)
    var_names = tuple(var_names)
    def fn(*args):
        if len(args) != len(var_names):
            raise ValueError(
                f"function expects {len(var_names)} arguments, got {len(args)}")
        env = {'np': np}
        for name, val in zip(var_names, args):
            env[name] = np.asarray(val, dtype=np.float32)
        return eval(code, {'__builtins__': {}}, env)
    fn.src = expr
    return fn


# ================================================================
#  Section 2 — PDE specification & parser
# ================================================================
class PDESpec:
    """Structured PDE parsed from a text string."""
    __slots__ = ('name', 'dim', 'params', 'bc', 'ic', 'source',
                 'equation', 'symbolic')
    def __init__(self):
        self.name      = ''
        self.dim       = 1
        self.params    = {}
        self.bc        = 'dirichlet'
        self.ic        = None
        self.source    = None
        self.equation  = ''
        self.symbolic  = None
    def __repr__(self):
        return f"<PDESpec {self.name or 'symbolic'} {self.dim}D {dict(self.params)}>"


class PDEParser:
    """Turn a text PDE description into a PDESpec.

    Two grammar forms are accepted.

    Preset form (named PDE):
        <name> <dim>D [ key=value ]...
        e.g.  'heat 1D alpha=0.1 L=1 T=0.5 nx=100 ic=sin(pi*x)'

    Symbolic form (free-form 1D linear PDE):
        du/dt = <coeff>*<op> [+|- <coeff>*<op>] ...
        e.g.  'du/dt = 0.1*d2u/dx2 - 0.5*du/dx - 0.2*u'
        Recognised operators: d2u/dx2, du/dx, u
        Coefficients may be numbers (0.5) or names (D, k)
        referenced from the same spec (e.g. D=0.05).
    """
    PRESETS = {'heat', 'wave', 'poisson', 'advection', 'burgers'}

    @staticmethod
    def parse(text: str) -> PDESpec:
        text = text.strip()
        if not text:
            raise ValueError("empty PDE spec")
        first = text.split()[0]
        if first in PDEParser.PRESETS:
            return PDEParser._parse_preset(text)
        if '=' in text and re.search(r'\bd[2]?u(?:/dx2?|/dt)?\b', text):
            return PDEParser._parse_symbolic(text)
        raise ValueError(f"cannot parse PDE spec: {text!r}")

    @staticmethod
    def _split_args(tokens):
        out = {}
        for tok in tokens:
            m = re.match(r'([A-Za-z_]\w*)=(.+)$', tok)
            if not m:
                raise ValueError(f"bad argument token: {tok!r}")
            key, raw = m.group(1), m.group(2)
            try:
                out[key] = float(raw)
            except ValueError:
                out[key] = raw
        return out

    @staticmethod
    def _parse_preset(text: str) -> PDESpec:
        tokens = text.split()
        spec = PDESpec()
        spec.name = tokens[0]
        dim_tok = tokens[1]
        if dim_tok not in ('1D', '2D'):
            raise ValueError(f"expected dimension token (1D/2D), got {dim_tok!r}")
        spec.dim = int(dim_tok[0])
        spec.params = PDEParser._split_args(tokens[2:])
        spec.bc = str(spec.params.get('bc', 'dirichlet'))
        var_names = ['x'] if spec.dim == 1 else ['x', 'y']
        if isinstance(spec.params.get('ic'), str):
            spec.ic = make_fn(spec.params['ic'], var_names)
        if isinstance(spec.params.get('source'), str):
            spec.source = make_fn(spec.params['source'], var_names)
        if spec.ic is None:
            spec.ic = make_fn('sin(pi*x)' if spec.dim == 1
                              else 'sin(pi*x)*sin(pi*y)', var_names)
        return spec

    @staticmethod
    def _parse_symbolic(text: str) -> PDESpec:
        spec = PDESpec()
        spec.name = 'symbolic-linear'
        spec.dim  = 1
        # Pull the first '=' so key=val params like 'D=0.05' stay intact.
        lhs, rhs = text.split('=', 1)
        spec.equation = lhs + '=' + rhs

        # Pull side parameters off the RHS first so they don't confuse
        # the term parser (e.g. 'D*d2u/dx2  D=0.05').
        param_pat = re.compile(r'\b([A-Za-z_]\w*)=(-?\d+(?:\.\d+)?)\b')
        spec.params = {m.group(1): float(m.group(2))
                       for m in param_pat.finditer(rhs)}
        eq_rhs = param_pat.sub('', rhs)

        # Term grammar:  [sign] coefficient '*' operator
        term_pat = re.compile(
            r'(?P<sign>[+-]?)'
            r'\s*(?:(?P<coef>\d+\.?\d*|[A-Za-z_]\w*)\s*\*?\s*)?'
            r'(?P<op>d2u/dx2|du/dx|u)(?!\w)'
        )
        coeffs = {}
        for m in term_pat.finditer(eq_rhs):
            sign = -1.0 if m.group('sign') == '-' else 1.0
            c_str = m.group('coef')
            if c_str is None:
                c = 1.0
            else:
                try:
                    c = float(c_str)
                except ValueError:
                    c = c_str  # symbolic name, resolved later
            op = m.group('op')
            coeffs[op] = coeffs.get(op, 0.0) + sign * c
        spec.symbolic = coeffs
        spec.bc = str(spec.params.get('bc', 'dirichlet'))
        if isinstance(spec.params.get('ic'), str):
            spec.ic = make_fn(spec.params['ic'], ['x'])
        else:
            spec.ic = make_fn('sin(pi*x)', ['x'])
        return spec


# ================================================================
#  Section 3 — Finite-difference operators
# ================================================================
class FD:
    """Central-difference stencils for 1D / 2D."""

    # ---- grids ---------------------------------------------------------
    @staticmethod
    def grid_1d(spec):
        L  = float(spec.params.get('L', 1.0))
        nx = int(spec.params.get('nx', 100))
        x  = np.linspace(0.0, L, nx, dtype=np.float32)
        dx = float(x[1] - x[0]) if nx > 1 else 1.0
        return x, dx

    @staticmethod
    def grid_2d(spec):
        L  = float(spec.params.get('L', 1.0))
        H  = float(spec.params.get('H', L))
        nx = int(spec.params.get('nx', 50))
        ny = int(spec.params.get('ny', nx))
        x  = np.linspace(0.0, L, nx, dtype=np.float32)
        y  = np.linspace(0.0, H, ny, dtype=np.float32)
        X, Y = np.meshgrid(x, y, indexing='xy')
        return x, y, X, Y, float(x[1] - x[0]), float(y[1] - y[0])

    # ---- operators -----------------------------------------------------
    @staticmethod
    def laplacian_1d(n, dx, bc='dirichlet'):
        """L such that (L u)_i ≈ u''(x_i) (central)."""
        L = np.zeros((n, n), dtype=np.float32)
        idx = np.arange(n)
        if bc == 'periodic':
            L[idx, (idx-1) % n]  =  1.0 / dx**2
            L[idx, idx]          = -2.0 / dx**2
            L[idx, (idx+1) % n]  =  1.0 / dx**2
        else:
            L[idx, idx]          = -2.0 / dx**2
            L[idx[1:],   idx[:-1]] =  1.0 / dx**2
            L[idx[:-1],  idx[1:]]  =  1.0 / dx**2
            if bc == 'neumann':
                # One-sided 2nd derivative at the boundary.
                L[0]    = 0
                L[0, 0] = -1.0 / dx**2
                L[0, 1] =  1.0 / dx**2
                L[-1]   = 0
                L[-1,-1] = -1.0 / dx**2
                L[-1,-2] =  1.0 / dx**2
        return L

    @staticmethod
    def laplacian_2d(nx, ny, dx, dy, bc='dirichlet'):
        Lx = FD.laplacian_1d(nx, dx, bc)
        Ly = FD.laplacian_1d(ny, dy, bc)
        Ix = np.eye(nx, dtype=np.float32)
        Iy = np.eye(ny, dtype=np.float32)
        # Vectorised in C-order: index = i*nx + j
        # ∂²/∂x² in flattened = kron(I_y, L_x); ∂²/∂y² = kron(L_y, I_x)
        return np.kron(Iy, Lx) + np.kron(Ly, Ix)

    @staticmethod
    def grad_1d(n, dx, bc='periodic'):
        D = np.zeros((n, n), dtype=np.float32)
        idx = np.arange(n)
        if bc == 'periodic':
            D[idx, (idx+1) % n]  =  0.5 / dx
            D[idx, (idx-1) % n]  = -0.5 / dx
        else:
            D[idx[1:],   idx[:-1]] = -0.5 / dx
            D[idx[:-1],  idx[1:]]  =  0.5 / dx
        return D

    @staticmethod
    def bdry_mask_1d(n, bc):
        m = np.zeros(n, dtype=bool)
        if bc in ('dirichlet', 'neumann'):
            m[0] = m[-1] = True
        return m

    @staticmethod
    def bdry_mask_2d(nx, ny, bc):
        m = np.zeros((ny, nx), dtype=bool)
        if bc in ('dirichlet', 'neumann'):
            m[0, :]    = m[-1, :]   = True
            m[:, 0]    = m[:, -1]   = True
        return m


# ================================================================
#  Section 4 — Solvers
# ================================================================
def _bc_reset(A, mask):
    """Set boundary rows of A to identity (Dirichlet)."""
    A[mask, :] = 0
    A[mask, mask] = 1
    return A

def _step_dt(spec):  return float(spec.params.get('dt', 1e-3))
def _step_tend(spec): return float(spec.params.get('T', 1.0))

def solve_heat_1d(spec):
    """u_t = α u_xx   (forward Euler)."""
    x, dx = FD.grid_1d(spec)
    nx = x.size
    alpha = float(spec.params.get('alpha', 0.1))
    dt = _step_dt(spec); T = _step_tend(spec)

    L_op = FD.laplacian_1d(nx, dx, spec.bc)
    A = np.eye(nx, dtype=np.float32) + alpha * dt * L_op
    if spec.bc == 'dirichlet':
        _bc_reset(A, FD.bdry_mask_1d(nx, 'dirichlet'))

    cfl = alpha * dt / (dx * dx)
    if cfl > 0.5:
        print(f"  [warn] CFL α·dt/dx² = {cfl:.3f} > 0.5 — explicit Euler unstable")

    u = np.asarray(spec.ic(x), dtype=np.float32)
    if spec.bc == 'dirichlet':
        u[0] = u[-1] = 0
    steps = max(1, int(round(T / dt)))
    snaps = [u.copy()]
    snap_every = max(1, steps // 12)
    for s in range(1, steps + 1):
        u = cache.matmul(A, u.reshape(nx, 1)).flatten()
        if s % snap_every == 0 or s == steps:
            snaps.append(u.copy())
    return {'x': x, 'snaps': snaps, 'spec': spec}

def solve_heat_2d(spec):
    """u_t = α (u_xx + u_yy)   (forward Euler)."""
    x, y, X, Y, dx, dy = FD.grid_2d(spec)
    nx, ny = x.size, y.size
    n = nx * ny
    alpha = float(spec.params.get('alpha', 0.1))
    dt = _step_dt(spec); T = _step_tend(spec)

    L_op = FD.laplacian_2d(nx, ny, dx, dy, spec.bc)
    A = np.eye(n, dtype=np.float32) + alpha * dt * L_op
    if spec.bc == 'dirichlet':
        _bc_reset(A, FD.bdry_mask_2d(nx, ny, 'dirichlet').flatten())

    cfl_x = alpha * dt / (dx * dx)
    cfl_y = alpha * dt / (dy * dy)
    if max(cfl_x, cfl_y) > 0.5:
        print(f"  [warn] CFL > 0.5  (α·dt/dx²={cfl_x:.3f},"
              f" α·dt/dy²={cfl_y:.3f}) — explicit Euler unstable")

    u0 = np.asarray(spec.ic(X, Y), dtype=np.float32).reshape(-1)
    if spec.bc == 'dirichlet':
        u0[FD.bdry_mask_2d(nx, ny, 'dirichlet').flatten()] = 0
    u = u0.reshape(n, 1).copy()

    steps = max(1, int(round(T / dt)))
    snaps = [u0.copy()]
    snap_every = max(1, steps // 6)
    for s in range(1, steps + 1):
        u = cache.matmul(A, u)
        if s % snap_every == 0 or s == steps:
            snaps.append(u.flatten().copy())
    return {'x': x, 'y': y, 'X': X, 'Y': Y, 'snaps': snaps, 'spec': spec}

def solve_wave_1d(spec):
    """u_tt = c² u_xx   (kick-drift)."""
    x, dx = FD.grid_1d(spec)
    nx = x.size
    c = float(spec.params.get('c', 1.0))
    dt = _step_dt(spec); T = _step_tend(spec)

    L_op = FD.laplacian_1d(nx, dx, spec.bc)
    c2L  = (c * c) * L_op
    if spec.bc == 'dirichlet':
        _bc_reset(c2L, FD.bdry_mask_1d(nx, 'dirichlet'))

    cfl = c * dt / dx
    if cfl > 1.0:
        print(f"  [warn] wave CFL c·dt/dx = {cfl:.3f} > 1 — unstable")

    u = np.asarray(spec.ic(x), dtype=np.float32)
    v = np.zeros_like(u)
    if spec.bc == 'dirichlet':
        u[0] = u[-1] = 0

    snaps = [u.copy()]
    steps = max(1, int(round(T / dt)))
    snap_every = max(1, steps // 12)
    for s in range(1, steps + 1):
        a = cache.matmul(c2L, u.reshape(nx, 1)).flatten()
        v += a * dt
        if spec.bc == 'dirichlet':
            v[0] = v[-1] = 0
        u += v * dt
        if spec.bc == 'dirichlet':
            u[0] = u[-1] = 0
        if s % snap_every == 0 or s == steps:
            snaps.append(u.copy())
    return {'x': x, 'snaps': snaps, 'spec': spec}

def solve_poisson_1d(spec):
    """∇²u = f    (direct solve with Dirichlet BCs)."""
    x, dx = FD.grid_1d(spec)
    nx = x.size
    L_op = FD.laplacian_1d(nx, dx, spec.bc)
    f = (np.asarray(spec.source(x), dtype=np.float32)
         if spec.source else np.zeros(nx, dtype=np.float32))
    if spec.bc == 'dirichlet' and nx > 2:
        A_int = L_op[1:-1, 1:-1]
        b_int = f[1:-1]
        u_int = np.linalg.solve(A_int, b_int)
        u = np.zeros(nx, dtype=np.float32)
        u[1:-1] = u_int
    else:
        u = np.linalg.solve(L_op, f)
    return {'x': x, 'u': u, 'spec': spec}

def solve_symbolic_linear_1d(spec):
    """du/dt = c_L L_op u + c_D D u + c_I u   (forward Euler).

    Coefficients come from the symbolic string; named ones (e.g. 'D')
    are looked up in spec.params with a default of 0.
    """
    x, dx = FD.grid_1d(spec)
    nx = x.size
    coeffs = {}
    for op, c in spec.symbolic.items():
        if isinstance(c, str):
            v = spec.params.get(c, 0.0)
            coeffs[op] = float(v) if not isinstance(v, str) else 0.0
        else:
            coeffs[op] = float(c)

    L_op = FD.laplacian_1d(nx, dx, spec.bc)
    D_op = FD.grad_1d(nx, dx, spec.bc)
    O = (coeffs.get('d2u/dx2', 0.0) * L_op
         + coeffs.get('du/dx',     0.0) * D_op
         + coeffs.get('u',         0.0) * np.eye(nx, dtype=np.float32))

    dt = _step_dt(spec); T = _step_tend(spec)
    A = np.eye(nx, dtype=np.float32) + dt * O
    if spec.bc == 'dirichlet':
        _bc_reset(A, FD.bdry_mask_1d(nx, 'dirichlet'))

    u = np.asarray(spec.ic(x), dtype=np.float32)
    if spec.bc == 'dirichlet':
        u[0] = u[-1] = 0

    steps = max(1, int(round(T / dt)))
    snap_every = max(1, steps // 12)
    snaps = [u.copy()]
    for s in range(1, steps + 1):
        u = cache.matmul(A, u.reshape(nx, 1)).flatten()
        if s % snap_every == 0 or s == steps:
            snaps.append(u.copy())
    return {'x': x, 'snaps': snaps, 'spec': spec, 'O': O, 'A': A}


# ================================================================
#  Section 5 — Dispatch: text spec -> solver
# ================================================================
_TABLE = {
    ('heat',             1): solve_heat_1d,
    ('heat',             2): solve_heat_2d,
    ('wave',             1): solve_wave_1d,
    ('poisson',          1): solve_poisson_1d,
    ('symbolic-linear',  1): solve_symbolic_linear_1d,
}

def solve(text: str):
    """Parse the text spec, run the matching solver, return its result."""
    spec = PDEParser.parse(text)
    key = (spec.name, spec.dim)
    if key not in _TABLE:
        raise ValueError(f"no solver registered for {key}")
    return _TABLE[key](spec)


# ================================================================
#  Section 6 — ASCII visualisation (no plotting library needed)
# ================================================================
def ascii_plot_1d(x, u, width=72, height=15, label=''):
    u = np.asarray(u).flatten()
    x = np.asarray(x).flatten()
    if u.size != x.size or x.size == 0:
        return f"[bad data: x={x.size}, u={u.size}]"
    umin, umax = float(u.min()), float(u.max())
    if umax == umin:
        umax = umin + 1
    W, H = width, height
    heights = np.zeros(W, dtype=int)
    cols = (np.round(np.linspace(0, x.size - 1, W)).astype(int)
            if W > 1 else np.array([0]))
    for c, i in enumerate(cols):
        v = float(u[i])
        r = round((v - umin) / (umax - umin) * (H - 1))
        heights[c] = max(0, min(H - 1, r))

    grid = [['.'] * W for _ in range(H)]
    for c, h_val in enumerate(heights):
        for r in range(H - 1 - h_val, H - 1 + 1):
            grid[r][c] = '#'

    if umin < 0 < umax:
        zr = H - 1 - round((0 - umin) / (umax - umin) * (H - 1))
        zr = max(0, min(H - 1, zr))
        for c in range(W):
            if grid[zr][c] == '.':
                grid[zr][c] = '·'

    lines = []
    if label:
        lines.append(label)
    for r in range(H):
        lines.append('|' + ''.join(grid[r]))
    lines.append('+' + '-' * W)
    lines.append(f"  range: [{umin:.4g}, {umax:.4g}]")
    return '\n'.join(lines)

def ascii_plot_2d(snap, nx, ny, width=60, height=22):
    arr = np.asarray(snap, dtype=np.float32).reshape(ny, nx)
    umin, umax = float(arr.min()), float(arr.max())
    if umax == umin:
        umax = umin + 1
    chars = ' .:-=+*#%@'
    nc = len(chars)
    rows = (np.round(np.linspace(ny - 1, 0, height)).astype(int)
            if height > 1 else np.array([0]))
    cols = (np.round(np.linspace(0, nx - 1, width)).astype(int)
            if width > 1 else np.array([0]))
    out = [f"  range: [{umin:.4g}, {umax:.4g}]"]
    for r in rows:
        line = ''
        for c in cols:
            v = (arr[r, c] - umin) / (umax - umin)
            ch = chars[min(nc - 1, int(v * nc))]
            line += ch * 2
        out.append(line)
    return '\n'.join(out)


# ================================================================
#  Section 7 — Command-line interface
# ================================================================
BANNER = (
    "============================================================\n"
    " VGPU PDE Solver — text-format PDEs, VGPU-cached matmul\n"
    "============================================================"
)

EXAMPLES = """
Examples (quote the whole PDE spec as one shell argument):

  heat 1D alpha=0.1 L=1 T=0.5 nx=100 dt=0.001 ic=sin(pi*x) bc=dirichlet
  heat 2D alpha=0.1 L=1 H=1 T=0.1 nx=40 ny=40 dt=0.0001 ic=sin(pi*x)*sin(pi*y)
  wave 1D c=1 L=2 T=2 nx=200 ic=sin(pi*x) dt=0.005
  poisson 1D L=1 nx=100 source=sin(pi*x) bc=dirichlet
  du/dt = 0.1*d2u/dx2 - 0.5*du/dx + 0.2*u
  du/dt = D*d2u/dx2  D=0.05
"""

def main(argv=None):
    if argv is None:
        argv = sys.argv[1:]
    print(BANNER)
    if not argv or argv[0] in ('-h', '--help', 'help'):
        print(EXAMPLES)
        return
    text = ' '.join(argv)
    print(f"\nPDE: {text}\n")

    result   = solve(text)
    spec     = result['spec']
    n_snap   = len(result.get('snaps', []))
    T_end    = float(spec.params.get('T', 1.0))

    print(cache.report())

    if 'snaps' in result:
        if 'X' in result:                                  # 2D heat
            nx, ny = result['x'].size, result['y'].size
            for i, u in enumerate(result['snaps']):
                t = T_end * i / max(1, n_snap - 1)
                print(f"\n--- frame {i+1}/{n_snap}  t = {t:.4f} ---")
                print(ascii_plot_2d(u, nx, ny))
        else:                                              # 1D evolution
            for i, u in enumerate(result['snaps']):
                t = T_end * i / max(1, n_snap - 1)
                print(f"\n--- frame {i+1}/{n_snap}  t = {t:.4f} ---")
                print(ascii_plot_1d(
                    result['x'], u, width=72, height=12,
                    label=f"u(x, t={t:.3f})"))
    elif 'u' in result:                                    # steady solve
        print("\nSteady solution u(x):")
        print(ascii_plot_1d(result['x'], result['u'],
                            width=72, height=18))


if __name__ == '__main__':
    main()