import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader, Subset
from datetime import datetime
import pandas as pd
#import tensorflow as tf
import numpy as np
from datetime import datetime, timedelta
import joblib,random,os
from collections import deque


__all__ = [
    "PositionalEncoding",
    "LearnablePositionalEncoding",
    "TransformerSingleStep",
    "TransformerMultiStep",
    "EarlyStopping",
    "SinglestepDataset",
    "ReplayBuffer",
    "CalibratorSingleStep",
    "CalibratorSingleStepV3",
    "CalibratorMultiStep",
    "MultiStepDataset",
    "TransformerMultistepMonth",
    "HybridTransformerModelV4",
    "MultiTimeframeTransformer",
    "TransformerSingleStepForEA",
    "ConvGRUTransformerHL",
    "Seq2SeqAttentionHL",
    "Seq2SeqAttentionHLWithConv",
    "HybridMultiBranchMultiStepAttention",
    "HybridMultiBranchMultiStepAttention",
    "LRFinderWrapper",
    "ConvGRUTransformerHLV2",
    "RobustScalerSimple",
    "MinMaxScaler"
]

#+-------------------------------------------------------------------+
#| Use in Training loop                                              |
#+-------------------------------------------------------------------+

class RobustScalerSimple:
    def __init__(self, quantile_range=(25,75)):
        self.quantile_range = quantile_range
        self.median_ = None
        self.iqr_ = None

    def fit(self, X):
        q_min, q_max = self.quantile_range
        self.median_ = np.median(X, axis=0)
        q1 = np.percentile(X, q_min, axis=0)
        q3 = np.percentile(X, q_max, axis=0)
        self.iqr_ = q3 - q1
        # ป้องกันหารด้วย 0
        self.iqr_[self.iqr_ == 0] = 1.0
        return self

    def transform(self, X):
        return (X - self.median_) / self.iqr_

    def fit_transform(self, X):
        self.fit(X)
        return self.transform(X)

    def inverse_transform(self, X_scaled):
        return X_scaled * self.iqr_ + self.median_

class MinMaxScaler:
    def __init__(self, feature_range=(-1,1)):
        self.min_val, self.max_val = feature_range
        self.min_ = None
        self.max_ = None
        self.scale_ = None

    def fit(self, X):
        self.min_ = X.min(axis=0)
        self.max_ = X.max(axis=0)
        self.scale_ = (self.max_val - self.min_val) / (self.max_ - self.min_ + 1e-8)
        return self

    def transform(self, X):
        return (X - self.min_) * self.scale_ + self.min_val

    def fit_transform(self, X):
        self.fit(X)
        return self.transform(X)

    def inverse_transform(self, X_scaled):
        return (X_scaled - self.min_val) / self.scale_ + self.min_
    
class EarlyStopping:
    def __init__(self, patience=7, delta=0, verbose=False):
        self.patience = patience
        self.delta = delta
        self.verbose = verbose
        self.counter = 0
        self.best_score = None
        self.early_stop = False
        self.val_loss_min = np.Inf
        
    def __call__(self, val_loss, model, path):
        score = -val_loss
        
        if self.best_score is None:
            self.best_score = score
            self.save_checkpoint(val_loss, model, path)
        elif score < self.best_score + self.delta:
            self.counter += 1
            print(f'EarlyStopping counter: {self.counter} out of {self.patience}')
            if self.counter >= self.patience:
                self.early_stop = True
        else:
            self.best_score = score
            self.save_checkpoint(val_loss, model, path)
            self.counter = 0
            
    def save_checkpoint(self, val_loss, model, path):
        if self.verbose:
            print(f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}). Saving model...')
        torch.save(model.state_dict(), path)
        self.val_loss_min = val_loss
    
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        pos = torch.arange(0, max_len).unsqueeze(1)
        div = torch.exp(torch.arange(0, d_model, 2) * -(np.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(pos * div)
        pe[:, 1::2] = torch.cos(pos * div)
        self.pe = pe.unsqueeze(0)
    def forward(self, x):
        x = x + self.pe[:, :x.size(1)].to(x.device)
        return x

#+-------------------------------------------------------------------+
#| 24082025 Inherit from nn.Module                                   |
#| Work with Transformer                                             |
#+-------------------------------------------------------------------+
class LearnablePositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=10000000):
        super().__init__()
        self.pos_embedding = nn.Embedding(max_len, d_model)
        self._init_weights()
        
    def _init_weights(self):
        #  initialization ที่เหมาะสมสำหรับ positional encoding
        nn.init.uniform_(self.pos_embedding.weight, -0.1, 0.1)
        
    def forward(self, x):
        # x.shape = [batch_size, seq_len, d_model]
        seq_len = x.size(1)
        pos = torch.arange(0, seq_len, device=x.device).unsqueeze(0)  # [1, seq_len]
        pos_embed = self.pos_embedding(pos)  # [1, seq_len, d_model]
        return x + pos_embed

class TransformerSingleStep(nn.Module):
    def __init__(self, input_size, d_model=64, nhead=4, num_layers=2, dropout=0.1, max_len=500):
        super().__init__()
        self.input_fc = nn.Linear(input_size, d_model)
        self.tanh = nn.Tanh()
        self.pos_encoder = LearnablePositionalEncoding(d_model, max_len=max_len)
        
        encoder_layer = nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward=128, 
                                                   dropout=dropout, batch_first=True)
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)
        self.output = nn.Linear(d_model, 2)

    def forward(self, x):
        # x.shape = [batch_size, seq_len, input_size]
        x = self.input_fc(x)
        x = self.tanh(x)
        x = self.pos_encoder(x)
        x = self.transformer(x)  # [B, seq_len, d_model]
        
        # ใช้เฉพาะ output ของ time step สุดท้าย
        last_output = x[:, -1, :]  # [B, d_model]
        
        # output layer แบ่ง high/low
        out = self.output(last_output)  # [B, 2] → col0=high, col1=low
        pred_high = out[:, 0].unsqueeze(1)  # [B,1]
        pred_low  = out[:, 1].unsqueeze(1)  # [B,1]
        
        return pred_high, pred_low

class TransformerMultiStep(nn.Module):
    def __init__(self, input_size, d_model=64, nhead=4, num_layers=2,
                 dropout=0.1, horizon=5, output_dropout=0.2, max_seq_len=12):
        super().__init__()
        self.horizon = horizon
        self.d_model = d_model
        
        # Input projection
        self.input_fc = nn.Linear(input_size, d_model)
        self.tanh = nn.Tanh()
        
        # Learnable positional encoding
        self.pos_encoding = nn.Parameter(torch.zeros(1, max_seq_len, d_model))
        nn.init.normal_(self.pos_encoding, mean=0.0, std=0.02)  # นี้ถูกต้องแล้ว
        
        # Transformer with causal masking
        encoder_layer = nn.TransformerEncoderLayer(
            d_model, nhead, dim_feedforward=128,
            dropout=dropout, batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)
        
        # Output layers
        self.output_high = nn.Sequential(
            nn.Linear(d_model, d_model//2),
            nn.Tanh(),
            nn.Dropout(output_dropout),
            nn.Linear(d_model//2, horizon)
        )
        
        self.output_low = nn.Sequential(
            nn.Linear(d_model, d_model//2),
            nn.Tanh(),
            nn.Dropout(output_dropout),
            nn.Linear(d_model//2, horizon)
        )

    def forward(self, x):
        batch_size, seq_len = x.size(0), x.size(1)
        
        x = self.input_fc(x)
        x = self.tanh(x)
        
        # Add positional encoding
        x = x + self.pos_encoding[:, :seq_len, :]
        
        # Causal mask
        mask = torch.triu(torch.ones(seq_len, seq_len) * float('-inf'), diagonal=1).to(x.device)
        
        x = self.transformer(x, mask=mask)
        last_output = x[:, -1, :]
        
        high_pred = self.output_high(last_output)
        low_pred = self.output_low(last_output)
        
        return high_pred, low_pred

#+-------------------------------------------------------------------+
#| 24082025                                                          |
#| Work with Online update                                           |
#+-------------------------------------------------------------------+
class ReplayBuffer:
    def __init__(self, capacity=200):
        self.capacity = capacity
        self.buffer = []

    def push(self, x, y):
        if len(self.buffer) >= self.capacity:
            self.buffer.pop(0)
        self.buffer.append((x.detach().clone(), y.detach().clone()))

    def sample(self, batch_size=5):
        if len(self.buffer) == 0:
            return None, None
        sample = random.sample(self.buffer, min(batch_size, len(self.buffer)))
        xs, ys = zip(*sample)
        return torch.stack(xs), torch.stack(ys)

#+-------------------------------------------------------------------+
#| 24082025                                                          |
#| Work with predict process and update predict price near real price|
#+-------------------------------------------------------------------+
class CalibratorSingleStep:
    def __init__(self, method='EMA', alpha=0.3):
        """
        method: 'EMA' หรือ 'linear'
        alpha: smoothing factor สำหรับ EMA
        """
        self.method = method
        self.alpha = alpha
        self.history_pred_high = []
        self.history_pred_low  = []
        self.history_true_high = []
        self.history_true_low  = []
        self.bias_high = 0.0
        self.bias_low  = 0.0
        self.model_high = None
        self.model_low  = None

    def update_history(self, pred_high, pred_low, true_high, true_low):
        self.history_pred_high.append(pred_high)
        self.history_pred_low.append(pred_low)
        self.history_true_high.append(true_high)
        self.history_true_low.append(true_low)

    def refit(self):
        if self.method == 'linear':
            if len(self.history_pred_high) >= 2:
                Xh = np.array(self.history_pred_high).reshape(-1,1)
                Yh = np.array(self.history_true_high)
                self.model_high = LinearRegression().fit(Xh, Yh)
            if len(self.history_pred_low) >= 2:
                Xl = np.array(self.history_pred_low).reshape(-1,1)
                Yl = np.array(self.history_true_low)
                self.model_low = LinearRegression().fit(Xl, Yl)
        # สำหรับ EMA ไม่ต้อง refit

    def calibrate(self, pred_high, pred_low):
        if self.method == 'EMA':
            ch = pred_high + self.alpha*(self.history_true_high[-1]-pred_high) if self.history_true_high else pred_high
            cl = pred_low  + self.alpha*(self.history_true_low[-1]-pred_low) if self.history_true_low else pred_low
            ch += self.bias_high
            cl += self.bias_low
        elif self.method == 'linear':
            ch = self.model_high.predict(np.array([[pred_high]]))[0] + self.bias_high if self.model_high else pred_high
            cl = self.model_low.predict(np.array([[pred_low]]))[0] + self.bias_low if self.model_low else pred_low
        else:
            ch, cl = pred_high, pred_low
        return ch, cl

    def update_bias(self, pred_high_calib, pred_low_calib, true_high, true_low):
        # update bias ด้วยค่า error ล่าสุด
        self.bias_high = true_high - pred_high_calib
        self.bias_low  = true_low  - pred_low_calib

#+-------------------------------------------------------------------+
#| 24082025                                                          |
#| Work with predict process and update predict price near real price|
#+-------------------------------------------------------------------+
class CalibratorSingleStepV3:
    """
    Bias calibration สำหรับ Online Update V3
    """
    def __init__(self, method='EMA', alpha=0.1):
        self.method = method
        self.alpha = alpha
        self.high_history = []
        self.low_history = []
        self.bias_high = 0.0
        self.bias_low = 0.0

    def update_history(self, pred_high, pred_low, true_high, true_low):
        self.high_history.append(true_high - pred_high)
        self.low_history.append(true_low - pred_low)

    def refit(self):
        if self.high_history:
            if self.method == 'EMA':
                self.bias_high = self.alpha * self.high_history[-1] + (1 - self.alpha) * self.bias_high
                self.bias_low  = self.alpha * self.low_history[-1]  + (1 - self.alpha) * self.bias_low
            elif self.method == 'mean':
                self.bias_high = sum(self.high_history)/len(self.high_history)
                self.bias_low  = sum(self.low_history)/len(self.low_history)

    def calibrate(self, pred_high, pred_low):
        return pred_high + self.bias_high, pred_low + self.bias_low

    def update_bias(self, pred_high_calib, pred_low_calib, true_high, true_low):
        self.update_history(pred_high_calib, pred_low_calib, true_high, true_low)
        self.refit()

#+-------------------------------------------------------------------+
#| 24082025 Inherit from Dataset                                     |
#| get Scaled data and create rolling window before use loader       |
#+-------------------------------------------------------------------+
class SinglestepDataset(Dataset):
    def __init__(self, X, y_high, y_low, lookback=4):
        """
        Dataset สำหรับ Time Series Forecasting แบบ Single-step (Rolling window)
        """
        self.X = torch.tensor(X, dtype=torch.float32)
        self.y_high = torch.tensor(y_high, dtype=torch.float32).squeeze(-1)  # ลบ dimension เกิน
        self.y_low = torch.tensor(y_low, dtype=torch.float32).squeeze(-1)
        self.lookback = lookback

        # index สุดท้ายที่ยังสร้าง sample ได้
        self.max_idx = len(self.X) - lookback

    def __len__(self):
        return self.max_idx

    def __getitem__(self, idx):
        x_seq = self.X[idx: idx + self.lookback]  # (lookback, features)
        y_h = self.y_high[idx + self.lookback].unsqueeze(0)  # shape: (1,)
        y_l = self.y_low[idx + self.lookback].unsqueeze(0)
        return x_seq, y_h, y_l

#+-------------------------------------------------------------------+
#| 24082025 Inherit from Dataset                                     |
#| get Scaled data and create rolling window before use loader       |
#+-------------------------------------------------------------------+
class MultiStepDataset(Dataset):
    def __init__(self, X, y_high, y_low, lookback=4, horizon=5):
        """
        X       : np.ndarray or torch.Tensor, shape (N, features)
        y_high  : np.ndarray or torch.Tensor, shape (N,)
        y_low   : np.ndarray or torch.Tensor, shape (N,)
        lookback: จำนวน step ที่ใช้เป็น input
        horizon : จำนวน step ข้างหน้าที่ต้องการทำนาย
        """
        self.X = torch.tensor(X, dtype=torch.float32)
        self.y_high = torch.tensor(y_high, dtype=torch.float32)
        self.y_low = torch.tensor(y_low, dtype=torch.float32)
        self.lookback = lookback
        self.horizon = horizon

        # index สุดท้ายที่ยังสร้าง sample ได้
        self.max_idx = len(self.X) - lookback - horizon + 1

    def __len__(self):
        return self.max_idx

    # Crete rolling window
    def __getitem__(self, idx):
        x_seq = self.X[idx: idx + self.lookback]   # (lookback, features)
        y_h_seq = self.y_high[idx + self.lookback : idx + self.lookback + self.horizon].squeeze(-1)  # (horizon,)
        y_l_seq = self.y_low[idx + self.lookback : idx + self.lookback + self.horizon].squeeze(-1)   # (horizon,)
        return x_seq, y_h_seq, y_l_seq
    
class CalibratorMultiStep:
    """
    Calibrator สำหรับ multistep/recursive prediction
    - เก็บ history ของ high/low
    - ปรับ bias ให้ step แรก match last_close
    - สามารถ smooth prediction ด้วย Gaussian filter
    """
    def __init__(self, last_close=None, max_history=50, sigma=1.5):
        self.last_close = last_close
        self.history_high = deque(maxlen=max_history)
        self.history_low  = deque(maxlen=max_history)
        self.sigma = sigma
    
    def update_history(self, raw_high, raw_low, calibrated_high=None, calibrated_low=None):
        """
        เก็บ history ของ step ล่าสุด
        - calibrated_high/low: optional สำหรับเก็บ version ปรับแล้ว
        """
        self.history_high.append(calibrated_high if calibrated_high is not None else raw_high)
        self.history_low.append(calibrated_low if calibrated_low is not None else raw_low)
    
    def refit(self):
        """
        ปรับค่าภายในถ้าจำเป็น (placeholder)
        - future: สามารถปรับ weight history, scale หรือ smoothing factor
        """
        pass
    
    def calibrate(self, raw_high, raw_low):
        """
        ปรับค่า bias ให้ step แรกตรง last_close
        - สำหรับ step ต่อไปจะ smooth ด้วย history
        """
        h, l = raw_high, raw_low
        
        # Step แรกของ fold → match last_close
        if self.last_close is not None and len(self.history_high) == 0:
            delta_h = h - self.last_close
            delta_l = l - self.last_close
            h -= delta_h
            l -= delta_l
        
        # Smooth ด้วย history (ถ้ามีมากกว่า 2 step)
        if len(self.history_high) > 1:
            h_array = np.array(list(self.history_high) + [h])
            l_array = np.array(list(self.history_low) + [l])
            h_smooth = gaussian_filter1d(h_array, sigma=self.sigma)[-1]
            l_smooth = gaussian_filter1d(l_array, sigma=self.sigma)[-1]
            h, l = h_smooth, l_smooth
        
        return h, l
    
class TransformerMultistepMonth(nn.Module):
    def __init__(self, input_size, d_model=64, nhead=4, num_layers=2,
                 dropout=0.1, horizon=4, conv_dropout=0.1, conv_kernel=2,
                 post_trans_dropout=0.1, output_dropout=0.1, dim_feedforward=128):
        super().__init__()
        self.horizon = horizon
        
        # Causal padding for conv1d
        self.causal_pad = nn.ConstantPad1d((conv_kernel - 1, 0), 0)
        self.conv1d = nn.Conv1d(in_channels=input_size, out_channels=d_model, kernel_size=conv_kernel, padding=0)
        self.tanh = nn.Tanh()
        self.norm1 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(conv_dropout)
        
        # Positional encoding
        self.pos_encoder = LearnablePositionalEncoding(d_model, max_len=200)
        
        # Transformer encoder
        encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead,
                                                   dim_feedforward=dim_feedforward,
                                                   dropout=dropout, batch_first=True)
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.post_trans_dropout = nn.Dropout(post_trans_dropout)
        
        # Output linear layers (sequence-wise)
        self.output_dropout_high = nn.Dropout(output_dropout)
        self.output_dropout_low = nn.Dropout(output_dropout)
        self.output_high = nn.Linear(d_model, horizon)
        self.output_low = nn.Linear(d_model, horizon)

    def forward(self, x):
        x = x.permute(0, 2, 1)
        x = self.causal_pad(x)
        x = self.conv1d(x)
        x = self.tanh(x)
        x = x.permute(0, 2, 1)
        x = self.norm1(x)
        x = self.dropout1(x)
        
        x = self.pos_encoder(x)
        x = self.transformer(x)
        x = self.post_trans_dropout(x)
        
        # last timestep
        x_last = x[:, -1, :]                  # (batch, d_model)
        out_high = self.output_dropout_high(x_last)
        out_low  = self.output_dropout_low(x_last)
        
        return self.output_high(out_high), self.output_low(out_low)

class LRFinderWrapper(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
    def forward(self, x):
        out_high, out_low = self.model(x)
        out = (out_high + out_low) / 2.0
        return out.squeeze(-1)
    
# --- Hybrid Transformer + LSTM + Causal CNN (optimized) ---
class HybridTransformerModelV4(nn.Module):
    def __init__(self, input_size, seq_len,
                 d_model=32, nhead=2, num_layers=1,
                 lstm_hidden=32, cnn_out=16, dropout=0.1):
        super().__init__()
        
        # --- CNN Branch ---
        self.cnn = CausalConv1d(in_channels=input_size, out_channels=cnn_out, kernel_size=3)
        self.cnn_bn = nn.BatchNorm1d(cnn_out)
        self.cnn_dropout = nn.Dropout(dropout)
        
        # --- LSTM Branch ---
        self.lstm = nn.LSTM(input_size, lstm_hidden, batch_first=True)
        self.lstm_dropout = nn.Dropout(dropout)
        
        # --- Transformer Branch ---
        self.input_projection = nn.Linear(input_size, d_model)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead, batch_first=True,
            dropout=dropout, activation='gelu'
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)
        self.trans_ln = nn.LayerNorm(d_model)
        
        # --- Fusion & Output ---
        fusion_size = cnn_out + lstm_hidden + d_model
        self.fc1 = nn.Linear(fusion_size, 64)
        self.fc1_ln = nn.LayerNorm(64)
        self.fc1_dropout = nn.Dropout(dropout)
        self.fc2 = nn.Linear(64, 2)  # high & low
    
    def forward(self, x):
        # CNN Branch - LeakyReLU
        cnn_x = x.transpose(1, 2)
        cnn_x = F.leaky_relu(self.cnn_bn(self.cnn(cnn_x)), 0.01)
        cnn_x = self.cnn_dropout(torch.mean(cnn_x, dim=2))
        
        # LSTM Branch - Tanh
        lstm_out, _ = self.lstm(x)
        lstm_x = self.lstm_dropout(lstm_out[:, -1, :])
        
        # Transformer Branch - GELU + LayerNorm
        trans_x = self.input_projection(x)
        trans_x = torch.mean(self.transformer(trans_x), dim=1)
        trans_x = self.trans_ln(trans_x)
        
        # Fusion - SELU
        fused = torch.cat([cnn_x, lstm_x, trans_x], dim=1)
        out = F.selu(self.fc1_ln(self.fc1(fused)))
        out = self.fc1_dropout(out)
        out = self.fc2(out)
        return out

class MultiTimeframeTransformer(nn.Module):
    """
    Multi-Timeframe Transformer:
    - 3 branches: Day / Week / Month
    - Residual connection ในแต่ละ branch + Fusion
    - Dropout + LayerNorm แทน BatchNorm1d เพื่อ train batch=1 ได้
    """
    def __init__(self, input_size, seq_len_list=[5,4,3], d_model=32, nhead=2, num_layers=2, dropout=0.1):
        super().__init__()
        self.seq_len_day, self.seq_len_week, self.seq_len_month = seq_len_list
        
        # Input projection + Dropout
        self.input_proj_day   = nn.Sequential(nn.Linear(input_size, d_model), nn.Dropout(dropout))
        self.input_proj_week  = nn.Sequential(nn.Linear(input_size, d_model), nn.Dropout(dropout))
        self.input_proj_month = nn.Sequential(nn.Linear(input_size, d_model), nn.Dropout(dropout))
        
        # Positional Encoding
        self.pos_day   = LearnablePositionalEncoding(d_model, max_len=self.seq_len_day)
        self.pos_week  = LearnablePositionalEncoding(d_model, max_len=self.seq_len_week)
        self.pos_month = LearnablePositionalEncoding(d_model, max_len=self.seq_len_month)
        
        # Transformer Encoder per branch
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead, dropout=dropout, batch_first=True, activation='gelu'
        )
        self.trans_day   = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.trans_week  = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.trans_month = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        
        # LayerNorm per branch
        self.ln_day   = nn.LayerNorm(d_model)
        self.ln_week  = nn.LayerNorm(d_model)
        self.ln_month = nn.LayerNorm(d_model)
        
        # Fusion Layer + LayerNorm + Dropout + Residual
        fusion_size = d_model * 3
        self.ln_fusion = nn.LayerNorm(fusion_size)
        self.dropout_fusion = nn.Dropout(dropout)
        self.fc1 = nn.Linear(fusion_size, 64)
        self.fc1_ln = nn.LayerNorm(64)
        self.fc2 = nn.Linear(64, 2)
        
    def forward_branch(self, x_proj, x_trans, ln):
        """
        x_proj: (B, seq_len, d_model)
        x_trans: Transformer Encoder
        ln: LayerNorm
        """
        # Transformer + Residual + LayerNorm + mean pooling
        res = x_proj
        out = x_trans(x_proj)
        out = out + res
        out = ln(out)
        out = torch.mean(out, dim=1)
        return out
        
    def forward(self, x_day, x_week, x_month):
        # Input projection + Positional Encoding
        day_proj   = self.pos_day(self.input_proj_day(x_day))
        week_proj  = self.pos_week(self.input_proj_week(x_week))
        month_proj = self.pos_month(self.input_proj_month(x_month))
        
        # Forward branch with residual
        day_feat   = self.forward_branch(day_proj, self.trans_day, self.ln_day)
        week_feat  = self.forward_branch(week_proj, self.trans_week, self.ln_week)
        month_feat = self.forward_branch(month_proj, self.trans_month, self.ln_month)
        
        # Concatenate branch features
        fused = torch.cat([day_feat, week_feat, month_feat], dim=1)
        fused_res = fused  # residual
        fused = self.ln_fusion(fused)
        fused = self.dropout_fusion(fused)
        fused = fused + fused_res  # fusion residual
        
        # Fully connected
        out = F.selu(self.fc1_ln(self.fc1(fused)))
        out = self.fc2(out)
        return out

class NormalizationLayer(nn.Module):
    """
    ใช้ normalize feature ด้วย mean/std ของ training data
    mean/std ต้องเป็น list/array ตามจำนวน feature
    """
    def __init__(self, mean, std):
        super().__init__()
        self.register_buffer("mean", torch.tensor(mean, dtype=torch.float32))
        self.register_buffer("std", torch.tensor(std, dtype=torch.float32))

    def forward(self, x):
        # x.shape = [batch, seq_len, n_features]
        return (x - self.mean) / self.std

class TransformerSingleStepForEA(nn.Module):
    def __init__(self, input_size, mean, std, d_model=64, nhead=4, num_layers=2,
                 dropout=0.1, max_len=500):
        super().__init__()
        # --- normalization layer ---
        self.norm_layer = NormalizationLayer(mean, std)

        # --- original backbone ---
        self.input_fc = nn.Linear(input_size, d_model)
        self.tanh = nn.Tanh()
        self.pos_encoder = LearnablePositionalEncoding(d_model, max_len=max_len)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=128,
            dropout=dropout,
            batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.global_pool = nn.AdaptiveAvgPool1d(1)
        self.output = nn.Linear(d_model, 2)

    def forward(self, x):
        # x.shape = [batch_size, seq_len, input_size]
        x = self.norm_layer(x)          # <-- normalize raw input
        x = self.input_fc(x)
        x = self.tanh(x)
        x = self.pos_encoder(x)
        x = self.transformer(x)
        x = x.permute(0, 2, 1)
        x = self.global_pool(x).squeeze(2)
        return self.output(x).unsqueeze(1)  # [B, 1, 2]

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 ConvGRUTransformerHL(nn.Module):
    def __init__(self, 
                 input_dim,             # features dimension
                 conv_channels=32,      # Conv1 output
                 gru_hidden=64,         # GRU Neturon
                 nhead=4,               # num_head
                 num_encoder_layers=1,  
                 dim_feedforward=128, 
                 seq_len=3,             # time step
                 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
        
        # --- Conv1D Layer (Causal) ---
        self.conv1 = nn.Conv1d(
            in_channels=input_dim, 
            out_channels=conv_channels, 
            kernel_size=self.kernel_size, 
            dilation=1,   # fixed
            padding=0     # causal padding ใส่เอง
        )
        self.conv_bn = nn.BatchNorm1d(conv_channels)
        
        # --- GRU Layer ---
        self.gru = nn.GRU(
            input_size=conv_channels, 
            hidden_size=gru_hidden, 
            batch_first=True
        )
        
        # --- 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"   # fixed GELU
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
        
        # --- Output Head (High / Low) ---
        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()
        
        # --- Conv1D ---
        x_conv = x.permute(0, 2, 1)  # (batch, features, seq_len)
        
        # Causal padding (dilation=1)
        pad = self.kernel_size - 1
        x_conv = F.pad(x_conv, (pad, 0))  # pad left only
        
        x_conv = self.conv1(x_conv)  # (batch, conv_channels, L_out)
        x_conv = self.conv_bn(x_conv)
        x_conv = torch.tanh(x_conv)  # keep tanh here
        
        # Slice ให้ตรง seq_len
        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)  # (batch, seq_len, gru_hidden)
        
        # Align sequence length
        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, :]
        
        # --- Positional Encoding ---
        gru_out = self.pos_encoder(gru_out)
        
        # --- Transformer Encoder ---
        trans_out = self.transformer(gru_out)  # (batch, seq_len, gru_hidden)
        
        # --- Output (last steps) ---
        last_steps = trans_out[:, -self.output_steps:, :]  # (batch, output_steps, hidden)
        
        high_out = self.fc_high(last_steps)  # (batch, output_steps, 1)
        low_out = self.fc_low(last_steps)    # (batch, output_steps, 1)
        
        return high_out.squeeze(-1), low_out.squeeze(-1)


class ConvGRUTransformerHLV2(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
        
        # --- 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
        )
        
        # --- 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 (High / Low) ---
        self.fc_high = nn.Linear(gru_hidden, output_steps) 
        self.fc_low  = nn.Linear(gru_hidden, output_steps)

        # --- Learnable bias for high/low ---
        self.bias_high = nn.Parameter(torch.zeros(1))
        self.bias_low  = nn.Parameter(torch.zeros(1))
    
    def forward(self, x):
        # x: (batch, seq_len, features)
        batch_size, seq_len, features = x.size()
        
        # --- 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))  # causal padding
        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)
        
        # --- GRU ---
        gru_out, _ = self.gru(x_conv)
        
        # Align sequence length
        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, :]
        
        # --- Positional Encoding ---
        gru_out = self.pos_encoder(gru_out)
        
        # --- Transformer Encoder ---
        trans_out = self.transformer(gru_out)
        
        # --- Output (last steps) ---
        last_steps = trans_out[:, -self.output_steps:, :]
        high_out = self.fc_high(last_steps) + self.bias_high
        low_out  = self.fc_low(last_steps)  + self.bias_low
        
        return high_out.squeeze(-1), low_out.squeeze(-1)



# ----------------------------
# Utility: ensure batch dim
# ----------------------------
def ensure_batch_dim(x):
    if x.dim() == 2:
        x = x.unsqueeze(0)
    return x

# ----------------------------
# Model (Seq2SeqAttentionHL) - เอาไว้ครั้งเดียว
# ----------------------------
class Seq2SeqAttentionHL(torch.nn.Module):
    def __init__(self, input_size, hidden_size, num_layers=1, dropout=0.2, device='cpu'):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.device = device
        self.encoder = torch.nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, dropout=0)
        self.decoder = torch.nn.LSTM(2, hidden_size, num_layers, batch_first=True, dropout=0)
        self.fc_high = torch.nn.Linear(hidden_size, 1)
        self.fc_low  = torch.nn.Linear(hidden_size, 1)
        self.dropout = torch.nn.Dropout(dropout)
        self.attention = torch.nn.Linear(hidden_size * 2, 1)

    def forward(self, src, horizon=4, tgt=None, teacher_forcing_ratio=0.5):
        src = ensure_batch_dim(src)
        batch_size, seq_len, _ = src.shape
        
        # Encoder
        encoder_outputs, (hidden, cell) = self.encoder(src)
        encoder_outputs = self.dropout(encoder_outputs)
        
        # Decoder initialization
        decoder_input = src[:, -1, :2].unsqueeze(1)
        out_high_seq = torch.zeros(batch_size, horizon, 1, device=self.device)
        out_low_seq = torch.zeros(batch_size, horizon, 1, device=self.device)
        
        for t in range(horizon):
            # Decoder step
            decoder_output, (hidden, cell) = self.decoder(decoder_input, (hidden, cell))
            
            # Attention mechanism
            decoder_output_expanded = decoder_output.expand(-1, seq_len, -1)
            attention_scores = torch.tanh(self.attention(
                torch.cat([decoder_output_expanded, encoder_outputs], dim=-1)
            ))
            attention_weights = torch.softmax(attention_scores, dim=1)
            context = torch.sum(attention_weights * encoder_outputs, dim=1, keepdim=True)
            
            # Combine context and decoder output
            combined = decoder_output + context
            
            # Output projections
            high = torch.tanh(self.fc_high(combined))
            low = torch.tanh(self.fc_low(combined))
            
            out_high_seq[:, t:t+1] = high
            out_low_seq[:, t:t+1] = low
            
            # Teacher forcing or use own predictions
            if tgt is not None and torch.rand(1) < teacher_forcing_ratio:
                decoder_input = tgt[:, t:t+1, :2]
            else:
                decoder_input = torch.cat([high, low], dim=-1)
                
        return out_high_seq, out_low_seq
    
class CausalConvBlockSafe(nn.Module):
    def __init__(self, input_size, dropout=0.1):
        super().__init__()
        self.kernel_size1 = 4
        self.kernel_size2 = 4

        self.conv1 = nn.Sequential(
            nn.ConstantPad1d((self.kernel_size1-1, 0), 0),
            nn.Conv1d(input_size, 128, kernel_size=self.kernel_size1),
            nn.BatchNorm1d(128),
            nn.Tanh(),                # <-- เปลี่ยนจาก LeakyReLU เป็น tanh
            nn.Dropout(dropout)
        )
        self.conv2 = nn.Sequential(
            nn.ConstantPad1d((self.kernel_size2-1, 0), 0),
            nn.Conv1d(128, 64, kernel_size=self.kernel_size2),
            nn.BatchNorm1d(64),
            nn.Tanh(),                # <-- เปลี่ยนจาก LeakyReLU เป็น tanh
            nn.Dropout(dropout)
        )
        self.layernorm = nn.LayerNorm(64)  # LayerNorm หลัง conv block

    def forward(self, x):
        conv_out = self.conv1(x)
        conv_out = self.conv2(conv_out)
        seq_len_in = x.size(2)
        seq_len_out = conv_out.size(2)
        if seq_len_out < seq_len_in:
            raise ValueError(f"Conv output shorter than input! {seq_len_out} < {seq_len_in}")
        conv_out = conv_out[:, :, -seq_len_in:]
        conv_out = conv_out.transpose(1,2)
        conv_out = self.layernorm(conv_out)
        return conv_out.transpose(1,2)

class Seq2SeqAttentionHLWithConv(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers=2, dropout=0.2, device='cpu'):
        super().__init__()
        self.hidden_size = hidden_size
        self.device = device

        # --- Causal Conv block ---
        self.conv_block = CausalConvBlockSafe(input_size, dropout=dropout)

        # --- LSTM layers ---
        self.encoder = nn.LSTM(64, hidden_size, num_layers=num_layers,
                               batch_first=True, dropout=dropout if num_layers > 1 else 0.0)
        self.decoder = nn.LSTM(2, hidden_size, num_layers=num_layers,
                               batch_first=True, dropout=dropout if num_layers > 1 else 0.0)

        # --- LayerNorm for LSTM outputs ---
        self.layernorm_enc = nn.LayerNorm(hidden_size)
        self.layernorm_dec = nn.LayerNorm(hidden_size)

        # --- Attention mechanism ---
        self.attention = nn.MultiheadAttention(hidden_size, num_heads=4,
                                               dropout=dropout, batch_first=True)
        self.layernorm_attn = nn.LayerNorm(hidden_size)

        # --- Output layers ---
        self.fc_high = nn.Sequential(
            nn.Linear(hidden_size, 32),
            nn.Tanh(),           # <-- เปลี่ยนจาก ReLU เป็น tanh
            nn.Dropout(dropout),
            nn.Linear(32, 1),
            nn.Tanh()
        )

        self.fc_low = nn.Sequential(
            nn.Linear(hidden_size, 32),
            nn.Tanh(),           # <-- เปลี่ยนจาก ReLU เป็น tanh
            nn.Dropout(dropout),
            nn.Linear(32, 1),
            nn.Tanh()
        )

    def forward(self, src, horizon=4, tgt=None, teacher_forcing_ratio=1):
        batch_size, seq_len, _ = src.shape

        # --- Conv block ---
        src_conv = src.transpose(1, 2)        # (B, features, seq_len)
        conv_out = self.conv_block(src_conv)  # (B, features, seq_len)
        conv_out = conv_out.transpose(1, 2)   # (B, seq_len, features)

        # --- Encoder ---
        encoder_outputs, (hidden, cell) = self.encoder(conv_out)
        encoder_outputs = self.layernorm_enc(encoder_outputs)

        # --- Decoder ---
        decoder_input = src[:, -1, :2].unsqueeze(1)  # initial input
        outputs_high, outputs_low = [], []

        for t in range(horizon):
            decoder_out, (hidden, cell) = self.decoder(decoder_input, (hidden, cell))
            decoder_out = self.layernorm_dec(decoder_out)

            # Multi-head attention with residual + LayerNorm
            attn_output, _ = self.attention(decoder_out, encoder_outputs, encoder_outputs)
            attn_output = self.layernorm_attn(attn_output + decoder_out)

            # Predict
            high = self.fc_high(attn_output)
            low = self.fc_low(attn_output)

            outputs_high.append(high)
            outputs_low.append(low)

            # Next input
            if tgt is not None and torch.rand(1) < teacher_forcing_ratio:
                decoder_input = tgt[:, t:t+1, :2]
            else:
                decoder_input = torch.cat([high, low], dim=-1)

        return torch.cat(outputs_high, dim=1), torch.cat(outputs_low, dim=1)


# --- Causal Conv1D ---
class CausalConv1d(nn.Conv1d):
    def __init__(self, in_channels, out_channels, kernel_size, **kwargs):
        super().__init__(in_channels, out_channels, kernel_size, **kwargs)
        self.left_padding = kernel_size - 1

    def forward(self, x):
        x = F.pad(x, (self.left_padding, 0))
        return super().forward(x)

# # ----------------------------
# # Causal Conv1d
# # ----------------------------
class CausalConv1d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, dropout=0.1):
        super().__init__()
        self.padding = kernel_size - 1  # causal padding
        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size)
        self.dropout = nn.Dropout(dropout)
        self.ln = nn.LayerNorm(out_channels)

    def forward(self, x):
        # x: [batch, seq_len, channels]
        x = x.permute(0, 2, 1)              # -> [batch, channels, seq_len]
        x = F.pad(x, (self.padding, 0))     # causal padding
        x = self.conv(x)                     # -> [batch, out_channels, seq_len]
        x = x.permute(0, 2, 1)              # -> [batch, seq_len, out_channels]
        x = self.dropout(x)
        x = self.ln(x)
        return torch.tanh(x)

# # ----------------------------
# # Seq2Seq Branch
# # ----------------------------
class Seq2SeqBranch(nn.Module):
    def __init__(self, in_channels, cnn_channels=32, lstm_hidden=64, kernel_size=3, dropout=0.1):
        super().__init__()
        self.cnn = CausalConv1d(in_channels, cnn_channels, kernel_size, dropout)
        self.lstm = nn.LSTM(cnn_channels, lstm_hidden, batch_first=True)
        self.pos_embed = nn.Embedding(1000, cnn_channels)  # positional embedding

    def forward(self, x):
        x = self.cnn(x)  # [batch, seq_len, cnn_channels]
        batch_size, seq_len, _ = x.shape
        positions = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(batch_size, -1)
        x = x + self.pos_embed(positions)
        output, _ = self.lstm(x)
        return output  # [batch, seq_len, lstm_hidden]

# ----------------------------
# Hybrid Multi-Branch Multi-Step Attention
# ----------------------------
class HybridMultiBranchMultiStepAttention(nn.Module):
    def __init__(self, horizon=4, cnn_channels=32, lstm_hidden=64, num_heads=4, dropout=0.1):
        super().__init__()
        self.horizon = horizon
        self.lstm_hidden = lstm_hidden

        # Branches: OHLC, ADX, ATR, Tick
        self.branches = nn.ModuleList([
            Seq2SeqBranch(4, cnn_channels, lstm_hidden, dropout=dropout),  # OHLC
            Seq2SeqBranch(1, cnn_channels, lstm_hidden, dropout=dropout),  # ADX
            Seq2SeqBranch(1, cnn_channels, lstm_hidden, dropout=dropout),  # ATR
            Seq2SeqBranch(1, cnn_channels, lstm_hidden, dropout=dropout),  # Tick
        ])

        # Multi-head attention across features
        self.multihead = nn.MultiheadAttention(
            embed_dim=lstm_hidden,
            num_heads=num_heads,
            batch_first=True,
            dropout=dropout
        )

        self.final_ln = nn.LayerNorm(lstm_hidden)
        
        # Decoder LSTM สำหรับ sequence-to-sequence prediction
        self.decoder_lstm = nn.LSTM(lstm_hidden, lstm_hidden, batch_first=True)
        
        # TimeDistributed FC layers สำหรับ multi-step prediction
        self.fc_high = nn.Linear(lstm_hidden, 1)
        self.fc_low  = nn.Linear(lstm_hidden, 1)

    def forward(self, x):
        # x shape: [batch, seq_len, 7] -> split features
        head_inputs = [
            x[:, :, 0:4],  # OHLC
            x[:, :, 4:5],  # ADX
            x[:, :, 5:6],  # ATR
            x[:, :, 6:7]   # Tick
        ]

        # Process each branch
        branch_outs = [branch(h) for branch, h in zip(self.branches, head_inputs)]
        
        # Stack along new dimension for features: [batch, seq_len, hidden, num_branches]
        branch_stack = torch.stack(branch_outs, dim=-1)
        
        # Merge branches (mean across branch dimension)
        branch_merged = branch_stack.mean(dim=-1)  # [batch, seq_len, hidden]

        # Multihead attention across time steps
        attn_out, _ = self.multihead(branch_merged, branch_merged, branch_merged)

        # Context: hidden state สุดท้าย
        context = attn_out[:, -1:, :]  # [batch, 1, hidden]

        # Repeat context สำหรับ horizon steps
        decoder_input = context.repeat(1, self.horizon, 1)  # [batch, horizon, hidden]

        # Decoder LSTM
        decoder_out, _ = self.decoder_lstm(decoder_input)  # [batch, horizon, hidden]
        decoder_out = self.final_ln(decoder_out)

        # --- TimeDistributed FC ---
        # reshape [batch*horizon, hidden] เพื่อให้ Linear ทำงานทีละ timestep
        batch, horizon, hidden = decoder_out.shape
        decoder_flat = decoder_out.reshape(batch*horizon, hidden)
        out_high_flat = self.fc_high(decoder_flat)
        out_low_flat  = self.fc_low(decoder_flat)
        # reshape กลับ [batch, horizon]
        out_high = out_high_flat.reshape(batch, horizon)
        out_low  = out_low_flat.reshape(batch, horizon)

        return out_high, out_low