"""DLinear: Decomposition Linear model for time series forecasting.

Reference: Are Transformers Effective for Time Series Forecasting? (DLinear)
Paper: https://arxiv.org/abs/2205.13504
"""

import torch
import torch.nn as nn
from dataclasses import dataclass
from typing import Literal

from ..base import BasePredictor, ModelConfig


@dataclass
class DLinearConfig(ModelConfig):
    """Configuration for DLinear model."""

    moving_avg_kernel: int = 25  # Kernel size for moving average decomposition
    individual: bool = False  # Whether to use individual linear layers per feature


class MovingAvg(nn.Module):
    """Moving average block for trend extraction."""

    def __init__(self, kernel_size: int):
        """Initialize moving average layer.

        Args:
            kernel_size: Size of the moving average window
        """
        super().__init__()
        self.kernel_size = kernel_size
        self.avg = nn.AvgPool1d(
            kernel_size=kernel_size,
            stride=1,
            padding=kernel_size // 2,
            count_include_pad=False
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply moving average.

        Args:
            x: Input tensor of shape (batch_size, seq_len, features)

        Returns:
            Smoothed tensor
        """
        # Padding at both ends
        front = x[:, 0:1, :].repeat(1, self.kernel_size // 2, 1)
        end = x[:, -1:, :].repeat(1, (self.kernel_size - 1) // 2, 1)
        x_padded = torch.cat([front, x, end], dim=1)

        # Apply average pooling (swap dimensions for AvgPool1d)
        x_padded = x_padded.permute(0, 2, 1)  # (batch, features, seq_len)
        x_smoothed = self.avg(x_padded)
        x_smoothed = x_smoothed.permute(0, 2, 1)  # (batch, seq_len, features)

        # Trim to original length
        return x_smoothed[:, :x.size(1), :]


class SeriesDecomposition(nn.Module):
    """Series decomposition block to extract trend and seasonal components."""

    def __init__(self, kernel_size: int):
        """Initialize decomposition layer.

        Args:
            kernel_size: Size of the moving average window for trend
        """
        super().__init__()
        self.moving_avg = MovingAvg(kernel_size)

    def forward(self, x: torch.Tensor) -> tuple:
        """Decompose series into trend and seasonal components.

        Args:
            x: Input tensor of shape (batch_size, seq_len, features)

        Returns:
            Tuple of (seasonal, trend) tensors
        """
        trend = self.moving_avg(x)
        seasonal = x - trend
        return seasonal, trend


class DLinear(BasePredictor):
    """DLinear model: Decomposition + Linear layers for forecasting."""

    def __init__(self, config: DLinearConfig):
        """Initialize DLinear model.

        Args:
            config: Model configuration
        """
        super().__init__(config)

        self.decomposition = SeriesDecomposition(config.moving_avg_kernel)

        if config.individual:
            # Individual linear layer for each feature
            self.linear_seasonal = nn.ModuleList([
                nn.Linear(config.lookback, 1)
                for _ in range(config.input_dim)
            ])
            self.linear_trend = nn.ModuleList([
                nn.Linear(config.lookback, 1)
                for _ in range(config.input_dim)
            ])
        else:
            # Shared linear layers
            self.linear_seasonal = nn.Linear(config.lookback, 1)
            self.linear_trend = nn.Linear(config.lookback, 1)

        # Final aggregation layer
        self.aggregation = nn.Sequential(
            nn.Linear(config.input_dim, config.hidden_dim),
            nn.ReLU(),
            nn.Dropout(config.dropout),
            nn.Linear(config.hidden_dim, 1)
        )

        # Output activation
        # 注意：对于分类任务，使用Identity而非Sigmoid
        # 因为BCEWithLogitsLoss内部已包含Sigmoid，避免双重激活
        if config.task == 'classification':
            self.output_activation = nn.Identity()  # 输出logits，由损失函数处理
        else:
            self.output_activation = nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass through DLinear.

        Args:
            x: Input tensor of shape (batch_size, lookback, input_dim)

        Returns:
            Output tensor of shape (batch_size, 1)
        """
        # Decompose into seasonal and trend
        seasonal, trend = self.decomposition(x)

        # Apply linear layers to each component
        if self.config.individual:
            seasonal_out = []
            trend_out = []
            for i in range(self.config.input_dim):
                # Extract feature i: (batch, seq_len)
                s_i = seasonal[:, :, i]
                t_i = trend[:, :, i]

                # Apply linear transformation: (batch, seq_len) -> (batch, 1)
                seasonal_out.append(self.linear_seasonal[i](s_i))
                trend_out.append(self.linear_trend[i](t_i))

            # Stack outputs: (batch, input_dim)
            seasonal_out = torch.stack(seasonal_out, dim=1).squeeze(-1)
            trend_out = torch.stack(trend_out, dim=1).squeeze(-1)
        else:
            # Shared linear layer
            # Swap dimensions: (batch, seq_len, features) -> (batch, features, seq_len)
            seasonal = seasonal.permute(0, 2, 1)
            trend = trend.permute(0, 2, 1)

            # Apply linear: (batch, features, seq_len) -> (batch, features, 1)
            seasonal_out = self.linear_seasonal(seasonal).squeeze(-1)
            trend_out = self.linear_trend(trend).squeeze(-1)

        # Combine seasonal and trend
        combined = seasonal_out + trend_out  # (batch, input_dim)

        # Aggregate and produce final prediction
        output = self.aggregation(combined)  # (batch, 1)
        output = self.output_activation(output)

        return output


if __name__ == "__main__":
    # Test DLinear model
    config = DLinearConfig(
        input_dim=30,
        lookback=168,
        task='classification',
        hidden_dim=512,
        moving_avg_kernel=25
    )

    model = DLinear(config)
    print(f"DLinear parameters: {model.count_parameters()}")

    # Test forward pass
    x = torch.randn(32, 168, 30)
    output = model(x)
    print(f"Output shape: {output.shape}")
    print(f"Output range: [{output.min():.4f}, {output.max():.4f}]")

    # Test ONNX export
    model.to_onnx('/tmp/dlinear_test.onnx')
    print("ONNX export successful!")