import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR

class HolographicMasterCodeTransformer(nn.Module):
    def __init__(self, input_dim=8, d_model=128, nhead=8, num_layers=6, output_dim=1, n_harmonics=4):
        super().__init__()
        self.d_model = d_model
        self.n_harmonics = n_harmonics
        self.alpha_0 = torch.tensor(1.0 / 137.035) # Convert to tensor during initialization
        
        self.input_proj = nn.Linear(input_dim, d_model)
        
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead, dim_feedforward=d_model*4,
            dropout=0.1, activation='gelu', batch_first=True, norm_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        
        self.harmonic_weights = nn.Parameter(torch.randn(n_harmonics, d_model))
        
        self.output_head = nn.Sequential(
            nn.LayerNorm(d_model),
            nn.Linear(d_model, d_model//2),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(d_model//2, output_dim)
        )
        
    def forward(self, x, alpha=None):
        # Ensure alpha_0 is on the correct device
        if self.alpha_0.device != x.device:
            self.alpha_0 = self.alpha_0.to(x.device)

        if alpha is None:
            alpha = self.alpha_0
        else:
            # Convert alpha to a tensor if it's not already, and move to the correct device
            alpha = torch.tensor(alpha, device=x.device, dtype=x.dtype) if not isinstance(alpha, torch.Tensor) else alpha.to(x.device)
            
        h = self.input_proj(x)
        h = h.unsqueeze(1)
        h = self.transformer(h)
        h = h.squeeze(1)
        
        # All operands for arithmetic operations should be tensors now
        alpha_diff = alpha - self.alpha_0
        
        psi = torch.zeros_like(h)
        for n in range(self.n_harmonics):
            # Ensure phase calculation uses tensors
            phase = 2 * torch.pi * (n + 1) * alpha_diff * 800.0
            psi += self.harmonic_weights[n] * torch.cos(h * (n + 1) + phase)
        
        alpha_dev = torch.abs(alpha_diff) # Now alpha_diff is guaranteed to be a tensor
        security_factor = torch.exp(-300.0 * alpha_dev)
        
        h = h + 0.35 * psi * security_factor
        return self.output_head(h)
    
    def holographic_reg(self):
        reg = sum(torch.norm(w, p=1) for w in self.harmonic_weights)
        return 0.0015 * reg / self.n_harmonics