"""N-BEATS模型实现：预测未来1小时BTC价格"""

import torch
import torch.nn as nn
import numpy as np


class NBeatsBlock(nn.Module):
    """N-BEATS基础块"""

    def __init__(self, input_size, theta_size, basis_function, layers, layer_size):
        super().__init__()
        self.input_size = input_size
        self.theta_size = theta_size
        self.basis_function = basis_function

        # 全连接层栈
        self.fc_stack = nn.ModuleList([
            nn.Linear(input_size if i == 0 else layer_size, layer_size)
            for i in range(layers)
        ])

        # Theta层（用于生成基函数参数）
        self.theta_b_fc = nn.Linear(layer_size, theta_size, bias=False)  # backcast
        self.theta_f_fc = nn.Linear(layer_size, theta_size, bias=False)  # forecast

        self.relu = nn.ReLU()

    def forward(self, x):
        """
        Args:
            x: (batch_size, input_size)
        Returns:
            backcast: (batch_size, input_size)
            forecast: (batch_size, forecast_size)
        """
        # 通过FC层栈
        for fc in self.fc_stack:
            x = self.relu(fc(x))

        # 生成theta参数
        theta_b = self.theta_b_fc(x)  # backcast coefficients
        theta_f = self.theta_f_fc(x)  # forecast coefficients

        # 应用基函数
        backcast = self.basis_function(theta_b, is_forecast=False)
        forecast = self.basis_function(theta_f, is_forecast=True)

        return backcast, forecast


class GenericBasis(nn.Module):
    """通用基函数"""

    def __init__(self, backcast_size, forecast_size):
        super().__init__()
        self.backcast_size = backcast_size
        self.forecast_size = forecast_size

    def forward(self, theta, is_forecast):
        if is_forecast:
            return theta[:, :self.forecast_size]
        else:
            return theta[:, :self.backcast_size]


class TrendBasis(nn.Module):
    """趋势基函数（多项式）"""

    def __init__(self, degree, backcast_size, forecast_size):
        super().__init__()
        self.backcast_size = backcast_size
        self.forecast_size = forecast_size

        # 多项式时间点
        backcast_time = torch.arange(backcast_size, dtype=torch.float32) / backcast_size
        forecast_time = torch.arange(forecast_size, dtype=torch.float32) / forecast_size

        # 构建多项式基
        self.backcast_basis = torch.stack([backcast_time ** i for i in range(degree)], dim=1)
        self.forecast_basis = torch.stack([forecast_time ** i for i in range(degree)], dim=1)

    def forward(self, theta, is_forecast):
        basis = self.forecast_basis if is_forecast else self.backcast_basis
        basis = basis.to(theta.device)
        return torch.matmul(theta, basis.T)


class SeasonalityBasis(nn.Module):
    """季节性基函数（傅里叶级数）"""

    def __init__(self, harmonics, backcast_size, forecast_size):
        super().__init__()
        self.backcast_size = backcast_size
        self.forecast_size = forecast_size

        backcast_time = 2 * np.pi * torch.arange(backcast_size, dtype=torch.float32) / backcast_size
        forecast_time = 2 * np.pi * torch.arange(forecast_size, dtype=torch.float32) / forecast_size

        # 构建傅里叶基（sin和cos）
        backcast_basis = []
        forecast_basis = []
        for i in range(1, harmonics + 1):
            backcast_basis.extend([torch.cos(i * backcast_time), torch.sin(i * backcast_time)])
            forecast_basis.extend([torch.cos(i * forecast_time), torch.sin(i * forecast_time)])

        self.backcast_basis = torch.stack(backcast_basis, dim=1)
        self.forecast_basis = torch.stack(forecast_basis, dim=1)

    def forward(self, theta, is_forecast):
        basis = self.forecast_basis if is_forecast else self.backcast_basis
        basis = basis.to(theta.device)
        return torch.matmul(theta, basis.T)


class NBeatsStack(nn.Module):
    """N-BEATS栈（多个块的组合）"""

    def __init__(self, input_size, forecast_size, num_blocks, block_layers,
                 layer_size, theta_size, basis_function):
        super().__init__()
        self.blocks = nn.ModuleList([
            NBeatsBlock(
                input_size=input_size,
                theta_size=theta_size,
                basis_function=basis_function,
                layers=block_layers,
                layer_size=layer_size
            )
            for _ in range(num_blocks)
        ])
        self.forecast_size = forecast_size

    def forward(self, x):
        """
        Args:
            x: (batch_size, input_size)
        Returns:
            forecast: (batch_size, forecast_size)
        """
        residual = x
        forecast = torch.zeros(x.size(0), self.forecast_size, device=x.device)

        for block in self.blocks:
            backcast, block_forecast = block(residual)
            residual = residual - backcast
            forecast = forecast + block_forecast

        return forecast


class NBeatsConfig:
    """N-BEATS模型配置"""

    def __init__(
        self,
        input_size=168,           # 输入序列长度
        forecast_size=1,          # 预测长度
        stacks=['trend', 'seasonality', 'generic'],  # 使用的栈类型
        num_blocks_per_stack=3,   # 每个栈的块数
        block_layers=4,           # 每个块的FC层数
        layer_size=256,           # FC层大小
        trend_degree=2,           # 趋势多项式阶数
        seasonality_harmonics=4,  # 季节性谐波数
    ):
        self.input_size = input_size
        self.forecast_size = forecast_size
        self.stacks = stacks
        self.num_blocks_per_stack = num_blocks_per_stack
        self.block_layers = block_layers
        self.layer_size = layer_size
        self.trend_degree = trend_degree
        self.seasonality_harmonics = seasonality_harmonics


class NBeats(nn.Module):
    """N-BEATS模型"""

    def __init__(self, config):
        super().__init__()
        self.config = config

        # 创建多个栈
        self.stacks = nn.ModuleList()

        for stack_type in config.stacks:
            if stack_type == 'trend':
                basis_function = TrendBasis(
                    degree=config.trend_degree,
                    backcast_size=config.input_size,
                    forecast_size=config.forecast_size
                )
                theta_size = config.trend_degree

            elif stack_type == 'seasonality':
                basis_function = SeasonalityBasis(
                    harmonics=config.seasonality_harmonics,
                    backcast_size=config.input_size,
                    forecast_size=config.forecast_size
                )
                theta_size = 2 * config.seasonality_harmonics

            else:  # generic
                basis_function = GenericBasis(
                    backcast_size=config.input_size,
                    forecast_size=config.forecast_size
                )
                theta_size = max(config.input_size, config.forecast_size)

            stack = NBeatsStack(
                input_size=config.input_size,
                forecast_size=config.forecast_size,
                num_blocks=config.num_blocks_per_stack,
                block_layers=config.block_layers,
                layer_size=config.layer_size,
                theta_size=theta_size,
                basis_function=basis_function
            )
            self.stacks.append(stack)

    def forward(self, x):
        """
        Args:
            x: (batch_size, lookback, features) 或 (batch_size, lookback)
        Returns:
            forecast: (batch_size, forecast_size)
        """
        # 如果输入是多特征，取最后一个特征（close price）
        if len(x.shape) == 3:
            x = x[:, :, -1]  # (batch_size, lookback)

        # 确保是2D
        if len(x.shape) == 1:
            x = x.unsqueeze(0)

        # 通过所有栈
        forecast = torch.zeros(x.size(0), self.config.forecast_size, device=x.device)
        residual = x

        for stack in self.stacks:
            stack_forecast = stack(residual)
            forecast = forecast + stack_forecast
            # 更新残差（用于下一个栈）
            # 注意：这里简化了，标准N-BEATS会从每个栈获取backcast

        return forecast.squeeze(-1)  # (batch_size,)

    def count_parameters(self):
        """统计模型参数数量"""
        return sum(p.numel() for p in self.parameters() if p.requires_grad)

    def save(self, path):
        """保存模型"""
        torch.save({
            'model_state_dict': self.state_dict(),
            'config': self.config
        }, path)

    @classmethod
    def load(cls, path, device='cpu'):
        """加载模型"""
        checkpoint = torch.load(path, map_location=device)
        model = cls(checkpoint['config'])
        model.load_state_dict(checkpoint['model_state_dict'])
        return model