#!/usr/bin/env python3
"""
NativeHelixLinear v0.1 - Differentiable Helix-Native Layer

This is the MINIMAL trainable Helix layer that:
1. Uses seeds (not full weight matrices) as parameters
2. Decodes weights via φ/FibPi math (differentiable in Torch)
3. Can actually learn (gradients flow back to seeds)

NO fancy codebooks, NO soft-attention, NO Se scheduling.
Just: seeds → φ decode → weights → matmul.

Key insight: The φ/FibPi substrate must be DIFFERENTIABLE.
All math is in pure PyTorch - no numpy, no detach.

Constants:
    PHI = (1 + sqrt(5)) / 2 ≈ 1.618034
    GOLDEN_ANGLE = 2π / φ² ≈ 2.399963 radians
"""

import math
import torch
import torch.nn as nn
from typing import Optional, Tuple

# Golden constants
PHI = (1.0 + math.sqrt(5.0)) / 2.0  # 1.618034...
GOLDEN_ANGLE = 2.0 * math.pi / (PHI ** 2)  # ~2.399963 radians


def decode_phi_torch(
    seeds: torch.Tensor,
    out_dim: int,
    n_harmonics: int = 8
) -> torch.Tensor:
    """
    Decode seeds into weight matrix using φ/FibPi basis.

    FULLY DIFFERENTIABLE - gradients flow back to seeds.

    Architecture:
        1. Each seed row encodes one output neuron's weight row
        2. Seeds are coefficients for φ-harmonic basis functions
        3. Basis = sin/cos at golden-angle-spaced frequencies

    The idea:
        W[i, j] = Σ_k (seed[i, k] * basis_k(j))

    Where basis_k uses golden-angle phase offsets and Fibonacci-scaled
    frequencies, creating a structured weight space.

    Args:
        seeds: [n_neurons, seed_dim] - trainable seed coefficients
        out_dim: Number of input features (columns in weight matrix)
        n_harmonics: Number of harmonic basis functions per seed element

    Returns:
        W: [n_neurons, out_dim] - decoded weight matrix
    """
    n_neurons, seed_dim = seeds.shape
    device = seeds.device
    dtype = seeds.dtype

    # Create position indices for output dimension
    # Shape: [out_dim]
    positions = torch.linspace(0, 1, out_dim, device=device, dtype=dtype)

    # Generate golden-angle-spaced phases
    # Shape: [n_harmonics]
    harmonic_indices = torch.arange(n_harmonics, device=device, dtype=dtype)
    phases = GOLDEN_ANGLE * harmonic_indices  # Golden angle accumulation

    # Generate Fibonacci-scaled frequencies (1, 1, 2, 3, 5, 8, 13, ...)
    # Shape: [n_harmonics]
    fib_freqs = _fibonacci_frequencies(n_harmonics, device, dtype)

    # Compute basis functions
    # For each position and harmonic: sin(2π * fib_freq * position + phase)
    # Shape: [out_dim, n_harmonics]
    angles = 2.0 * math.pi * fib_freqs.unsqueeze(0) * positions.unsqueeze(1) + phases.unsqueeze(0)
    basis_sin = torch.sin(angles)  # [out_dim, n_harmonics]
    basis_cos = torch.cos(angles)  # [out_dim, n_harmonics]

    # Interleave sin/cos to get 2*n_harmonics basis functions
    # Shape: [out_dim, 2*n_harmonics]
    basis = torch.zeros(out_dim, 2 * n_harmonics, device=device, dtype=dtype)
    basis[:, 0::2] = basis_sin
    basis[:, 1::2] = basis_cos

    # Pad or truncate basis to match seed_dim
    if 2 * n_harmonics > seed_dim:
        basis = basis[:, :seed_dim]
    elif 2 * n_harmonics < seed_dim:
        # Pad with additional harmonics - enough to cover seed_dim
        extra_needed = seed_dim - 2 * n_harmonics
        n_extra_harmonics = (extra_needed + 1) // 2 + 1  # enough for sin+cos pairs

        extra_indices = torch.arange(n_extra_harmonics, device=device, dtype=dtype)
        extra_phases = GOLDEN_ANGLE * (extra_indices + n_harmonics)
        extra_freqs = _fibonacci_frequencies(n_extra_harmonics, device, dtype) * PHI
        extra_angles = 2.0 * math.pi * extra_freqs.unsqueeze(0) * positions.unsqueeze(1) + extra_phases.unsqueeze(0)

        extra_sin = torch.sin(extra_angles)  # [out_dim, n_extra_harmonics]
        extra_cos = torch.cos(extra_angles)  # [out_dim, n_extra_harmonics]

        # Interleave extra sin/cos
        extra_basis = torch.zeros(out_dim, 2 * n_extra_harmonics, device=device, dtype=dtype)
        extra_basis[:, 0::2] = extra_sin
        extra_basis[:, 1::2] = extra_cos

        basis = torch.cat([basis, extra_basis[:, :extra_needed]], dim=1)

    # Decode: W = seeds @ basis.T
    # [n_neurons, seed_dim] @ [seed_dim, out_dim] → [n_neurons, out_dim]
    W = seeds @ basis.t()

    return W


def _fibonacci_frequencies(n: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
    """Generate first n Fibonacci numbers as frequencies, normalized."""
    fib = [1.0, 1.0]
    for _ in range(n - 2):
        fib.append(fib[-1] + fib[-2])
    fib_tensor = torch.tensor(fib[:n], device=device, dtype=dtype)
    # Normalize to reasonable range [0.5, ~8]
    return 0.5 + fib_tensor / fib_tensor.max() * 7.5


class NativeHelixLinear(nn.Module):
    """
    Native Helix Linear Layer - trainable, differentiable, compressed.

    Instead of storing [out_features, in_features] weights directly,
    we store [out_features, seed_dim] seeds and decode via φ/FibPi.

    Compression ratio: in_features / seed_dim

    Example:
        layer = NativeHelixLinear(784, 128, seed_dim=32)
        # Dense would have 784*128 = 100,352 params
        # Native has 128*32 = 4,096 params (24.5:1 compression)
        # Plus bias: 128 params

    The layer IS trainable - gradients flow to self.seeds.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        seed_dim: int = 16,
        bias: bool = True,
        n_harmonics: int = 8
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.seed_dim = seed_dim
        self.n_harmonics = n_harmonics

        # Trainable seeds - the actual parameters
        self.seeds = nn.Parameter(torch.randn(out_features, seed_dim) * 0.1)

        # Optional bias (not compressed - just regular bias)
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features))
        else:
            self.register_parameter('bias', None)

        # Stats for comparison
        self._dense_params = in_features * out_features + (out_features if bias else 0)
        self._native_params = out_features * seed_dim + (out_features if bias else 0)

    @property
    def compression_ratio(self) -> float:
        """Return compression ratio vs dense layer."""
        return self._dense_params / self._native_params

    def decode_weights(self) -> torch.Tensor:
        """Decode seeds into full weight matrix (differentiable)."""
        return decode_phi_torch(self.seeds, self.in_features, self.n_harmonics)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass: decode weights, then matmul.

        Args:
            x: [batch, in_features]

        Returns:
            out: [batch, out_features]
        """
        W = self.decode_weights()  # [out_features, in_features]
        out = x @ W.t()  # [batch, out_features]
        if self.bias is not None:
            out = out + self.bias
        return out

    def extra_repr(self) -> str:
        return (
            f"in_features={self.in_features}, "
            f"out_features={self.out_features}, "
            f"seed_dim={self.seed_dim}, "
            f"compression={self.compression_ratio:.1f}:1, "
            f"bias={self.bias is not None}"
        )


class NativeHelixMLP(nn.Module):
    """
    MLP using NativeHelixLinear layers.

    For testing on simple tasks like spiral classification.
    """

    def __init__(
        self,
        in_features: int,
        hidden_features: int,
        out_features: int,
        seed_dim: int = 16,
        n_harmonics: int = 8
    ):
        super().__init__()
        self.fc1 = NativeHelixLinear(in_features, hidden_features, seed_dim, n_harmonics=n_harmonics)
        self.fc2 = NativeHelixLinear(hidden_features, out_features, seed_dim, n_harmonics=n_harmonics)
        self.relu = nn.ReLU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

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

    @property
    def compression_ratio(self) -> float:
        dense_params = (
            self.fc1.in_features * self.fc1.out_features + self.fc1.out_features +
            self.fc2.in_features * self.fc2.out_features + self.fc2.out_features
        )
        return dense_params / self.total_params


class DenseBaseline(nn.Module):
    """Dense MLP baseline for comparison."""

    def __init__(self, in_features: int, hidden_features: int, out_features: int):
        super().__init__()
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.relu = nn.ReLU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

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


def verify_gradients():
    """Verify that gradients flow through decode_phi_torch."""
    print("=" * 60)
    print("GRADIENT VERIFICATION TEST")
    print("=" * 60)

    # Create a simple test
    seeds = torch.randn(4, 8, requires_grad=True)
    W = decode_phi_torch(seeds, 16)

    # Compute a scalar loss
    loss = W.sum()

    # Backprop
    loss.backward()

    # Check gradients exist
    assert seeds.grad is not None, "FAIL: No gradients on seeds!"
    assert seeds.grad.abs().sum() > 0, "FAIL: Gradients are all zeros!"

    print(f"  Seeds shape: {seeds.shape}")
    print(f"  Weights shape: {W.shape}")
    print(f"  Grad shape: {seeds.grad.shape}")
    print(f"  Grad norm: {seeds.grad.norm().item():.6f}")
    print("  PASS: Gradients flow correctly!")
    print()

    return True


def demo():
    """Quick demo of NativeHelixLinear."""
    print("=" * 60)
    print("NATIVE HELIX LINEAR DEMO")
    print("=" * 60)

    # Verify gradients first
    verify_gradients()

    # Compare parameter counts
    in_f, hidden, out_f = 784, 128, 10
    seed_dim = 16

    native = NativeHelixMLP(in_f, hidden, out_f, seed_dim=seed_dim)
    dense = DenseBaseline(in_f, hidden, out_f)

    print("PARAMETER COMPARISON:")
    print(f"  Dense MLP:   {dense.total_params:,} params")
    print(f"  Native MLP:  {native.total_params:,} params")
    print(f"  Compression: {native.compression_ratio:.1f}:1")
    print()

    # Forward pass test
    x = torch.randn(32, in_f)
    y_native = native(x)
    y_dense = dense(x)

    print("FORWARD PASS:")
    print(f"  Input shape:  {x.shape}")
    print(f"  Native out:   {y_native.shape}")
    print(f"  Dense out:    {y_dense.shape}")
    print()

    # Gradient test on native
    loss = y_native.sum()
    loss.backward()

    has_grads = all(p.grad is not None and p.grad.abs().sum() > 0 for p in native.parameters())
    print(f"  Gradients on all native params: {'PASS' if has_grads else 'FAIL'}")
    print()

    print("=" * 60)
    print("DEMO COMPLETE - Ready for training experiments")
    print("=" * 60)


if __name__ == "__main__":
    demo()
