import torch
import torch.nn as nn
import torch.nn.functional as F


__all__ = [
    "FeatureAttention",
    "TimeAttention",
    "PositionEncodeLearnable",
    "ConvGRUTransformerHLV10"
]


class FeatureAttention(nn.Module):
    def __init__(self, feature_dim):
        super().__init__()
        self.attn = nn.Linear(feature_dim, feature_dim)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x):
        # x: [batch, seq_len, features]
        weights = self.softmax(self.attn(x))  # [batch, seq_len, features]
        out = x * weights
        return out.sum(dim=-1)  # [batch, seq_len]

class TimeAttention(nn.Module):
    def __init__(self, seq_len):
        super().__init__()
        self.attn = nn.Linear(seq_len, seq_len)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x):
        # x: [batch, seq_len, features]
        x_t = x.transpose(1, 2)  # [batch, features, seq_len]
        
        # ต้อง reshape ก่อนเข้า linear layer
        batch_size, features, seq_len = x_t.shape
        x_flat = x_t.reshape(batch_size * features, seq_len)  # [batch*features, seq_len]
        
        weights = self.softmax(self.attn(x_flat))  # [batch*features, seq_len]
        weights = weights.reshape(batch_size, features, seq_len)  # [batch, features, seq_len]
        
        out = x_t * weights
        return out.sum(dim=-1)  # [batch, features]

class PositionEncodeLearnable(nn.Module):
    def __init__(self, seq_len, d_model):
        super().__init__()
        self.pos_embedding = nn.Parameter(torch.randn(1, seq_len, d_model))
    
    def forward(self, x):
        return x + self.pos_embedding

class ConvGRUTransformerHLV10(nn.Module):
    def __init__(self, 
                 input_dim,             
                 conv_channels=32,      
                 gru_hidden=64,         
                 nhead=4,               
                 num_encoder_layers=1,  
                 dim_feedforward=128, 
                 seq_len=3,             
                 kernel_size=3, 
                 dropout=0.1, 
                 output_steps=1
                 ):
        super().__init__()
        self.seq_len = seq_len
        self.kernel_size = kernel_size
        self.output_steps = output_steps
        
        # --- Attention Modules ---
        self.feature_attn = FeatureAttention(input_dim)
        
        
        # --- Conv1D Layer (Causal) ---
        self.conv1 = nn.Conv1d(
            in_channels=input_dim, 
            out_channels=conv_channels, 
            kernel_size=self.kernel_size, 
            dilation=1,   
            padding=0     
        )
        self.conv_bn = nn.BatchNorm1d(conv_channels)
        
        # --- GRU Layer ---
        self.gru = nn.GRU(
            input_size=conv_channels, 
            hidden_size=gru_hidden, 
            batch_first=True
        )

        self.time_attn = TimeAttention(seq_len)
        
        # --- Learnable Positional Encoding ---
        self.pos_encoder = PositionEncodeLearnable(seq_len, gru_hidden)
        
        # --- Transformer Encoder ---
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=gru_hidden, 
            nhead=nhead, 
            dim_feedforward=dim_feedforward, 
            dropout=dropout, 
            batch_first=True,
            activation="gelu"
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
        
        # --- Output Head ---
        self.fc_high = nn.Linear(gru_hidden, output_steps) 
        self.fc_low  = nn.Linear(gru_hidden, output_steps)
    
    def forward(self, x):
        # x: (batch, seq_len, features)
        batch_size, seq_len, features = x.size()
        
        # --- Feature-wise Attention ก่อน Conv ---
        feat_attn = self.feature_attn(x)  # (batch, seq_len)
        x = x * feat_attn.unsqueeze(-1)    # apply attention weight
        
        # --- Conv1D ---
        x_conv = x.permute(0, 2, 1)  # (batch, features, seq_len)
        pad = self.kernel_size - 1
        x_conv = F.pad(x_conv, (pad, 0))  
        x_conv = self.conv1(x_conv)  
        x_conv = self.conv_bn(x_conv)
        x_conv = torch.tanh(x_conv)  
        x_conv = x_conv[:, :, -seq_len:]
        x_conv = x_conv.permute(0, 2, 1)  # (batch, seq_len, conv_channels)
        
        # --- GRU ---
        gru_out, _ = self.gru(x_conv)
        if gru_out.size(1) != self.seq_len:
            current_seq_len = gru_out.size(1)
            if current_seq_len < self.seq_len:
                pad_size = self.seq_len - current_seq_len
                gru_out = F.pad(gru_out, (0, 0, 0, pad_size))
            else:
                gru_out = gru_out[:, :self.seq_len, :]
        
        # --- Time-wise Attention หลัง GRU ---
        time_attn_out = self.time_attn(gru_out)  # (batch, gru_hidden)
        x_attn = time_attn_out.unsqueeze(1).repeat(1, self.seq_len, 1)
        
        # --- Positional Encoding ---
        x_attn = self.pos_encoder(x_attn)
        
        # --- Transformer Encoder ---
        trans_out = self.transformer(x_attn)  
        
        # --- Output ---
        last_steps = trans_out[:, -self.output_steps:, :]  
        high_out = self.fc_high(last_steps)  
        low_out = self.fc_low(last_steps)    
        
        return high_out.squeeze(-1), low_out.squeeze(-1)