"""MLX dual-timescale matrix state with denoise and commit lifetimes."""

from __future__ import annotations

from dataclasses import dataclass

import mlx.core as mx
from mlx import nn

from .mlx_gdn2_memory import GDN2Memory


@dataclass(frozen=True)
class GDN2TrajectoryState:
    cells: mx.array
    row: mx.array
    persistent: mx.array
    seen: mx.array


    def clear_refill(self, lengths: mx.array, head: mx.array) -> "GDN2TrajectoryState":
        batch, canvas = self.seen.shape
        if lengths.shape != (batch,) or head.shape != (batch,):
            raise ValueError("Commit lengths and canvas heads must be per row.")
        physical = mx.arange(canvas)[None, :]
        overwritten = ((physical - head[:, None]) % canvas) < lengths[:, None]
        return GDN2TrajectoryState(
            mx.where(overwritten[..., None, None, None], 0.0, self.cells),
            self.row, self.persistent, self.seen & ~overwritten,
        )


class GDN2TrajectoryMemory(nn.Module):
    def __init__(self, width: int, *, probes: int = 4,
                 working_heads: int = 16, working_key: int = 64,
                 working_value: int = 64, persistent_heads: int = 16,
                 persistent_key: int = 128, persistent_value: int = 128,
                 persistent_observation_dim: int | None = None) -> None:
        super().__init__()
        if probes <= 0:
            raise ValueError("Probe count must be positive.")
        self.probes = probes
        self.cell = GDN2Memory(width, working_heads, working_key, working_value)
        self.row = GDN2Memory(width, working_heads, working_key, working_value)
        self.persistent = GDN2Memory(width, persistent_heads, persistent_key, persistent_value,
                                    observation_dim=persistent_observation_dim)
        self.probe_embed = nn.Embedding(probes, width)
        self.probe_embed.weight = mx.random.normal((probes, width)) * 0.02

    def empty(self, batch: int, canvas: int) -> GDN2TrajectoryState:
        return GDN2TrajectoryState(
            self.cell.empty(batch, canvas), self.row.empty(batch),
            self.persistent.empty(batch), mx.zeros((batch, canvas), mx.bool_),
        )

    def read(self, state: GDN2TrajectoryState, query: mx.array) -> mx.array:
        result = (self.cell.read(state.cells, query)
                  + self.row.read_shared(state.row, query)
                  + self.persistent.read_shared(state.persistent, query))
        return mx.where(state.seen[..., None], result, 0.0)

    def observe(self, state: GDN2TrajectoryState, observation: mx.array,
                live: mx.array, head: mx.array) -> GDN2TrajectoryState:
        batch, canvas, width = observation.shape
        if live.shape != (batch, canvas) or head.shape != (batch,):
            raise ValueError("Observation mask and head have incorrect shapes.")
        cells = self.cell.transition(state.cells, observation, live)
        seen = state.seen | live
        logical_idx = (head[:, None] + mx.arange(canvas)[None, :]) % canvas
        logical = mx.take_along_axis(observation,
                                     mx.broadcast_to(logical_idx[..., None], observation.shape), axis=1)
        logical_live = mx.take_along_axis(live, mx.stop_gradient(logical_idx), axis=1)
        row_state = state.row
        for probe in range(self.probes):
            lo = canvas * probe // self.probes
            hi = canvas * (probe + 1) // self.probes
            selected = logical_live[:, lo:hi]
            count = mx.sum(selected.astype(mx.float32), axis=1, keepdims=True)
            pooled = mx.sum(logical[:, lo:hi].astype(mx.float32)
                            * selected[..., None], axis=1)
            pooled = (pooled / mx.maximum(count, 1.0)).astype(observation.dtype)
            pooled = pooled + self.probe_embed.weight[probe].astype(observation.dtype)
            row_state = self.row.transition(row_state, pooled, count[:, 0] > 0)
        return GDN2TrajectoryState(cells, row_state, state.persistent, seen)


__all__ = ["GDN2TrajectoryMemory", "GDN2TrajectoryState"]