import numpy as _np_cpu
from .status import Status
from .train import _to_gpu, cd_binary_binary, cd_gaussian_binary, cd_gaussian_gaussian, cd_binary_gaussian
from .stack import Stack, StackType
from .matrix import Mat, np, rms_error_accu, convert
from .entity import Entity
from .layer import Layer
from .entity import EntityParams, TrainingParams

_CD_FUNC = {
    Entity.Type.BB_RBM: cd_binary_binary,
    Entity.Type.GB_RBM: cd_gaussian_binary,
    Entity.Type.GG_RBM: cd_gaussian_gaussian,
    Entity.Type.BG_RBM: cd_binary_gaussian,
}


class StackRnn(Stack):
    """Recurrent RBM stack.

    Two modes selected by the number of appended layers:

    Shared weights  (1 layer):
        The same Entity processes every time step.
        Sequences may have any length.

    Unrolled / own weights  (N layers):
        layers[t % N] processes time step t — each position has its own W.
        Training requires sequences of exactly length N.
    """

    def __init__(self, name: str, work_dir: str = '.'):
        Stack.__init__(self, StackType.Rnn, name, work_dir)
        self._h: Mat | None = None  # context / hidden state: (batch, h_size)
        self._t: int = 0             # current time-step counter

    # ── Convenience factories ─────────────────────────────────────────────

    @staticmethod
    def make_layer(name: str, sensory_size: int, h_size: int,
                   entity_params: EntityParams, training_params: TrainingParams) -> Layer:
        """Single layer for shared-weights mode."""
        return Layer(name, (1, sensory_size, h_size, h_size), entity_params, training_params)

    @staticmethod
    def make_unrolled(time_steps: int, sensory_size: int, h_size: int,
                      entity_params: EntityParams, training_params: TrainingParams) -> list[Layer]:
        """N layers for unrolled (own-weights) mode — one per time step.

        Usage::

            for layer in StackRnn.make_unrolled(T, ...):
                rnn.append(layer)
        """
        return [
            Layer(f"t{t}", (1, sensory_size, h_size, h_size), entity_params, training_params)
            for t in range(time_steps)
        ]

    # ── Mode ─────────────────────────────────────────────────────────────

    @property
    def is_shared(self) -> bool:
        return self.num_layers() == 1

    def _entity_at(self, t: int) -> Entity:
        return self.from_index(t % self.num_layers()).entity

    def next_entity(self) -> Entity:
        """Entity that will be used by the next step() call."""
        return self._entity_at(self._t)

    def current_entity(self) -> Entity:
        """Entity used by the most recent step() call."""
        return self._entity_at(max(0, self._t - 1))

    # ── Derived sizes ─────────────────────────────────────────────────────

    def h_size(self, layer_idx: int = 0) -> int:
        return self.from_index(layer_idx).entity.shape[1]

    def sensory_size(self, layer_idx: int = 0) -> int:
        e = self.from_index(layer_idx).entity
        return e.shape[0] - e.shape[1]

    # ── State management ──────────────────────────────────────────────────

    def reset(self, batch_size: int = 1):
        """Zero the context vector and reset the time-step counter."""
        self._h = np.zeros((batch_size, self.h_size()))
        self._t = 0

    # ── Inference ─────────────────────────────────────────────────────────

    def step(self, x: Mat) -> Mat:
        """One time step forward.

        In shared mode  uses the single entity.
        In unrolled mode uses layers[_t % N] and advances _t.

        x: (batch_size, sensory_size) or (sensory_size,)
        Returns the new context vector h_t.
        """
        if x.ndim == 1:
            x = x[None, :]
        batch_size = x.shape[0]
        if self._h is None or self._h.shape[0] != batch_size:
            self.reset(batch_size)

        entity  = self._entity_at(self._t)
        visible = np.concatenate([self._h, _to_gpu(x)], axis=1)
        self._h = entity.forward(visible)
        self._t += 1
        return self._h

    def reconstruct(self, h: Mat) -> Mat:
        """Decode h → visible, returning only the sensory portion.

        Uses the entity from the most recent step() call.
        """
        entity  = self.current_entity()
        visible = entity.reconstruct(h)
        return visible[:, entity.shape[1]:]

    # ── Training ──────────────────────────────────────────────────────────

    def train(self, sequences: Mat, status: Status = None):
        """Train on sequences.

        sequences: (T, sensory_size)           — single sequence
                   (num_seq, T, sensory_size)  — batch of sequences

        Shared mode   (1 layer): T may be any value.
        Unrolled mode (N layers): T must equal N.
        """
        if status is None:
            status = Status()

        seqs = sequences if sequences.ndim == 3 else sequences[None, :]
        num_seq, T, _ = seqs.shape

        if self.is_shared:
            entity = self.from_index(0).entity
            print(f"Train shared ({entity.name}) "
                  f"for {entity.training_params.num_epochs} epochs")
            if entity.enable_training and entity.training_params is not None:
                self._train_shared(entity, seqs, num_seq, T, status)
        else:
            assert T == self.num_layers(), (
                f"Unrolled mode: sequence length T={T} "
                f"must equal num_layers={self.num_layers()}"
            )
            params = self.from_index(0).entity.training_params
            print(f"Train unrolled ({self.num_layers()} layers) "
                  f"for {params.num_epochs} epochs")
            self._train_unrolled(seqs, num_seq, T, status)

    def _train_shared(self, entity: Entity, seqs: Mat,
                      num_seq: int, T: int, status: Status):
        """One entity, reused at every time step."""
        cd_func = _CD_FUNC[entity.type]
        params  = entity.training_params
        h_sz    = entity.shape[1]

        d_progress   = 100.0 / params.num_epochs
        progress     = 0.0
        keep_running = True

        entity.grad_zero()
        status.on_change(entity)

        for epoch in range(params.num_epochs):
            if not keep_running:
                break

            h = np.zeros((num_seq, h_sz))
            err_total = 0.0

            for t in range(T):
                x_t     = _to_gpu(seqs[:, t, :])
                visible = np.concatenate([h, x_t], axis=1)

                dwhv, dbv, dbh = cd_func(entity, visible)
                grad = entity.grad_compute(dbv, dbh, dwhv)
                entity.state_adjust(grad, 1.0 / num_seq)

                h = entity.forward(visible)
                err_total += rms_error_accu(visible - entity.reconstruct(h))

            progress += d_progress
            if status.want_report(round(progress)):
                if not status.on_change(entity, {
                    "progress": {"value": round(progress), "unit": "%"},
                    "err_rms":  {"value": err_total / T,   "unit": ""},
                }):
                    keep_running = False
                    break

        h = np.zeros((num_seq, h_sz))
        err_total = 0.0
        for t in range(T):
            x_t     = _to_gpu(seqs[:, t, :])
            visible = np.concatenate([h, x_t], axis=1)
            h = entity.forward(visible)
            err_total += rms_error_accu(visible - entity.reconstruct(h))
        status.on_change(entity, {
            "progress":      {"value": 100,           "unit": "%"},
            "err_rms_total": {"value": err_total / T, "unit": ""},
        })

    def _train_unrolled(self, seqs: Mat, num_seq: int, T: int, status: Status):
        """N entities, one per time step — each has its own W, b_v, b_h.

        Uses temporal-shift padding (matching C++ RnnStack):
          Flatten (num_seq, T, sensory_size) → (N, sensory_size), append T-1 zero
          rows, then layer t trains on batch_padded[t : N+t] — a one-step delay.

        Joint training: all layers are updated together each epoch.
          Context c from layer t feeds layer t+1 within the same epoch pass,
          so gradients propagate through the full temporal chain.
        """
        params       = self.from_index(0).entity.training_params
        h_sz         = self.from_index(0).entity.shape[1]
        s_sz         = seqs.shape[2]
        d_progress   = 100.0 / params.num_epochs
        progress     = 0.0
        keep_running = True

        # Flatten to (N, s_sz) keeping sequence-major order, on CPU for slicing
        flat      = _np_cpu.asarray(convert(seqs) if hasattr(seqs, 'get') else seqs)
        flat      = flat.reshape(-1, s_sz)
        N         = flat.shape[0]
        pad       = _np_cpu.zeros((T - 1, s_sz), dtype=flat.dtype)
        batch_pad = _np_cpu.concatenate([flat, pad], axis=0)   # (N+T-1, s_sz)

        for layer in self.layers:
            layer.entity.grad_zero()
        status.on_change(self.from_index(0).entity)

        for epoch in range(params.num_epochs):
            if not keep_running:
                break

            c = np.zeros((N, h_sz))
            err_total = 0.0

            for t, layer in enumerate(self.layers):
                entity  = layer.entity
                cd_func = _CD_FUNC[entity.type]

                x_t     = _to_gpu(batch_pad[t : N + t])   # shifted slice, (N, s_sz)
                visible = np.concatenate([c, x_t], axis=1)

                dwhv, dbv, dbh = cd_func(entity, visible)
                grad = entity.grad_compute(dbv, dbh, dwhv)
                entity.state_adjust(grad, 1.0 / N)

                c = entity.forward(visible)
                err_total += rms_error_accu(visible - entity.reconstruct(c))

            progress += d_progress
            if status.want_report(round(progress)):
                if not status.on_change(self.from_index(0).entity, {
                    "progress": {"value": round(progress), "unit": "%"},
                    "err_rms":  {"value": err_total / T,   "unit": ""},
                }):
                    keep_running = False
                    break

        # Final report pass
        c = np.zeros((N, h_sz))
        err_total = 0.0
        for t, layer in enumerate(self.layers):
            entity  = layer.entity
            x_t     = _to_gpu(batch_pad[t : N + t])
            visible = np.concatenate([c, x_t], axis=1)
            c = entity.forward(visible)
            err_total += rms_error_accu(visible - entity.reconstruct(c))
        status.on_change(self.from_index(0).entity, {
            "progress":      {"value": 100,           "unit": "%"},
            "err_rms_total": {"value": err_total / T, "unit": ""},
        })
