"""
evaluate() — speed stack (fastest to slowest):
  1. numexpr  — Intel VML + cache-chunked eval, 2-8x faster than Numba for trig/log/exp
  2. C++ JIT  — hand-optimized for unsupported functions
  3. Numba    — emergency fallback (rarely needed)
"""

import re
import ctypes
import math
import numpy as np
import numba as nb
import numexpr as ne
from functools import lru_cache

# ──────────────────────────────────────────────
# C++ backend
# ──────────────────────────────────────────────
try:
    _lib = ctypes.CDLL('./engine.so')
    _lib.batch_evaluate_expr.argtypes = [
        ctypes.c_char_p,
        ctypes.POINTER(ctypes.c_double),
        ctypes.POINTER(ctypes.c_double),
        ctypes.c_int,
    ]
    _lib.bulk_run_expr.argtypes = [
        ctypes.c_char_p,
        ctypes.POINTER(ctypes.c_double),
        ctypes.POINTER(ctypes.c_double),
        ctypes.c_int,
        ctypes.c_int,
    ]
    _CPP_AVAILABLE = True
except OSError:
    _CPP_AVAILABLE = False



def _cpp_evaluate(expr_cpp: str, x_arr: np.ndarray) -> np.ndarray:
    n = len(x_arr)
    res = np.empty(n, dtype=np.float64)
    _lib.batch_evaluate_expr(
        expr_cpp.encode(),
        x_arr.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
        res.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
        ctypes.c_int(n),
    )
    return res


# ──────────────────────────────────────────────
# Token parser for ^ and **
# ──────────────────────────────────────────────
def _find_token_start(s: str, end_pos: int) -> int:
    """Walk backwards from end_pos to find start of last complete token."""
    i = end_pos - 1
    while i >= 0 and s[i] == ' ':
        i -= 1
    if i >= 0 and s[i] == ')':
        depth = 1
        i -= 1
        while i >= 0 and depth > 0:
            depth += (s[i] == ')') - (s[i] == '(')
            i -= 1
        # Grab function name before (
        while i >= 0 and (s[i].isalnum() or s[i] == '_'):
            i -= 1
        return i + 1
    while i >= 0 and (s[i].isalnum() or s[i] in '._'):
        i -= 1
    return i + 1


def _find_token_end(s: str, start_pos: int) -> int:
    """Walk forwards from start_pos to find end of next complete token."""
    i = start_pos
    while i < len(s) and s[i] == ' ':
        i += 1
    if i < len(s) and s[i] == '(':
        depth = 1
        i += 1
        while i < len(s) and depth > 0:
            depth += (s[i] == '(') - (s[i] == ')')
            i += 1
        return i
    while i < len(s) and (s[i].isalnum() or s[i] in '._'):
        i += 1
    if i < len(s) and s[i] == '(':
        depth = 1
        i += 1
        while i < len(s) and depth > 0:
            depth += (s[i] == '(') - (s[i] == ')')
            i += 1
    return i


def _replace_pow(expr: str, target: str) -> str:
    """Replace ^ and ** operators. target='**' for Python, 'pow' for C++."""
    e = expr
    for op in ('**', '^'):
        result = []
        i = 0
        while i < len(e):
            is_op = (e[i:i+2] == '**') if op == '**' else (e[i] == '^')
            if is_op:
                step = 2 if op == '**' else 1
                result_str = ''.join(result)
                tok_start  = _find_token_start(result_str, len(result_str))
                base       = result_str[tok_start:]
                result     = list(result_str[:tok_start])
                i += step
                while i < len(e) and e[i] == ' ':
                    i += 1
                exp_end = _find_token_end(e, i)
                exp     = e[i:exp_end]
                i       = exp_end
                if target == 'pow':
                    result.append(f'pow({base}, {exp})')
                else:
                    result.append(f'{base}**{exp}')
            else:
                result.append(e[i])
                i += 1
        e = ''.join(result)
    return e


# ──────────────────────────────────────────────
# Expression translators
# ──────────────────────────────────────────────

# numexpr doesn't support these — fall through to C++
_NE_UNSUPPORTED = re.compile(r'\b(atan2|ceil|floor|sinh|cosh|tanh)\b')

_NE_MAP = {"ln": "log", "asin": "arcsin", "acos": "arccos", "atan": "arctan", "fabs": "abs"}
_NE_PATTERN = re.compile(r'\b(' + '|'.join(sorted(_NE_MAP, key=len, reverse=True)) + r')\b')

def _to_numexpr(expr: str) -> str:
    e = _replace_pow(expr.strip(), target='**')
    return _NE_PATTERN.sub(lambda m: _NE_MAP[m.group(0)], e)


_CPP_MAP = {
    "log10": "log10", "log2": "log2", "log": "log", "ln": "log",
    "exp": "exp",     "sqrt": "sqrt", "abs": "fabs", "fabs": "fabs",
    "sin": "sin",     "cos": "cos",   "tan": "tan",
    "asin": "asin",   "acos": "acos", "atan": "atan", "atan2": "atan2",
    "sinh": "sinh",   "cosh": "cosh", "tanh": "tanh",
    "ceil": "ceil",   "floor": "floor",
}
_CPP_PATTERN = re.compile(r'\b(' + '|'.join(sorted(_CPP_MAP, key=len, reverse=True)) + r')\b')

def _to_cpp_expr(expr: str) -> str:
    e = _replace_pow(expr.strip(), target='pow')
    e = re.sub(r'\bx\b', 'val', e)
    return _CPP_PATTERN.sub(lambda m: _CPP_MAP[m.group(0)], e)


# Numba (emergency fallback only)
_NB_MAP = {
    "log10": "math.log10", "log2": "math.log2", "log": "math.log", "ln": "math.log",
    "exp": "math.exp",     "sqrt": "math.sqrt",  "abs": "math.fabs", "fabs": "math.fabs",
    "sin": "math.sin",     "cos": "math.cos",    "tan": "math.tan",
    "asin": "math.asin",   "acos": "math.acos",  "atan": "math.atan",  "atan2": "math.atan2",
    "sinh": "math.sinh",   "cosh": "math.cosh",  "tanh": "math.tanh",
    "ceil": "math.ceil",   "floor": "math.floor","pow": "math.pow",
    "pi": "math.pi",       "e": "math.e",
}
_NB_PATTERN = re.compile(r'\b(' + '|'.join(sorted(_NB_MAP, key=len, reverse=True)) + r')\b')

def _to_numba_expr(expr: str) -> str:
    e = _replace_pow(expr.strip(), target='**')
    return _NB_PATTERN.sub(lambda m: _NB_MAP[m.group(0)], e)


@lru_cache(maxsize=64)
def _build_numba_kernel(expr_key: str):
    nb_expr = _to_numba_expr(expr_key)
    src = f"""
import math
import numba as _nb
import numpy as _np

@_nb.njit(parallel=True, fastmath=True, cache=True, boundscheck=False)
def _kernel(x: _np.ndarray, out: _np.ndarray) -> None:
    for i in _nb.prange(x.shape[0]):
        out[i] = {nb_expr}
"""
    ns: dict = {}
    exec(compile(src, "<jit_kernel>", "exec"), ns)  # noqa: S102
    return ns["_kernel"]


def _canonicalize(expr: str) -> str:
    return re.sub(r'\s+', ' ', expr.strip().lower())


# ──────────────────────────────────────────────
# Public API
# ──────────────────────────────────────────────
def evaluate(expr: str, x_input) -> np.ndarray:
    """
    Evaluate a math expression over an array of x values.

    Speed hierarchy:
      1. numexpr  — Intel VML, cache-optimized, 2-8x faster than Numba for trig/exp/log
      2. C++ JIT  — for unsupported functions (sinh/cosh/tanh/ceil/floor/atan2)
      3. Numba    — emergency fallback

    Supports: +  -  *  /  **  ^  sin  cos  tan  asin  acos  atan  atan2
              sinh cosh tanh  log  log2  log10  ln  exp  sqrt  abs  ceil  floor
    """
    x_arr = np.ascontiguousarray(x_input, dtype=np.float64)
    canon = _canonicalize(expr)

    # ── Tier 1: numexpr ──────────────────────────────────────────────────
    if not _NE_UNSUPPORTED.search(canon):
        try:
            return ne.evaluate(
                _to_numexpr(canon),
                local_dict={'x': x_arr},
                optimization='aggressive',
            )
        except Exception:
            pass

    # ── Tier 2: C++ JIT ──────────────────────────────────────────────────
    if _CPP_AVAILABLE:
        try:
            return _cpp_evaluate(_to_cpp_expr(expr), x_arr)
        except Exception:
            pass

    # ── Tier 3: Numba (fallback) ─────────────────────────────────────────
    try:
        out = np.empty(x_arr.shape[0], dtype=np.float64)
        _build_numba_kernel(canon)(x_arr, out)
        return out
    except Exception as e:
        raise RuntimeError(f"All backends failed: {e}") from e
    
def bulk(expr: str, x_input, iterations: int = 1) -> np.ndarray:
    x_arr = np.ascontiguousarray(x_input, dtype=np.float64)
    res = np.empty(x_arr.shape[0], dtype=np.float64)

    if _CPP_AVAILABLE:
        try:
            _lib.bulk_run_expr(
                _to_cpp_expr(expr).encode(),
                x_arr.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
                res.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
                ctypes.c_int(x_arr.shape[0]),
                ctypes.c_int(iterations),
            )
            return res
        except Exception:
            pass

    for _ in range(iterations):
        res = evaluate(expr, x_arr)
    return res