###############################################################################
# TRAINING: PILE dataset (streaming) with wave-scalar computing
###############################################################################
import optax

def load_pile_streaming(batch_size, seq_len, seed=42):
    from datasets import load_dataset
    ds = load_dataset("monology/pile-uncopyrighted", streaming=True, split="train")
    def tokenize(example):
        text = example.get("text", "")
        if not text: text = ""
        tokens = list(text.encode('utf-8'))[:seq_len]
        tokens = tokens + [0] * (seq_len - len(tokens))
        return {"input_ids": tokens}
    ds = ds.map(tokenize)
    ds = ds.shuffle(seed=seed, buffer_size=1000)
    return ds

def create_batches(ds, batch_size, seq_len, d_model, n_batches=50):
    batches = []
    it = iter(ds)
    for _ in range(n_batches):
        batch = []
        for _ in range(batch_size):
            try:
                example = next(it)
                ids = example.get("input_ids", [0]*seq_len)
                if len(ids) < seq_len: ids = ids + [0]*(seq_len-len(ids))
                batch.append(ids[:seq_len])
            except StopIteration: batch.append([0]*seq_len)
        ids_arr = np.array(batch, dtype=np.float32)
        x = ids_arr / 256.0
        x = np.tile(x[..., None], (1,1,d_model//seq_len+1))[:,:,:d_model]
        x = x * (0.5 + 0.5*np.random.randn(batch_size,seq_len,d_model)*0.1)
        batches.append(jnp.array(x.astype(np.float32)))
    return batches

def train_model(model, d_model, n_steps=30, lr=1e-4):
    """Train CogniThread model on PILE data with wave-scalar computing."""
    print(f"  Training on PILE (streaming) for {n_steps} steps...")
    try:
        ds = load_pile_streaming(batch_size=2, seq_len=8)
        batches = create_batches(ds, batch_size=2, seq_len=8, d_model=d_model, n_batches=n_steps)
        print(f"  Loaded {len(batches)} batches from PILE")
    except Exception as e:
        print(f"  PILE loading failed ({e}), using synthetic data")
        key = jax.random.key(0)
        batches = [jax.random.normal(jax.random.fold_in(key, i), (2, 8, d_model))
                   for i in range(n_steps)]

    # Use gradient clipping to prevent NaN from exploding SSM dynamics
    tx = optax.chain(
        optax.clip_by_global_norm(1.0),
        optax.adam(lr),
    )
    opt = nn.ModelAndOptimizer(model, tx, wrt=nn.Param)

    @nn.jit
    def train_step(opt, x, step_key):
        def loss_fn(model):
            y_final, var_y, w_f, mu = model(x, key=step_key)
            # Use L2 loss on output (autoencoder-style)
            # Replace inf/nan with large but finite value
            y_sq = jnp.sum(y_final ** 2)
            y_sq = jnp.where(jnp.isfinite(y_sq), y_sq, jnp.array(1e6))
            var_loss = jnp.mean(var_y) * 0.01
            ent_loss = -jnp.mean(mu * jnp.log(mu + 1e-8)) * 0.001
            return jnp.mean(y_sq) + var_loss + ent_loss
        loss, grads = nn.value_and_grad(loss_fn)(opt.model)
        # Replace NaN gradients with zeros
        grads = jax.tree.map(lambda g: jnp.where(jnp.isfinite(g), g, jnp.zeros_like(g)), grads)
        opt.update(grads)
        return jnp.where(jnp.isfinite(loss), loss, jnp.array(1e6))

    losses = []
    for step, x_batch in enumerate(batches):
        step_key = jax.random.key(step)
        loss = train_step(opt, x_batch, step_key)
        losses.append(float(loss))
        if step % 10 == 0 or step == n_steps - 1:
            print(f"    Step {step+1}/{n_steps}: loss={float(loss):.4f}")

    print(f"  Training complete. Final loss: {losses[-1]:.4f}")
    return opt.model, losses