#!/usr/bin/env python3
"""
helix_cdc/fibpi_runtime/decoders/simple_block_decoder.py

WO-FIBPI-DNA-05 + WO-FIBPI-DNA-11:
Tiny neural decoder with Se (Symbolic Entropy) steering:
(z, anchor, band, se) -> tensor_block

v2: Random projection z forces decoder to use anchor/band/Se.
Se steering tells decoder how loud/smooth/noisy the region should be.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Optional, Tuple

import numpy as np

try:
    import torch
    from torch import nn
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False
    nn = None


@dataclass
class SimpleDecoderConfig:
    latent_dim: int = 32
    anchor_dim: int = 3
    block_shape: Tuple[int, int] = (64, 64)
    hidden_dim: int = 256
    num_bands: int = 8
    se_dim: int = 6  # NEW: Se steering dimension


if TORCH_AVAILABLE:
    class SimpleBlockDecoder(nn.Module):
        def __init__(self, cfg: SimpleDecoderConfig):
            super().__init__()
            self.cfg = cfg
            # Input: z + anchor + band_onehot + se_features
            input_dim = cfg.latent_dim + cfg.anchor_dim + cfg.num_bands + cfg.se_dim
            output_dim = cfg.block_shape[0] * cfg.block_shape[1]

            self.net = nn.Sequential(
                nn.Linear(input_dim, cfg.hidden_dim),
                nn.ReLU(),
                nn.LayerNorm(cfg.hidden_dim),
                nn.Linear(cfg.hidden_dim, cfg.hidden_dim),
                nn.ReLU(),
                nn.LayerNorm(cfg.hidden_dim),
                nn.Linear(cfg.hidden_dim, cfg.hidden_dim // 2),
                nn.ReLU(),
                nn.Linear(cfg.hidden_dim // 2, output_dim),
            )

        def forward(
            self,
            z: torch.Tensor,
            anchor: torch.Tensor,
            band_onehot: torch.Tensor,
            se_feats: Optional[torch.Tensor] = None,
        ) -> torch.Tensor:
            """
            z: [B, latent_dim]
            anchor: [B, anchor_dim]
            band_onehot: [B, num_bands]
            se_feats: [B, se_dim] (optional, defaults to zeros)
            returns: [B, H*W]
            """
            if se_feats is None:
                se_feats = torch.zeros(z.shape[0], self.cfg.se_dim, device=z.device)
            x = torch.cat([z, anchor, band_onehot, se_feats], dim=-1)
            return self.net(x)

        def decode_block(
            self,
            z: np.ndarray,
            anchor: np.ndarray,
            band_id: int,
            se_feats: Optional[np.ndarray] = None,
        ) -> np.ndarray:
            """
            Decode a single block from latent z, anchor coords, band_id, and Se features.
            Returns numpy array of shape block_shape.
            """
            z_t = torch.from_numpy(z.astype(np.float32)).unsqueeze(0)
            anchor_t = torch.from_numpy(anchor.astype(np.float32)).unsqueeze(0)
            band_oh = band_to_onehot(band_id, self.cfg.num_bands)
            band_t = torch.from_numpy(band_oh).unsqueeze(0)

            if se_feats is not None:
                se_t = torch.from_numpy(se_feats.astype(np.float32)).unsqueeze(0)
            else:
                se_t = torch.zeros(1, self.cfg.se_dim)

            with torch.no_grad():
                out = self.forward(z_t, anchor_t, band_t, se_t)

            return out.squeeze(0).numpy().reshape(self.cfg.block_shape)
else:
    class SimpleBlockDecoder:
        """Stub decoder when torch is not available."""
        def __init__(self, cfg: SimpleDecoderConfig):
            self.cfg = cfg
            print("[WARNING] torch not available, using stub decoder")

        def decode_block(
            self,
            z: np.ndarray,
            anchor: np.ndarray,
            band_id: int,
            se_feats: Optional[np.ndarray] = None,
        ) -> np.ndarray:
            """Return zeros as stub."""
            return np.zeros(self.cfg.block_shape, dtype=np.float32)


def encode_block_to_latent(block: np.ndarray, cfg: SimpleDecoderConfig) -> np.ndarray:
    """
    v1: dumb encoder: just take the first latent_dim entries of flattened block.
    This is intentionally stupid but gives us a testable path.
    """
    flat = block.astype("float32").flatten()
    if flat.size < cfg.latent_dim:
        flat = np.pad(flat, (0, cfg.latent_dim - flat.size))
    return flat[: cfg.latent_dim]


def band_to_onehot(band_id: int, num_bands: int) -> np.ndarray:
    v = np.zeros((num_bands,), dtype="float32")
    v[band_id] = 1.0
    return v


def decode_helix64_to_latent(code: str, latent_dim: int = 32) -> np.ndarray:
    """
    Decode Helix64 (base64) encoded float32 bytes back to latent vector.
    """
    import base64
    data = base64.b64decode(code)
    arr = np.frombuffer(data, dtype=np.float32)
    if arr.size < latent_dim:
        arr = np.pad(arr, (0, latent_dim - arr.size))
    return arr[:latent_dim]


def main() -> None:
    """Test the decoder with random inputs."""
    cfg = SimpleDecoderConfig()
    decoder = SimpleBlockDecoder(cfg)

    # Random test
    z = np.random.randn(cfg.latent_dim).astype(np.float32)
    anchor = np.array([0.5, 0.5, 0.5], dtype=np.float32)
    band_id = 4

    block = decoder.decode_block(z, anchor, band_id)
    print(f"[SIMPLE-DECODER] Output shape: {block.shape}")
    print(f"[SIMPLE-DECODER] Output range: [{block.min():.4f}, {block.max():.4f}]")


if __name__ == "__main__":
    main()
