import numpy as np

try:
    import torch
    import torch.nn as nn
    import torch.optim as optim
    from torch.utils.data import DataLoader, TensorDataset
    HAS_TORCH = True
    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
except ImportError:
    HAS_TORCH = False
    DEVICE = None

class TimeSeriesTransformer(nn.Module if HAS_TORCH else object):
    """
    Advanced Multi-Head Attention Transformer for Deep Sequence Forecasting.
    Designed to catch non-linear microscopic patterns in RNG outputs.
    """
    def __init__(self, input_dim=1, d_model=64, nhead=4, num_layers=3, dropout=0.1):
        if not HAS_TORCH:
            raise ImportError("PyTorch is required for the Transformer module.")
        super().__init__()
        self.input_projection = nn.Linear(input_dim, d_model)
        
        # Positional Encoding
        self.pos_encoder = PositionalEncoding(d_model, dropout)
        
        encoder_layers = nn.TransformerEncoderLayer(d_model, nhead, d_model*4, dropout, batch_first=True)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layers, num_layers)
        
        self.decoder = nn.Sequential(
            nn.Linear(d_model, d_model // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_model // 2, 1)
        )

    def forward(self, src):
        # src shape: [batch_size, sequence_length, input_dim]
        src = self.input_projection(src)
        src = self.pos_encoder(src)
        output = self.transformer_encoder(src)
        # Take the last sequence output to predict the next step
        last_out = output[:, -1, :]
        return self.decoder(last_out)

class PositionalEncoding(nn.Module if HAS_TORCH else object):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        if not HAS_TORCH: return
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)

        position = torch.arange(max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) * (-np.log(10000.0) / d_model))
        pe = torch.zeros(max_len, 1, d_model)
        pe[:, 0, 0::2] = torch.sin(position * div_term)
        pe[:, 0, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe)

    def forward(self, x):
        """
        Arguments:
            x: Tensor, shape ``[batch_size, seq_len, embedding_dim]``
        """
        x = x + self.pe[:x.size(1)].transpose(0, 1)
        return self.dropout(x)

class TransformerPredictor:
    """Wrapper to easily train and infer using the Transformer."""
    def __init__(self, sequence_length=20, epochs=50, lr=0.001):
        self.sequence_length = sequence_length
        self.epochs = epochs
        self.lr = lr
        self.model = None
        self.is_trained = False
        
        if HAS_TORCH:
            self.model = TimeSeriesTransformer(input_dim=1).to(DEVICE)
            self.optimizer = optim.Adam(self.model.parameters(), lr=self.lr)
            self.criterion = nn.MSELoss()

    def _prepare_data(self, history):
        if len(history) <= self.sequence_length:
            return None, None
            
        X, y = [], []
        for i in range(len(history) - self.sequence_length):
            raw_x = history[i:i+self.sequence_length]
            raw_y = history[i+self.sequence_length]
            
            flat_x = [float(np.ravel(val)[0]) for val in raw_x]
            flat_y = float(np.ravel(raw_y)[0])
            
            X.append(flat_x)
            y.append(flat_y)
            
        # Normalize to 0-1 for better gradients
        X = np.array(X, dtype=np.float32) / 100.0
        y = np.array(y, dtype=np.float32) / 100.0
        
        X = torch.tensor(X).unsqueeze(-1).to(DEVICE) # [Batch, Seq, 1]
        y = torch.tensor(y).unsqueeze(-1).to(DEVICE) # [Batch, 1]
        return X, y

    def train(self, history):
        if not HAS_TORCH or len(history) < self.sequence_length * 2:
            return False
            
        X, y = self._prepare_data(history)
        if X is None: return False
        
        dataset = TensorDataset(X, y)
        loader = DataLoader(dataset, batch_size=32, shuffle=True)
        
        self.model.train()
        for epoch in range(self.epochs):
            total_loss = 0
            for batch_x, batch_y in loader:
                self.optimizer.zero_grad()
                output = self.model(batch_x)
                loss = self.criterion(output, batch_y)
                loss.backward()
                self.optimizer.step()
                total_loss += loss.item()
                
        self.is_trained = True
        return True
        
    def predict(self, recent_history):
        if not self.is_trained or not HAS_TORCH:
            return None
        if len(recent_history) < self.sequence_length:
            return None
            
        # Take the most recent sequences and flatten them in case they are nested
        raw_seq = recent_history[-self.sequence_length:]
        flat_seq = [float(np.ravel(x)[0]) for x in raw_seq]
        seq = np.array(flat_seq, dtype=np.float32) / 100.0
        seq = torch.tensor(seq).unsqueeze(0).unsqueeze(-1).to(DEVICE) # [1, Seq, 1]
        
        self.model.eval()
        with torch.no_grad():
            pred = self.model(seq)
            
        # Denormalize
        return float(pred.item() * 100.0)
