"""Base model interface for all prediction models."""

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Literal, Optional
import torch
import torch.nn as nn


@dataclass
class ModelConfig:
    """Base configuration for all models."""

    input_dim: int  # Number of input features
    lookback: int  # Sequence length
    task: Literal['classification', 'regression'] = 'classification'
    hidden_dim: int = 512
    dropout: float = 0.1


class BasePredictor(nn.Module, ABC):
    """Abstract base class for all prediction models."""

    def __init__(self, config: ModelConfig):
        """Initialize model with configuration.

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

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

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

        Returns:
            Output tensor. Shape depends on task:
                - Classification: (batch_size, 1) with sigmoid output
                - Regression: (batch_size, 1) with raw value
        """
        pass

    def predict(self, x: torch.Tensor) -> torch.Tensor:
        """Make predictions (inference mode).

        Args:
            x: Input tensor

        Returns:
            Predictions
        """
        self.eval()
        with torch.no_grad():
            return self.forward(x)

    def get_loss_function(self, pos_weight=None, device='cpu') -> nn.Module:
        """Get appropriate loss function for the task.

        Args:
            pos_weight: Weight for positive class (for imbalanced data)
            device: Device to place tensors on

        Returns:
            Loss function
        """
        if self.config.task == 'classification':
            if pos_weight is not None:
                # 使用float32以兼容MPS (Mac M1)
                pos_weight_tensor = torch.tensor([pos_weight], dtype=torch.float32).to(device)
                return nn.BCEWithLogitsLoss(pos_weight=pos_weight_tensor)
            else:
                return nn.BCELoss()
        else:  # regression
            return nn.MSELoss()

    def save(self, path: str):
        """Save model checkpoint.

        Args:
            path: Path to save checkpoint
        """
        torch.save({
            'model_state_dict': self.state_dict(),
            'config': self.config
        }, path)

    @classmethod
    def load(cls, path: str, device: Optional[str] = None):
        """Load model from checkpoint.

        Args:
            path: Path to checkpoint
            device: Device to load model to

        Returns:
            Loaded model
        """
        checkpoint = torch.load(path, map_location=device)
        model = cls(checkpoint['config'])
        model.load_state_dict(checkpoint['model_state_dict'])
        return model

    def to_onnx(self, path: str, dummy_input: Optional[torch.Tensor] = None):
        """Export model to ONNX format.

        Args:
            path: Path to save ONNX model
            dummy_input: Example input tensor. If None, will create default.
        """
        # Move to CPU for ONNX export (critical for M1)
        self.cpu()
        self.eval()

        # Create dummy input if not provided
        if dummy_input is None:
            dummy_input = torch.randn(
                1,
                self.config.lookback,
                self.config.input_dim
            )

        # Ensure dummy input is on CPU
        dummy_input = dummy_input.cpu()

        # Export to ONNX
        torch.onnx.export(
            self,
            dummy_input,
            path,
            export_params=True,
            opset_version=14,
            do_constant_folding=True,
            input_names=['input'],
            output_names=['output'],
            dynamic_axes={
                'input': {0: 'batch_size'},
                'output': {0: 'batch_size'}
            }
        )

        print(f"Model exported to ONNX: {path}")

    def count_parameters(self) -> int:
        """Count trainable parameters.

        Returns:
            Number of trainable parameters
        """
        return sum(p.numel() for p in self.parameters() if p.requires_grad)


if __name__ == "__main__":
    # Test base class (with a dummy implementation)
    class DummyModel(BasePredictor):
        def __init__(self, config: ModelConfig):
            super().__init__(config)
            self.fc = nn.Linear(config.input_dim * config.lookback, 1)
            if config.task == 'classification':
                self.activation = nn.Sigmoid()
            else:
                self.activation = nn.Identity()

        def forward(self, x):
            batch_size = x.size(0)
            x = x.view(batch_size, -1)
            x = self.fc(x)
            return self.activation(x)

    # Test
    config = ModelConfig(input_dim=10, lookback=168, task='classification')
    model = DummyModel(config)

    print(f"Parameters: {model.count_parameters()}")

    # Test forward pass
    x = torch.randn(32, 168, 10)
    output = model(x)
    print(f"Output shape: {output.shape}")