#!/usr/bin/env python3
"""
flowtorch_helix_ops.py
======================

FlowTorch-compatible differentiable Helix substrate operations.

This module provides PyTorch-native layers that implement the Helix
native generation paradigm with backpropagation support.

KEY INSIGHT (from native_helix_ratio_test.py):
- Native Helix achieves 510:1 by SHARING codebook across layers
- Each layer only needs a tiny seed (25-30 bytes)
- The codebook is amortized across ALL layers

This module makes that DIFFERENTIABLE:
- HelixCodebook: Shared learnable codebook (one per model)
- HelixLinear: Linear layer that generates weights from seed + codebook
- HelixMLP: Full MLP using native Helix substrate

The effective parameter count is:
  P_effective = codebook_params + sum(seed_params_per_layer)

For N layers:
  P_effective ≈ codebook_size * pattern_dim + N * seed_dim

This is MUCH smaller than N * (in_features * out_features) for dense.
"""

import math
import hashlib
import struct
from typing import Tuple, Optional, List

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

# Import the REAL Se and substrate from user's existing code
try:
    from helix_cdc.helix64_phi_codebook import compute_se, se_penalty
    from helix_cdc.scheduler.wave_scheduler import fibpi_wave_schedule
    SE_AVAILABLE = True
except ImportError:
    SE_AVAILABLE = False

# Constants from Helix substrate
PHI = 1.618033988749895
PI = math.pi
GOLDEN_ANGLE = 2.39996322972865332


# ============================================================================
# Se Steering for Training (Uses User's Real Se = log1p(H × C × D))
# ============================================================================

def compute_se_torch(tensor: torch.Tensor) -> float:
    """
    Compute Se = log1p(H × C × D) on a PyTorch tensor.

    Uses the user's real formula from helix64_phi_codebook.py
    """
    if SE_AVAILABLE:
        return compute_se(tensor.detach().cpu().numpy())
    else:
        # Fallback: simple entropy measure
        v = tensor.detach().flatten()
        H = torch.std(v).item() * 2  # Proxy for entropy
        return float(np.log1p(H))


def se_penalty_torch(tensor: torch.Tensor, target_band: Tuple[float, float] = (0.3, 0.7)) -> float:
    """
    Compute Se penalty for steering weights into healthy band.

    This is the constraint that keeps generated weights well-behaved.
    """
    if SE_AVAILABLE:
        return se_penalty(tensor.detach().cpu().numpy(), target_band)
    else:
        se = compute_se_torch(tensor)
        se_low, se_high = target_band
        if se_low <= se <= se_high:
            return 0.0
        elif se < se_low:
            return (se_low - se) ** 2
        else:
            return (se - se_high) ** 2


class SeSteeredLoss(nn.Module):
    """
    Loss wrapper that adds Se steering penalty.

    total_loss = base_loss + se_weight * sum(se_penalty(layer_weights))

    This steers the learned codebook and seeds toward generating
    weights with healthy Se values (0.3-0.7 target band).
    """

    def __init__(self, base_criterion: nn.Module, se_weight: float = 0.01, target_band: Tuple[float, float] = (0.3, 0.7)):
        super().__init__()
        self.base_criterion = base_criterion
        self.se_weight = se_weight
        self.target_band = target_band

    def forward(self, output: torch.Tensor, target: torch.Tensor, model: nn.Module) -> torch.Tensor:
        base_loss = self.base_criterion(output, target)

        # Compute Se penalty over all generated weights
        se_loss = 0.0
        for name, param in model.named_parameters():
            if 'weight' in name.lower() or 'pattern' in name.lower():
                se_loss += se_penalty_torch(param, self.target_band)

        return base_loss + self.se_weight * se_loss


# ============================================================================
# FibPi Learning Rate Scheduler (Se-Controlled)
# ============================================================================

class FibPiLRScheduler:
    """
    Fibonacci/Pi spiral learning rate scheduler with Se control.

    KEY INSIGHT from user's files:
    - "Use Se as a scheduler, not just a physics toy"
    - Learning rate follows golden ratio spiral
    - Se band fraction controls adaptation speed

    The schedule combines:
    1. FibPi spiral: lr modulated by φ-spiral phase
    2. Se feedback: in_band_fraction adjusts lr magnitude
    3. Warmup: gradual ramp-up at start
    """

    def __init__(
        self,
        optimizer,
        total_steps: int,
        lr_max: float = 1e-2,
        lr_min: float = 1e-5,
        warmup_steps: int = 100,
        target_band: Tuple[float, float] = (0.3, 0.7),
    ):
        self.optimizer = optimizer
        self.total_steps = total_steps
        self.lr_max = lr_max
        self.lr_min = lr_min
        self.warmup_steps = warmup_steps
        self.target_band = target_band

        # Tracking state
        self.current_step = 0
        self.in_band_count = 0
        self.total_count = 0
        self.se_history = []

        # FibPi constants
        self.phi = PHI
        self.golden_angle = GOLDEN_ANGLE

        # Generate spiral schedule
        self._precompute_spiral()

    def _precompute_spiral(self):
        """Precompute FibPi spiral modulation for all steps."""
        self.spiral_mod = []
        for step in range(self.total_steps):
            # Normalized progress
            t = step / max(self.total_steps - 1, 1)

            # FibPi spiral phase
            phase = t * self.phi * self.golden_angle * 2 * PI

            # Spiral modulation (oscillates between 0 and 1)
            mod = 0.5 * (1 + math.cos(phase))

            # Apply golden ratio decay envelope
            envelope = 1.0 / (1.0 + t * self.phi)

            self.spiral_mod.append(mod * envelope)

    def get_in_band_fraction(self) -> float:
        """Get fraction of steps where Se was in target band."""
        if self.total_count == 0:
            return 0.5  # Neutral starting point
        return self.in_band_count / self.total_count

    def update_se(self, se_value: float):
        """Update Se tracking with new measurement."""
        self.se_history.append(se_value)
        self.total_count += 1

        se_low, se_high = self.target_band
        if se_low <= se_value <= se_high:
            self.in_band_count += 1

    def get_lr(self) -> float:
        """Get current learning rate with FibPi + Se modulation."""
        step = min(self.current_step, self.total_steps - 1)

        # Warmup phase
        if step < self.warmup_steps:
            warmup_factor = step / self.warmup_steps
            base_lr = self.lr_min + (self.lr_max - self.lr_min) * warmup_factor
        else:
            # FibPi spiral schedule
            spiral = self.spiral_mod[step]
            base_lr = self.lr_min + (self.lr_max - self.lr_min) * spiral

        # Se feedback modulation
        in_band_frac = self.get_in_band_fraction()

        # When in-band is high, we're stable -> can use higher lr
        # When in-band is low, Se is unstable -> reduce lr
        se_factor = 0.5 + 0.5 * in_band_frac  # Range: 0.5 to 1.0

        return base_lr * se_factor

    def step(self, se_value: Optional[float] = None):
        """Advance scheduler by one step."""
        if se_value is not None:
            self.update_se(se_value)

        lr = self.get_lr()
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = lr

        self.current_step += 1
        return lr

    def state_dict(self) -> dict:
        """Return scheduler state for checkpointing."""
        return {
            'current_step': self.current_step,
            'in_band_count': self.in_band_count,
            'total_count': self.total_count,
            'se_history': self.se_history[-100:],  # Keep last 100
        }

    def load_state_dict(self, state: dict):
        """Load scheduler state from checkpoint."""
        self.current_step = state['current_step']
        self.in_band_count = state['in_band_count']
        self.total_count = state['total_count']
        self.se_history = state.get('se_history', [])


class SeScheduledTrainer:
    """
    Training loop with Se-controlled scheduling.

    This implements the "fluid dynamics" approach where:
    - Se is measured on model weights each batch
    - in_band_fraction controls learning rate
    - FibPi spiral provides base schedule

    The key insight: Se acts as a control signal, not just a penalty.
    """

    def __init__(
        self,
        model: nn.Module,
        optimizer,
        criterion: nn.Module,
        lr_max: float = 1e-2,
        lr_min: float = 1e-5,
        target_band: Tuple[float, float] = (0.3, 0.7),
        se_weight: float = 0.01,
    ):
        self.model = model
        self.optimizer = optimizer
        self.base_criterion = criterion
        self.target_band = target_band
        self.se_weight = se_weight

        # Will be initialized on first train call
        self.scheduler = None
        self.lr_max = lr_max
        self.lr_min = lr_min

        # Metrics tracking
        self.metrics = {
            'train_loss': [],
            'val_loss': [],
            'train_acc': [],
            'val_acc': [],
            'se_values': [],
            'in_band_fraction': [],
            'learning_rates': [],
        }

    def compute_model_se(self) -> float:
        """Compute average Se across all weight tensors."""
        se_values = []
        for name, param in self.model.named_parameters():
            if 'weight' in name.lower() or 'pattern' in name.lower():
                se = compute_se_torch(param)
                se_values.append(se)

        return np.mean(se_values) if se_values else 0.5

    def train_epoch(self, dataloader) -> Tuple[float, float]:
        """Train for one epoch with Se scheduling."""
        self.model.train()
        total_loss = 0.0
        correct = 0
        total = 0

        for batch_idx, (data, target) in enumerate(dataloader):
            self.optimizer.zero_grad()

            # Forward pass
            output = self.model(data)

            # Base loss
            loss = self.base_criterion(output, target)

            # Se penalty (adds steering)
            se_loss = 0.0
            for name, param in self.model.named_parameters():
                if 'weight' in name.lower() or 'pattern' in name.lower():
                    se_loss += se_penalty_torch(param, self.target_band)

            total_batch_loss = loss + self.se_weight * se_loss

            # Backward pass
            total_batch_loss.backward()
            self.optimizer.step()

            # Measure Se AFTER update and feed to scheduler
            current_se = self.compute_model_se()
            lr = self.scheduler.step(current_se)

            # Track metrics
            total_loss += loss.item()
            _, predicted = output.max(1)
            total += target.size(0)
            correct += predicted.eq(target).sum().item()

            self.metrics['se_values'].append(current_se)
            self.metrics['learning_rates'].append(lr)

        avg_loss = total_loss / len(dataloader)
        accuracy = correct / total

        self.metrics['in_band_fraction'].append(self.scheduler.get_in_band_fraction())

        return avg_loss, accuracy

    def validate(self, dataloader) -> Tuple[float, float]:
        """Validate model."""
        self.model.eval()
        total_loss = 0.0
        correct = 0
        total = 0

        with torch.no_grad():
            for data, target in dataloader:
                output = self.model(data)
                loss = self.base_criterion(output, target)

                total_loss += loss.item()
                _, predicted = output.max(1)
                total += target.size(0)
                correct += predicted.eq(target).sum().item()

        return total_loss / len(dataloader), correct / total

    def train(
        self,
        train_loader,
        val_loader,
        epochs: int,
        verbose: bool = True,
    ) -> dict:
        """Full training loop with Se scheduling."""
        total_steps = epochs * len(train_loader)

        # Initialize scheduler
        self.scheduler = FibPiLRScheduler(
            self.optimizer,
            total_steps=total_steps,
            lr_max=self.lr_max,
            lr_min=self.lr_min,
            warmup_steps=min(100, total_steps // 10),
            target_band=self.target_band,
        )

        best_val_loss = float('inf')
        best_val_acc = 0.0

        for epoch in range(epochs):
            train_loss, train_acc = self.train_epoch(train_loader)
            val_loss, val_acc = self.validate(val_loader)

            self.metrics['train_loss'].append(train_loss)
            self.metrics['val_loss'].append(val_loss)
            self.metrics['train_acc'].append(train_acc)
            self.metrics['val_acc'].append(val_acc)

            if val_loss < best_val_loss:
                best_val_loss = val_loss
            if val_acc > best_val_acc:
                best_val_acc = val_acc

            if verbose:
                in_band = self.scheduler.get_in_band_fraction()
                current_lr = self.scheduler.get_lr()
                print(f"Epoch {epoch+1:3d}/{epochs}: "
                      f"Train Loss: {train_loss:.4f}, Acc: {train_acc:.3f} | "
                      f"Val Loss: {val_loss:.4f}, Acc: {val_acc:.3f} | "
                      f"Se in-band: {in_band:.2%}, LR: {current_lr:.2e}")

        return {
            'final_train_loss': train_loss,
            'final_val_loss': val_loss,
            'final_train_acc': train_acc,
            'final_val_acc': val_acc,
            'best_val_loss': best_val_loss,
            'best_val_acc': best_val_acc,
            'in_band_fraction': self.scheduler.get_in_band_fraction(),
            'history': self.metrics,
        }


class HelixCodebook(nn.Module):
    """
    Shared learnable codebook for Native Helix generation.

    This is the "Einstein layer" - structure that's shared across all layers.
    One codebook instance should be shared by all HelixLinear layers in a model.

    The codebook contains pattern_dim-dimensional patterns that can be
    combined via learned mixing weights to generate any weight matrix.

    Storage cost: codebook_size * pattern_dim * 4 bytes
    """

    def __init__(
        self,
        codebook_size: int = 256,
        pattern_dim: int = 64,
        init_scale: float = 1.0,  # FIXED: Was 0.01, too small for gradient flow
    ):
        super().__init__()
        self.codebook_size = codebook_size
        self.pattern_dim = pattern_dim

        # Learnable codebook patterns
        # Initialize with FibPi3D-like geometry for structure
        # Note: init_scale now defaults to 1.0 for proper gradient flow
        patterns = self._fibpi3d_init(codebook_size, pattern_dim)
        self.patterns = nn.Parameter(torch.from_numpy(patterns * init_scale))

    def _fibpi3d_init(self, size: int, dim: int) -> np.ndarray:
        """Initialize patterns using FibPi3D geometry."""
        patterns = np.zeros((size, dim), dtype=np.float32)

        for i in range(size):
            # FibPi3D coordinate
            t = i * PHI
            angle = (t * GOLDEN_ANGLE) % (2 * PI)
            z = (i % size) / size
            r = math.sqrt(1 - z * z) if z < 1 else 0
            x, y = r * math.cos(angle), r * math.sin(angle)

            # Use coordinates to seed deterministic pattern
            seed_val = int(abs(x * 1e6) + abs(y * 1e6) + abs(z * 1e6)) % (2**31)
            rng = np.random.RandomState(seed_val)

            pattern = rng.randn(dim).astype(np.float32)
            pattern /= (np.linalg.norm(pattern) + 1e-10)
            patterns[i] = pattern

        return patterns

    def get_pattern(self, idx: torch.Tensor) -> torch.Tensor:
        """Get patterns by index (differentiable via embedding lookup)."""
        return F.embedding(idx, self.patterns)

    def param_count(self) -> int:
        """Return number of parameters in codebook."""
        return self.codebook_size * self.pattern_dim

    def storage_bytes(self) -> int:
        """Return storage size in bytes."""
        return self.param_count() * 4  # float32


class HelixSeed(nn.Module):
    """
    Learnable seed parameters for a single layer.

    The seed determines how to mix codebook patterns to generate weights.
    This is the per-layer "addressing" into the shared codebook.

    Storage cost: n_mix * 2 * 4 bytes (indices + weights)
    """

    def __init__(
        self,
        n_mix: int = 32,
        codebook_size: int = 256,
        init_scale: float = 1.0,  # FIXED: Was 0.1, too small
    ):
        super().__init__()
        self.n_mix = n_mix
        self.codebook_size = codebook_size

        # Soft indices into codebook (continuous for differentiability)
        # We'll use softmax over these to get soft-attention over patterns
        # Initialize with small values so softmax starts near uniform
        self.pattern_logits = nn.Parameter(torch.randn(n_mix, codebook_size) * 0.1)

        # Mixing weights - initialize larger for meaningful contribution
        # FIXED: init_scale now 1.0 by default
        self.mix_weights = nn.Parameter(torch.randn(n_mix) * init_scale)

    def get_soft_patterns(self, codebook: HelixCodebook) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Get soft-selected patterns and their weights.

        Returns:
            patterns: (n_mix, pattern_dim) - soft-selected patterns
            weights: (n_mix,) - mixing weights for each pattern
        """
        # Soft attention over codebook
        attn = F.softmax(self.pattern_logits, dim=-1)  # (n_mix, codebook_size)

        # Soft pattern selection via attention
        patterns = torch.matmul(attn, codebook.patterns)  # (n_mix, pattern_dim)

        return patterns, self.mix_weights

    def param_count(self) -> int:
        """Return number of seed parameters."""
        return self.n_mix * self.codebook_size + self.n_mix

    def storage_bytes(self) -> int:
        """Return storage size in bytes."""
        return self.param_count() * 4


class HelixLinear(nn.Module):
    """
    Linear layer that generates weights from Helix substrate.

    Instead of storing (in_features * out_features) parameters,
    this layer generates weights from:
    - A shared codebook (passed in)
    - A small seed (n_mix soft pattern selections + weights)

    The weight generation is:
        W = seed.generate(codebook, shape) * scale

    This is fully differentiable - gradients flow through to:
    - codebook.patterns (shared)
    - seed.pattern_logits (per-layer)
    - seed.mix_weights (per-layer)

    Effective parameter count:
        seed: n_mix * (codebook_size + 1)

    vs Dense Linear:
        weight: in_features * out_features
        bias: out_features
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        codebook: HelixCodebook,
        n_mix: int = 32,
        bias: bool = True,
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.codebook = codebook  # Shared, not owned

        # Per-layer seed
        self.seed = HelixSeed(n_mix=n_mix, codebook_size=codebook.codebook_size)

        # Learnable scale (for proper initialization)
        # FIXED: Use Xavier-like scaling but start at a meaningful value
        fan_in = in_features
        # Xavier: 1/sqrt(fan_in) for MNIST first layer = 1/sqrt(784) = 0.036
        # We want the generated weights to have std ~= Xavier init, which is sqrt(2/(fan_in+fan_out))
        # But the weight generation already produces unit-variance patterns mixed together
        # So we scale DOWN to achieve Xavier-like std
        self.scale = nn.Parameter(torch.tensor(math.sqrt(2.0 / (fan_in + out_features))))

        # Optional bias (always per-layer, small)
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features))
        else:
            self.register_parameter('bias', None)

    def generate_weight(self) -> torch.Tensor:
        """Generate the weight matrix from seed + codebook."""
        patterns, mix_weights = self.seed.get_soft_patterns(self.codebook)

        # patterns: (n_mix, pattern_dim)
        # mix_weights: (n_mix,)
        pattern_dim = patterns.shape[1]
        n_mix = patterns.shape[0]

        # Target size
        total_elements = self.in_features * self.out_features

        # Generate weight by tiling and mixing patterns
        # This mirrors the NativeHelixGenerator logic but differentiably

        # Strategy: create weight from weighted sum of tiled patterns
        weight = torch.zeros(total_elements, device=patterns.device, dtype=patterns.dtype)

        # Tile each pattern across the weight and mix
        for i in range(n_mix):
            # Tile pattern to cover all elements
            n_tiles = (total_elements + pattern_dim - 1) // pattern_dim
            tiled = patterns[i].repeat(n_tiles)[:total_elements]

            # Position-modulated mixing (FibPi-inspired)
            positions = torch.arange(total_elements, device=patterns.device, dtype=patterns.dtype)
            phase = (positions * PHI / total_elements * GOLDEN_ANGLE) % (2 * PI)
            pos_mod = torch.cos(phase + i * GOLDEN_ANGLE)

            weight += mix_weights[i] * tiled * pos_mod

        # Reshape to matrix
        weight = weight.view(self.out_features, self.in_features)

        return weight * self.scale

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        weight = self.generate_weight()
        return F.linear(x, weight, self.bias)

    def native_param_count(self) -> int:
        """Return native (seed) parameter count (excluding shared codebook)."""
        count = self.seed.param_count() + 1  # +1 for scale
        if self.bias is not None:
            count += self.out_features
        return count

    def dense_param_count(self) -> int:
        """Return equivalent dense layer parameter count."""
        count = self.in_features * self.out_features
        if self.bias is not None:
            count += self.out_features
        return count


class HelixMLP(nn.Module):
    """
    Multi-layer perceptron using Native Helix substrate.

    All layers share a single codebook - this is where the
    massive parameter savings come from.

    Parameter count comparison:
    - Dense MLP: sum(layer_i_params) for all layers
    - Helix MLP: codebook_params + sum(seed_i_params) for all layers

    The codebook is amortized across layers, so more layers = bigger savings.
    """

    def __init__(
        self,
        layer_dims: List[int],
        codebook_size: int = 256,
        pattern_dim: int = 64,
        n_mix: int = 32,
        activation: str = "relu",
    ):
        super().__init__()

        # Shared codebook (the "Einstein" structure)
        self.codebook = HelixCodebook(
            codebook_size=codebook_size,
            pattern_dim=pattern_dim,
        )

        # Build layers
        self.layers = nn.ModuleList()
        for i in range(len(layer_dims) - 1):
            layer = HelixLinear(
                in_features=layer_dims[i],
                out_features=layer_dims[i + 1],
                codebook=self.codebook,
                n_mix=n_mix,
                bias=True,
            )
            self.layers.append(layer)

        # Activation
        if activation == "relu":
            self.act = nn.ReLU()
        elif activation == "gelu":
            self.act = nn.GELU()
        elif activation == "tanh":
            self.act = nn.Tanh()
        else:
            self.act = nn.ReLU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for i, layer in enumerate(self.layers):
            x = layer(x)
            if i < len(self.layers) - 1:  # No activation on last layer
                x = self.act(x)
        return x

    def native_param_count(self) -> int:
        """Total native parameter count (codebook + all seeds)."""
        total = self.codebook.param_count()
        for layer in self.layers:
            total += layer.native_param_count()
        return total

    def dense_equivalent_param_count(self) -> int:
        """Equivalent dense MLP parameter count."""
        total = 0
        for layer in self.layers:
            total += layer.dense_param_count()
        return total

    def compression_ratio(self) -> float:
        """Compression ratio vs equivalent dense MLP."""
        native = self.native_param_count()
        dense = self.dense_equivalent_param_count()
        return dense / native if native > 0 else 0.0

    def summary(self) -> str:
        """Return a summary of parameter counts."""
        native = self.native_param_count()
        dense = self.dense_equivalent_param_count()
        ratio = self.compression_ratio()

        lines = [
            "=" * 50,
            "HelixMLP Parameter Summary",
            "=" * 50,
            f"Codebook params:     {self.codebook.param_count():,}",
        ]

        for i, layer in enumerate(self.layers):
            lines.append(f"Layer {i} seed params: {layer.native_param_count():,}")

        lines.extend([
            "-" * 50,
            f"Total Native params: {native:,}",
            f"Dense equivalent:    {dense:,}",
            f"Compression ratio:   {ratio:.1f}:1",
            "=" * 50,
        ])

        return "\n".join(lines)


class DenseMLPBaseline(nn.Module):
    """
    Standard dense MLP for comparison.
    Same architecture as HelixMLP but with regular Linear layers.
    """

    def __init__(
        self,
        layer_dims: List[int],
        activation: str = "relu",
    ):
        super().__init__()

        self.layers = nn.ModuleList()
        for i in range(len(layer_dims) - 1):
            self.layers.append(nn.Linear(layer_dims[i], layer_dims[i + 1]))

        if activation == "relu":
            self.act = nn.ReLU()
        elif activation == "gelu":
            self.act = nn.GELU()
        elif activation == "tanh":
            self.act = nn.Tanh()
        else:
            self.act = nn.ReLU()

        self._init_weights()

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for i, layer in enumerate(self.layers):
            x = layer(x)
            if i < len(self.layers) - 1:
                x = self.act(x)
        return x

    def param_count(self) -> int:
        return sum(p.numel() for p in self.parameters())


# Convenience function for creating matched pairs
def create_matched_models(
    layer_dims: List[int],
    codebook_size: int = 256,
    pattern_dim: int = 64,
    n_mix: int = 32,
) -> Tuple[HelixMLP, DenseMLPBaseline]:
    """
    Create a Helix MLP and a matched Dense MLP baseline for comparison.

    Returns:
        helix_model: HelixMLP with native substrate
        dense_model: DenseMLPBaseline with same architecture
    """
    helix = HelixMLP(
        layer_dims=layer_dims,
        codebook_size=codebook_size,
        pattern_dim=pattern_dim,
        n_mix=n_mix,
    )

    dense = DenseMLPBaseline(layer_dims=layer_dims)

    return helix, dense
