#!/usr/bin/env python3
"""
HelixLobe - Native Helix residual module for existing linear layers.

Used to grow new capabilities on top of a frozen base model
(e.g., TinyLlama GENOME₀) with very few additional parameters.

Core idea:
    y = W_base(x) + alpha * W_helix(x)

Where:
    - W_base: frozen dense weights from compressed/regenerated model
    - W_helix: NativeHelixLinear using φ/FibPi seeds as parameters
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Optional, Dict, Any

import torch
import torch.nn as nn

import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from helix_cdc.native_helix_linear import NativeHelixLinear


@dataclass
class HelixLobeConfig:
    seed_dim: int = 64          # dimension of seed vector for NativeHelixLinear
    n_harmonics: int = 8        # number of phi harmonics
    use_gate: bool = True       # whether to use a trainable scalar gate α
    init_gate: float = 0.1      # initial α value
    train_base: bool = False    # normally False: base frozen
    name: str = "helix_lobe"    # lobe name for receipts


class HelixLobe(nn.Module):
    """
    Wraps an existing nn.Linear with a Helix-native residual.

    Example:
        base_linear = model.layers[3].mlp.down_proj
        lobe = HelixLobe.from_linear(base_linear, HelixLobeConfig(seed_dim=64))
        model.layers[3].mlp.down_proj = lobe
    """

    def __init__(
        self,
        base_linear: nn.Linear,
        config: HelixLobeConfig,
    ):
        super().__init__()
        self.config = config

        # Keep a reference to the base linear layer
        self.base = base_linear

        self.in_features = base_linear.in_features
        self.out_features = base_linear.out_features

        # Optionally freeze base params
        if not config.train_base:
            for p in self.base.parameters():
                p.requires_grad = False

        # Helix-native residual
        self.helix = NativeHelixLinear(
            in_features=self.in_features,
            out_features=self.out_features,
            seed_dim=config.seed_dim,
            bias=False,
            n_harmonics=config.n_harmonics,
        )

        # Optional scalar gate α
        if config.use_gate:
            self.alpha = nn.Parameter(
                torch.tensor(config.init_gate, dtype=torch.float32)
            )
        else:
            self.register_buffer(
                "alpha", torch.tensor(1.0, dtype=torch.float32)
            )

        # Cache param counts
        self._dense_params = sum(p.numel() for p in self.base.parameters())
        self._helix_params = sum(p.numel() for p in self.helix.parameters())
        if config.use_gate:
            self._helix_params += 1  # alpha

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        base_out = self.base(x)
        lobe_out = self.helix(x)
        return base_out + self.alpha * lobe_out

    def trainable_parameters(self):
        """Yield only trainable parameters (helix seeds + alpha)."""
        yield from self.helix.parameters()
        if self.config.use_gate:
            yield self.alpha

    def extra_state_dict(self) -> Dict[str, Any]:
        """
        Extra metadata for receipts (config, compression info, etc.).
        """
        compression = self._dense_params / max(self._helix_params, 1)

        return {
            "config": self.config.__dict__,
            "in_features": self.in_features,
            "out_features": self.out_features,
            "dense_params": self._dense_params,
            "helix_params": self._helix_params,
            "compression": compression,
            "alpha": float(self.alpha.item()),
        }

    @property
    def compression(self) -> float:
        """Compression ratio vs dense adapter."""
        dense_adapter = self.in_features * self.out_features
        return dense_adapter / max(self._helix_params, 1)

    @classmethod
    def from_linear(
        cls,
        base_linear: nn.Linear,
        config: Optional[HelixLobeConfig] = None,
    ) -> "HelixLobe":
        if config is None:
            config = HelixLobeConfig()
        return cls(base_linear, config)

    def extra_repr(self) -> str:
        return (
            f"in={self.in_features}, out={self.out_features}, "
            f"seed_dim={self.config.seed_dim}, "
            f"alpha={self.alpha.item():.4f}, "
            f"compression={self.compression:.1f}:1"
        )


def get_module_by_name(root: nn.Module, name: str) -> nn.Module:
    """Get a submodule by dot-separated path."""
    parts = name.split(".")
    module = root
    for p in parts:
        if p.isdigit():
            module = module[int(p)]
        else:
            module = getattr(module, p)
    return module


def set_module_by_name(root: nn.Module, name: str, new_module: nn.Module) -> None:
    """Set a submodule by dot-separated path."""
    parts = name.split(".")
    module = root
    for p in parts[:-1]:
        if p.isdigit():
            module = module[int(p)]
        else:
            module = getattr(module, p)
    last = parts[-1]
    if last.isdigit():
        module[int(last)] = new_module
    else:
        setattr(module, last, new_module)


def demo():
    """Quick demo of HelixLobe."""
    print("=" * 70)
    print("HELIX LOBE DEMO")
    print("=" * 70)

    # Create a base linear
    base = nn.Linear(768, 768)
    base_params = sum(p.numel() for p in base.parameters())
    print(f"\nBase linear: 768x768 = {base_params:,} params")

    # Wrap with HelixLobe
    config = HelixLobeConfig(seed_dim=16, n_harmonics=8)
    lobe = HelixLobe.from_linear(base, config)

    print(f"\nHelixLobe stats:")
    for k, v in lobe.extra_state_dict().items():
        if isinstance(v, float):
            print(f"  {k}: {v:.4f}")
        elif isinstance(v, dict):
            print(f"  {k}: {v}")
        else:
            print(f"  {k}: {v}")

    # Forward pass test
    x = torch.randn(2, 128, 768)
    y = lobe(x)
    print(f"\nForward pass:")
    print(f"  Input:  {x.shape}")
    print(f"  Output: {y.shape}")

    # Gradient test
    loss = y.sum()
    loss.backward()

    # Check base is frozen
    base_grads = [p.grad for p in lobe.base.parameters() if p.grad is not None]
    helix_grads = [p.grad for p in lobe.helix.parameters() if p.grad is not None]

    print(f"\nGradient check:")
    print(f"  Base grads:  {len(base_grads)} (should be 0 - frozen)")
    print(f"  Helix grads: {len(helix_grads)} (should be >0 - trainable)")
    print(f"  Alpha grad:  {lobe.alpha.grad is not None}")

    # Trainable params count
    trainable = sum(p.numel() for p in lobe.trainable_parameters())
    print(f"\nTrainable params: {trainable:,}")
    print(f"Compression vs dense adapter: {lobe.compression:.1f}:1")

    print("\n" + "=" * 70)
    print("HELIX LOBE DEMO COMPLETE")
    print("=" * 70)


if __name__ == "__main__":
    demo()
