import os, sys, math, struct
os.environ['JAX_PLATFORMS'] = 'cpu'
import jax
import jax.numpy as jnp
import flax.nnx as nn
import numpy as np
from typing import Tuple

print(f"=== CogniThread SSM Architecture === JAX {jax.__version__}")

###############################################################################
# 1. WAVESCALAR — Thermal Analog Wave Representation
###############################################################################
class WaveScalar(nn.Module):
    """
    Encodes each scalar as a radionics scalar wave with 5 parameters:
    amplitude, phase_shift, frequency, period (1/freq), fractal_dim.
    v(t) = A * sin(2*pi*f*t + phi) * fractal_mod(t, D)
    fractal_mod = 1 + sum_{n=1}^{N} (1/n^D) * sin(2*pi*n*f*t + n*phi)
    """
    def __init__(self, shape: Tuple[int, ...], *, rngs: nn.Rngs,
                 init_amp_scale: float = 1.0,
                 freq_range: Tuple[float, float] = (0.1, 10.0),
                 fractal_range: Tuple[float, float] = (1.0, 2.0)):
        super().__init__()
        self.shape = shape
        key = rngs.params()
        fan_in = max(math.prod(shape), 1)
        std = init_amp_scale / math.sqrt(fan_in)
        self.amplitude = nn.Param(jax.random.normal(key, shape) * std)
        key, sk = jax.random.split(key)
        self.phase_shift = nn.Param(jax.random.uniform(sk, shape, minval=0, maxval=2*jnp.pi))
        key, sk = jax.random.split(key)
        lf = jax.random.uniform(sk, shape, minval=jnp.log(freq_range[0]), maxval=jnp.log(freq_range[1]))
        self.frequency = nn.Param(jnp.exp(lf))
        key, sk = jax.random.split(key)
        self.fractal_dim = nn.Param(jax.random.uniform(sk, shape, minval=fractal_range[0], maxval=fractal_range[1]))

    @property
    def period(self):
        return 1.0 / self.frequency.value

    def fractal_mod(self, t, D, n_harm=8):
        harms = jnp.arange(1, n_harm+1, dtype=jnp.float32)
        w = 1.0 / jnp.power(harms, D[..., None])
        f, phi = self.frequency.value, self.phase_shift.value
        sigs = jnp.sin(2*jnp.pi*harms*f[..., None]*t + harms*phi[..., None])
        mod = 1.0 + jnp.sum(w * sigs, axis=-1)
        return mod / jnp.sqrt(1.0 + jnp.sum(w**2, axis=-1))

    def evaluate(self, t=None):
        if t is None:
            t = jnp.array(0.0)
        A = self.amplitude.value
        f = self.frequency.value
        phi = self.phase_shift.value
        D = self.fractal_dim.value
        base = A * jnp.sin(2*jnp.pi*f*t + phi)
        return base * self.fractal_mod(t, D)

    def decode(self):
        return self.evaluate(t=jnp.array(0.0))

    def rms_value(self):
        A, D = self.amplitude.value, self.fractal_dim.value
        harms = jnp.arange(1, 9, dtype=jnp.float32)
        w = 1.0 / jnp.power(harms, D[..., None])
        rf = jnp.sqrt(1.0 + jnp.sum(w**2, axis=-1)/2.0)
        nm = jnp.sqrt(1.0 + jnp.sum(w**2, axis=-1))
        return A / jnp.sqrt(2.0) * rf / nm


class WaveScalarLinear(nn.Module):
    def __init__(self, in_f: int, out_f: int, *, rngs: nn.Rngs):
        super().__init__()
        self.weight = WaveScalar((out_f, in_f), rngs=rngs)
        self.bias = WaveScalar((out_f,), rngs=rngs, init_amp_scale=0.01)

    def __call__(self, x):
        return x @ self.weight.decode().T + self.bias.decode()

print("[1/6] WaveScalar defined")