"""
AI Voxel Renderer
==================
Renders a 2D slice of the AI Voxel on GPU (OpenGL 4.3+ compute shader)
or CPU (NumPy fallback). Visualizes density, flux, and edge sharpness.

Usage:
    python ai_voxel_render.py              # 256×256, slice dims 0,1
    python ai_voxel_render.py 512 0 2      # 512×512, slice dims 0,2

Requirements:
    numpy, matplotlib
    moderngl  (optional — for GPU rendering:  pip install moderngl)
"""

import os
import sys
import time
import numpy as np


# ════════════════════════════════════════════════════════════════════════
# GLSL COMPUTE SHADER  (Stage 6 — matches the AI Voxel framework)
# ════════════════════════════════════════════════════════════════════════
GLSL_SHADER = r"""#version 430 core

layout(local_size_x = 64) in;

uniform int   u_K;
uniform int   u_D;
uniform float u_T;
uniform int   u_steps;
uniform int   u_Dhidden;

layout(std430, binding = 0) readonly buffer FieldOmega  { float omega[];    };
layout(std430, binding = 1) readonly buffer FieldA      { float aCoeff[];   };
layout(std430, binding = 2) readonly buffer FieldB      { float bCoeff[];   };
layout(std430, binding = 3) readonly buffer BoundaryW1  { float bW1[];      };
layout(std430, binding = 4) readonly buffer BoundaryB1  { float bB1[];      };
layout(std430, binding = 5) readonly buffer BoundaryW2  { float bW2[];      };
layout(std430, binding = 6) readonly buffer BoundaryB2  { float bB2[];      };
layout(std430, binding = 7) readonly buffer QueryPoints { float queryPts[]; };
layout(std430, binding = 8) writeonly buffer Output     { float outData[];  };

const int MAX_D = 128;
const int MAX_K = 256;

// ── Vector field: F(h) = Σ_k a_k sin(ω_k·h) + b_k cos(ω_k·h) ──
void evaluateField(in float h[MAX_D], out float F[MAX_D]) {
    for (int d = 0; d < u_D; d++) F[d] = 0.0;
    for (int k = 0; k < u_K; k++) {
        float projection = 0.0;
        for (int d = 0; d < u_D; d++)
            projection += omega[k * u_D + d] * h[d];
        float sinP = sin(projection);
        float cosP = cos(projection);
        for (int d = 0; d < u_D; d++)
            F[d] += aCoeff[k * u_D + d] * sinP + bCoeff[k * u_D + d] * cosP;
    }
}

// ── Boundary: S(p) = W2·tanh(W1·p + b1) + b2 ──
float evaluateBoundary(in float h[MAX_D]) {
    float hidden[MAX_D];
    for (int j = 0; j < u_Dhidden; j++) {
        float sum = bB1[j];
        for (int d = 0; d < u_D; d++)
            sum += bW1[j * u_D + d] * h[d];
        hidden[j] = tanh(sum);
    }
    float S = bB2[0];
    for (int j = 0; j < u_Dhidden; j++)
        S += bW2[j] * hidden[j];
    return S;
}

// ── Numerical gradient of S via central differences ──
void gradientBoundary(in float h[MAX_D], out float gradS[MAX_D]) {
    float eps = 0.001;
    float hPlus[MAX_D], hMinus[MAX_D];
    for (int d = 0; d < u_D; d++) {
        for (int i = 0; i < u_D; i++) {
            hPlus[i] = h[i];
            hMinus[i] = h[i];
        }
        hPlus[d] += eps;
        hMinus[d] -= eps;
        gradS[d] = (evaluateBoundary(hPlus) - evaluateBoundary(hMinus)) / (2.0 * eps);
    }
}

// ── Flux: 𝔽(p) = ∇S · F(p)  (THE MISSING INFORMATION) ──
float computeFlux(in float h[MAX_D]) {
    float gradS[MAX_D], F[MAX_D];
    gradientBoundary(h, gradS);
    evaluateField(h, F);
    float flux = 0.0;
    for (int d = 0; d < u_D; d++)
        flux += gradS[d] * F[d];
    return flux;
}

// ── RK4 integration of dh/dt = F(h) ──
void integrateRK4(inout float h[MAX_D], float dt, int numSteps) {
    float k1[MAX_D], k2[MAX_D], k3[MAX_D], k4[MAX_D], hTemp[MAX_D];
    for (int step = 0; step < numSteps; step++) {
        evaluateField(h, k1);
        for (int d = 0; d < u_D; d++) hTemp[d] = h[d] + 0.5 * dt * k1[d];
        evaluateField(hTemp, k2);
        for (int d = 0; d < u_D; d++) hTemp[d] = h[d] + 0.5 * dt * k2[d];
        evaluateField(hTemp, k3);
        for (int d = 0; d < u_D; d++) hTemp[d] = h[d] + dt * k3[d];
        evaluateField(hTemp, k4);
        for (int d = 0; d < u_D; d++)
            h[d] += dt / 6.0 * (k1[d] + 2.0*k2[d] + 2.0*k3[d] + k4[d]);
    }
}

void main() {
    uint tid = gl_GlobalInvocationID.x;

    // Load query point from buffer
    float h[MAX_D];
    for (int d = 0; d < u_D; d++)
        h[d] = queryPts[tid * u_D + d];

    // Integrate ODE
    float dt = u_T / float(u_steps);
    integrateRK4(h, dt, u_steps);

    // Evaluate boundary and flux
    float S = evaluateBoundary(h);
    float flux = computeFlux(h);

    // Outputs: density, r, g, b, flux, edgeSharpness
    float density = 1.0 / (1.0 + exp(-S * 10.0));
    float r = 0.5 + 0.5 * h[0];
    float g = 0.5 + 0.5 * (u_D > 1 ? h[1] : 0.0);
    float b = 0.5 + 0.5 * (u_D > 2 ? h[2] : 0.0);
    float edgeSharpness = abs(flux) / (1.0 + abs(S));

    uint outIdx = tid * 6;
    outData[outIdx + 0] = density;
    outData[outIdx + 1] = r;
    outData[outIdx + 2] = g;
    outData[outIdx + 3] = b;
    outData[outIdx + 4] = flux;
    outData[outIdx + 5] = edgeSharpness;
}
"""


# ════════════════════════════════════════════════════════════════════════
# COEFFICIENT LOADING
# ════════════════════════════════════════════════════════════════════════
def load_coefficients(path="ai_voxel_coeffs.npz"):
    """Load AI Voxel coefficients saved by ai_voxel_experiment.py."""
    c = np.load(path)
    Dhidden = int(c['Dhidden']) if 'Dhidden' in c else int(c['bW1'].shape[0])
    bB2 = (c['bB2'].astype(np.float32) if 'bB2' in c
           else np.zeros(1, dtype=np.float32))
    return {
        'omega':   c['omega'].astype(np.float32),   # [K, D]
        'a':       c['a'].astype(np.float32),        # [K, D]
        'b':       c['b'].astype(np.float32),        # [K, D]
        'bW1':     c['bW1'].astype(np.float32),      # [DHIDDEN, D]
        'bB1':     c['bB1'].astype(np.float32),      # [DHIDDEN]
        'bW2':     c['bW2'].astype(np.float32),      # [DHIDDEN]
        'bB2':     bB2,                               # [1]
        'D':       int(c['D']),
        'K':       int(c['K']),
        'Dhidden': Dhidden,
        'steps':   int(c['steps']),
        'T':       float(c['T']),
    }


# ════════════════════════════════════════════════════════════════════════
# GPU RENDERING  (OpenGL 4.3+ compute shader via moderngl)
# ════════════════════════════════════════════════════════════════════════
def render_gpu(coeffs, query_points):
    """Render on GPU using OpenGL compute shader. Returns (density, colors, flux, edge)."""
    import moderngl

    D = coeffs['D']
    K = coeffs['K']
    N = query_points.shape[0]

    ctx = moderngl.create_standalone_context(require=430)
    shader = ctx.compute_shader(GLSL_SHADER)

    # Pad bB2 to at least 16 bytes (some drivers dislike tiny SSBOs)
    bB2_padded = np.zeros(4, dtype=np.float32)
    bB2_padded[0] = coeffs['bB2'][0]

    # Create + bind SSBOs
    bufs = [
        ctx.buffer(coeffs['omega'].tobytes()),                       # 0
        ctx.buffer(coeffs['a'].tobytes()),                           # 1
        ctx.buffer(coeffs['b'].tobytes()),                           # 2
        ctx.buffer(coeffs['bW1'].tobytes()),                         # 3
        ctx.buffer(coeffs['bB1'].tobytes()),                         # 4
        ctx.buffer(coeffs['bW2'].tobytes()),                         # 5
        ctx.buffer(bB2_padded.tobytes()),                            # 6
        ctx.buffer(query_points.tobytes()),                          # 7
        ctx.buffer(np.zeros(N * 6, dtype=np.float32).tobytes()),     # 8
    ]
    for i, buf in enumerate(bufs):
        buf.bind_to_storage_buffer(i)

    # Uniforms
    shader['u_K'].value = K
    shader['u_D'].value = D
    shader['u_T'].value = coeffs['T']
    shader['u_steps'].value = coeffs['steps']
    shader['u_Dhidden'].value = coeffs['Dhidden']

    # Dispatch
    num_groups = (N + 63) // 64
    shader.run(group_x=num_groups)

    # Read back
    raw = np.frombuffer(bufs[8].read(), dtype=np.float32).reshape(-1, 6)
    return raw[:, 0], raw[:, 1:4], raw[:, 4], raw[:, 5]


# ════════════════════════════════════════════════════════════════════════
# CPU RENDERING  (NumPy fallback — same math, vectorised)
# ════════════════════════════════════════════════════════════════════════
def _field_np(h, omega, a, b):
    """F(h) = sin(Ω·h)·a + cos(Ω·h)·b   h:[N,D] → [N,D]"""
    proj = h @ omega.T          # [N, K]
    return np.sin(proj) @ a + np.cos(proj) @ b

def _boundary_np(h, bW1, bB1, bW2, bB2):
    """S(p) = W2·tanh(W1·p + b1) + b2   h:[N,D] → [N]"""
    hidden = np.tanh(h @ bW1.T + bB1)   # [N, DHIDDEN]
    return hidden @ bW2 + bB2[0]        # [N]

def _grad_boundary_np(h, bW1, bB1, bW2, bB2, eps=0.001):
    """∇S via central differences, vectorised over all D dims."""
    N, D = h.shape
    h_plus  = np.broadcast_to(h[:, None, :], (N, D, D)).copy()
    h_minus = np.broadcast_to(h[:, None, :], (N, D, D)).copy()
    idx = np.arange(D)
    h_plus[:, idx, idx]  += eps
    h_minus[:, idx, idx] -= eps
    s_plus  = _boundary_np(h_plus.reshape(N * D, D),  bW1, bB1, bW2, bB2).reshape(N, D)
    s_minus = _boundary_np(h_minus.reshape(N * D, D), bW1, bB1, bW2, bB2).reshape(N, D)
    return (s_plus - s_minus) / (2.0 * eps)

def render_cpu(coeffs, query_points):
    """Render on CPU using NumPy. Returns (density, colors, flux, edge)."""
    omega, a, b = coeffs['omega'], coeffs['a'], coeffs['b']
    bW1, bB1, bW2, bB2 = coeffs['bW1'], coeffs['bB1'], coeffs['bW2'], coeffs['bB2']
    T, steps, D = coeffs['T'], coeffs['steps'], coeffs['D']
    N = query_points.shape[0]

    # RK4 integration
    h = query_points.copy()
    dt = T / steps
    for _ in range(steps):
        k1 = _field_np(h, omega, a, b)
        k2 = _field_np(h + 0.5 * dt * k1, omega, a, b)
        k3 = _field_np(h + 0.5 * dt * k2, omega, a, b)
        k4 = _field_np(h + dt * k3, omega, a, b)
        h = h + (dt / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)

    # Boundary, gradient, flux
    S     = _boundary_np(h, bW1, bB1, bW2, bB2)
    gradS = _grad_boundary_np(h, bW1, bB1, bW2, bB2)
    F     = _field_np(h, omega, a, b)
    flux  = (gradS * F).sum(axis=-1)

    # Outputs (match GLSL shader exactly)
    density = 1.0 / (1.0 + np.exp(-S * 10.0))
    chans = []
    for d in range(3):
        if d < D:
            chans.append(0.5 + 0.5 * h[:, d])
        else:
            chans.append(np.full(N, 0.5, dtype=np.float32))
    colors = np.stack(chans, axis=-1)
    edge = np.abs(flux) / (1.0 + np.abs(S))

    return density, colors, flux, edge


# ════════════════════════════════════════════════════════════════════════
# VISUALIZATION
# ════════════════════════════════════════════════════════════════════════
def visualize(density, colors, flux, edge, grid_size, dim0, dim1,
              backend, ms):
    import matplotlib.pyplot as plt

    density = density.reshape(grid_size, grid_size)
    flux    = flux.reshape(grid_size, grid_size)
    edge    = edge.reshape(grid_size, grid_size)
    colors  = colors.reshape(grid_size, grid_size, 3)

    fig, axes = plt.subplots(2, 2, figsize=(12, 10))
    fig.suptitle(
        f"AI Voxel Render  —  {backend}  —  {ms:.1f} ms  —  "
        f"{grid_size}×{grid_size} = {grid_size**2:,} queries",
        fontsize=13,
    )
    ext = [-1, 1, -1, 1]

    im0 = axes[0, 0].imshow(density, cmap='viridis', origin='lower', extent=ext)
    axes[0, 0].set_title('Density  σ(S·10)')
    axes[0, 0].set_xlabel(f'dim {dim0}'); axes[0, 0].set_ylabel(f'dim {dim1}')
    plt.colorbar(im0, ax=axes[0, 0], fraction=0.046)

    im1 = axes[0, 1].imshow(flux, cmap='RdBu_r', origin='lower', extent=ext)
    axes[0, 1].set_title('Flux  ∇S·F  (missing information)')
    axes[0, 1].set_xlabel(f'dim {dim0}')
    plt.colorbar(im1, ax=axes[0, 1], fraction=0.046)

    im2 = axes[1, 0].imshow(edge, cmap='hot', origin='lower', extent=ext)
    axes[1, 0].set_title('Edge Sharpness  |flux| / (1+|S|)')
    axes[1, 0].set_xlabel(f'dim {dim0}'); axes[1, 0].set_ylabel(f'dim {dim1}')
    plt.colorbar(im2, ax=axes[1, 0], fraction=0.046)

    axes[1, 1].imshow(np.clip(colors, 0, 1), origin='lower', extent=ext)
    axes[1, 1].set_title('Color (RGB from latent)')
    axes[1, 1].set_xlabel(f'dim {dim0}')

    plt.tight_layout()
    out_path = "ai_voxel_render.png"
    plt.savefig(out_path, dpi=150, bbox_inches='tight')
    print(f"  Saved: {out_path}")
    plt.show()


# ════════════════════════════════════════════════════════════════════════
# MAIN
# ════════════════════════════════════════════════════════════════════════
def main():
    # Parse args
    grid_size = 256
    dim0, dim1 = 0, 1
    if len(sys.argv) >= 2:
        grid_size = int(sys.argv[1])
    if len(sys.argv) >= 4:
        dim0 = int(sys.argv[2])
        dim1 = int(sys.argv[3])

    print("=" * 56)
    print("  AI VOXEL RENDERER")
    print("=" * 56)

    # Load coefficients
    coeff_path = "ai_voxel_coeffs.npz"
    if not os.path.exists(coeff_path):
        print(f"\n  Error: {coeff_path} not found.")
        print("  Run ai_voxel_experiment.py first to train + save coefficients.")
        sys.exit(1)

    coeffs = load_coefficients(coeff_path)
    D = coeffs['D']
    print(f"  D={D}  K={coeffs['K']}  Dhidden={coeffs['Dhidden']}")
    print(f"  steps={coeffs['steps']}  T={coeffs['T']}")
    print(f"  Slice: dim {dim0} × dim {dim1}")
    print(f"  Grid:  {grid_size}×{grid_size} = {grid_size**2:,} queries")

    # Generate query points (2D slice, remaining dims = 0)
    x = np.linspace(-1, 1, grid_size, dtype=np.float32)
    y = np.linspace(-1, 1, grid_size, dtype=np.float32)
    xx, yy = np.meshgrid(x, y)
    N = grid_size * grid_size
    query_points = np.zeros((N, D), dtype=np.float32)
    query_points[:, dim0] = xx.flatten()
    query_points[:, dim1] = yy.flatten()

    # Try GPU first, fall back to CPU
    backend = None
    try:
        import moderngl  # noqa: F401
        print("\n  Trying GPU (OpenGL 4.3 compute shader)...")
        t0 = time.perf_counter()
        density, colors, flux, edge = render_gpu(coeffs, query_points)
        ms = (time.perf_counter() - t0) * 1000
        backend = "GPU (OpenGL 4.3)"
        print(f"  GPU render: {ms:.1f} ms  ({N / (ms / 1e3):,.0f} queries/sec)")
    except ImportError:
        print("\n  moderngl not installed — using CPU (NumPy).")
        print("  For GPU rendering:  pip install moderngl")
        backend = "CPU (NumPy)"
        t0 = time.perf_counter()
        density, colors, flux, edge = render_cpu(coeffs, query_points)
        ms = (time.perf_counter() - t0) * 1000
        print(f"  CPU render: {ms:.1f} ms  ({N / (ms / 1e3):,.0f} queries/sec)")
    except Exception as e:
        print(f"\n  GPU failed: {e}")
        print("  Falling back to CPU (NumPy)...")
        backend = "CPU (NumPy, fallback)"
        t0 = time.perf_counter()
        density, colors, flux, edge = render_cpu(coeffs, query_points)
        ms = (time.perf_counter() - t0) * 1000
        print(f"  CPU render: {ms:.1f} ms  ({N / (ms / 1e3):,.0f} queries/sec)")

    print(f"\n  Density: [{density.min():.4f}, {density.max():.4f}]")
    print(f"  Flux:    [{flux.min():.4f}, {flux.max():.4f}]")
    print(f"  Edge:    [{edge.min():.4f}, {edge.max():.4f}]")

    # Visualize
    print("\n  Generating visualization...")
    visualize(density, colors, flux, edge, grid_size, dim0, dim1,
              backend, ms)

    print("\n" + "=" * 56)
    print("  Done.")
    print("=" * 56)


if __name__ == "__main__":
    main()