#!/usr/bin/env python3
"""
PVA (Price Variation Attention) Model
======================================
Complete PVA model combining transformer attention encoder with XGBoost head.

This module implements the full PVA architecture:
1. Transformer encoder (PriceFocusedAttention) for sequence representation
2. XGBoost head for final prediction
3. Outputs: direction, magnitude, confidence

Architecture:
    Input (seq_len, n_features)
        |
        v
    [PriceFocusedAttention Encoder]
    (4 layers, 256 d_model, 8 heads)
        |
        v
    [Sequence Pooling: last/mean]
        |
        v
    [XGBoost Head]
        |
        v
    Output: (direction, magnitude, confidence)

Author: ML-Specialist (NEXUS v4.0)
Version: 1.0.0
Created: 2026-01-25
"""

from typing import Tuple, List, Optional, Dict, Any, Union
from dataclasses import dataclass, field
from pathlib import Path
import json

import numpy as np
import pandas as pd
from loguru import logger

try:
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False
    logger.warning("PyTorch not available")

try:
    import xgboost as xgb
    XGB_AVAILABLE = True
except ImportError:
    XGB_AVAILABLE = False
    logger.warning("XGBoost not available")

try:
    import joblib
    JOBLIB_AVAILABLE = True
except ImportError:
    JOBLIB_AVAILABLE = False

# Import attention module from parent
from ...attention import PriceFocusedAttention, PriceAttentionConfig


@dataclass
class PVAConfig:
    """
    Configuration for PVA (Price Variation Attention) Model.

    Contains all hyperparameters for both the transformer encoder
    and the XGBoost prediction head.
    """

    # Transformer Encoder Config
    d_model: int = 256
    n_heads: int = 8
    n_layers: int = 4
    d_ff: int = 1024
    dropout: float = 0.1
    max_seq_len: int = 512
    attention_dropout: float = 0.1

    # Input dimensions
    input_features: int = 30  # Determined by feature engineering
    sequence_length: int = 100

    # Pooling strategy for sequence representation
    pooling: str = "last"  # "last", "mean", "max", "first"

    # XGBoost Head Config
    xgb_n_estimators: int = 200
    xgb_max_depth: int = 6
    xgb_learning_rate: float = 0.05
    xgb_subsample: float = 0.8
    xgb_colsample_bytree: float = 0.8
    xgb_reg_alpha: float = 0.1
    xgb_reg_lambda: float = 1.0
    xgb_min_child_weight: int = 3
    xgb_gamma: float = 0.1

    # Confidence calibration
    confidence_temperature: float = 1.5

    # Output configuration
    output_direction: bool = True
    output_magnitude: bool = True
    output_confidence: bool = True

    # Training device
    device: str = "cpu"

    def to_attention_config(self) -> PriceAttentionConfig:
        """Convert to PriceAttentionConfig for attention module."""
        return PriceAttentionConfig(
            d_model=self.d_model,
            n_heads=self.n_heads,
            n_layers=self.n_layers,
            d_ff=self.d_ff,
            dropout=self.dropout,
            attention_dropout=self.attention_dropout,
            max_seq_len=self.max_seq_len,
            input_features=self.input_features
        )

    def to_dict(self) -> Dict[str, Any]:
        """Convert config to dictionary."""
        return {
            'd_model': self.d_model,
            'n_heads': self.n_heads,
            'n_layers': self.n_layers,
            'd_ff': self.d_ff,
            'dropout': self.dropout,
            'max_seq_len': self.max_seq_len,
            'attention_dropout': self.attention_dropout,
            'input_features': self.input_features,
            'sequence_length': self.sequence_length,
            'pooling': self.pooling,
            'xgb_n_estimators': self.xgb_n_estimators,
            'xgb_max_depth': self.xgb_max_depth,
            'xgb_learning_rate': self.xgb_learning_rate,
            'xgb_subsample': self.xgb_subsample,
            'xgb_colsample_bytree': self.xgb_colsample_bytree,
            'xgb_reg_alpha': self.xgb_reg_alpha,
            'xgb_reg_lambda': self.xgb_reg_lambda,
            'xgb_min_child_weight': self.xgb_min_child_weight,
            'xgb_gamma': self.xgb_gamma,
            'confidence_temperature': self.confidence_temperature,
            'output_direction': self.output_direction,
            'output_magnitude': self.output_magnitude,
            'output_confidence': self.output_confidence,
            'device': self.device
        }

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> 'PVAConfig':
        """Create config from dictionary."""
        return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__})


@dataclass
class PVAPrediction:
    """Output prediction from PVA model."""

    direction: float  # -1 to 1 (bearish to bullish)
    magnitude: float  # Expected return magnitude (absolute)
    confidence: float  # 0 to 1 confidence score
    raw_prediction: float  # Raw model output
    attention_weights: Optional[np.ndarray] = None  # Attention for interpretability

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            'direction': float(self.direction),
            'magnitude': float(self.magnitude),
            'confidence': float(self.confidence),
            'raw_prediction': float(self.raw_prediction)
        }


class PVAModel(nn.Module):
    """
    PVA (Price Variation Attention) Model.

    Combines a transformer-based attention encoder for sequence
    representation with an XGBoost head for robust predictions.

    The model is designed for price movement prediction without
    relying on temporal features (hour, day, session).

    Architecture:
    1. PriceFocusedAttention encoder processes sequences
    2. Pooling layer extracts fixed-size representation
    3. XGBoost head predicts direction and magnitude

    Usage:
        config = PVAConfig(input_features=30, sequence_length=100)
        model = PVAModel(config)

        # Training (two-stage)
        model.fit_encoder(X_train, y_train, X_val, y_val)
        model.fit_xgboost(X_train, y_train)

        # Inference
        prediction = model.predict(X_test)
        direction, magnitude, confidence = prediction.direction, ...

        # Get attention scores for interpretability
        attention_scores = model.get_attention_scores(X_test)
    """

    def __init__(self, config: Optional[PVAConfig] = None):
        """
        Initialize PVA model.

        Args:
            config: Model configuration (uses defaults if None)
        """
        if not TORCH_AVAILABLE:
            raise ImportError("PyTorch is required. Install with: pip install torch")

        super().__init__()

        self.config = config or PVAConfig()
        self.device = torch.device(self.config.device)

        # Initialize attention encoder
        attention_config = self.config.to_attention_config()
        self.attention_encoder = PriceFocusedAttention(
            config=attention_config,
            input_features=self.config.input_features
        )

        # XGBoost head (initialized later during training)
        self._xgb_direction: Optional[xgb.XGBClassifier] = None
        self._xgb_magnitude: Optional[xgb.XGBRegressor] = None

        # Training state
        self._is_trained = False
        self._encoder_trained = False

        # Move to device
        self.to(self.device)

        logger.info(f"PVAModel initialized")
        logger.info(f"  Device: {self.device}")
        logger.info(f"  Encoder: {self.config.n_layers} layers, {self.config.d_model} dim")
        logger.info(f"  Pooling: {self.config.pooling}")

    def forward(
        self,
        x: torch.Tensor,
        return_attention: bool = False
    ) -> Tuple[torch.Tensor, Optional[List[torch.Tensor]]]:
        """
        Forward pass through attention encoder.

        Args:
            x: Input tensor of shape (batch, seq_len, n_features)
            return_attention: Whether to return attention weights

        Returns:
            encoded: Encoded representation (batch, d_model)
            attentions: Optional list of attention weight tensors
        """
        # Ensure input is on correct device
        x = x.to(self.device)

        # Pass through attention encoder
        encoded_seq, attentions = self.attention_encoder(
            x, return_all_attentions=return_attention
        )

        # Apply pooling to get fixed-size representation
        if self.config.pooling == "last":
            encoded = encoded_seq[:, -1, :]
        elif self.config.pooling == "first":
            encoded = encoded_seq[:, 0, :]
        elif self.config.pooling == "mean":
            encoded = encoded_seq.mean(dim=1)
        elif self.config.pooling == "max":
            encoded = encoded_seq.max(dim=1)[0]
        else:
            raise ValueError(f"Unknown pooling: {self.config.pooling}")

        if return_attention:
            return encoded, attentions
        return encoded, None

    def encode(self, x: Union[np.ndarray, torch.Tensor]) -> np.ndarray:
        """
        Encode sequences to fixed-size representations.

        Used to generate features for XGBoost head.

        Args:
            x: Input sequences (batch, seq_len, features) or (seq_len, features)

        Returns:
            Encoded representations as numpy array (batch, d_model)
        """
        self.eval()

        # Handle single sequence
        if isinstance(x, np.ndarray):
            x = torch.from_numpy(x).float()

        if x.dim() == 2:
            x = x.unsqueeze(0)

        with torch.no_grad():
            x = x.to(self.device)
            encoded, _ = self.forward(x, return_attention=False)

        return encoded.cpu().numpy()

    def fit_encoder(
        self,
        X_train: np.ndarray,
        y_train: np.ndarray,
        X_val: Optional[np.ndarray] = None,
        y_val: Optional[np.ndarray] = None,
        epochs: int = 50,
        batch_size: int = 64,
        learning_rate: float = 1e-4,
        early_stopping_patience: int = 10
    ) -> Dict[str, List[float]]:
        """
        Train the attention encoder with supervised learning.

        Uses MSE loss for return prediction as pre-training task.

        Args:
            X_train: Training sequences (n_samples, seq_len, n_features)
            y_train: Training targets (n_samples,)
            X_val: Validation sequences (optional)
            y_val: Validation targets (optional)
            epochs: Number of training epochs
            batch_size: Batch size
            learning_rate: Learning rate
            early_stopping_patience: Patience for early stopping

        Returns:
            Dictionary with training history
        """
        logger.info(f"Training encoder: {epochs} epochs, batch_size={batch_size}")

        # Create prediction head for encoder training
        encoder_head = nn.Linear(self.config.d_model, 1).to(self.device)

        # Combine parameters
        all_params = list(self.attention_encoder.parameters()) + list(encoder_head.parameters())
        optimizer = torch.optim.AdamW(all_params, lr=learning_rate, weight_decay=0.01)
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer, mode='min', factor=0.5, patience=5
        )
        criterion = nn.MSELoss()

        # Convert to tensors
        X_train_t = torch.from_numpy(X_train).float()
        y_train_t = torch.from_numpy(y_train).float().view(-1, 1)

        if X_val is not None and y_val is not None:
            X_val_t = torch.from_numpy(X_val).float().to(self.device)
            y_val_t = torch.from_numpy(y_val).float().view(-1, 1).to(self.device)
        else:
            X_val_t, y_val_t = None, None

        # Training loop
        history = {'train_loss': [], 'val_loss': []}
        best_val_loss = float('inf')
        patience_counter = 0
        best_state = None

        n_batches = (len(X_train) + batch_size - 1) // batch_size

        for epoch in range(epochs):
            self.train()
            encoder_head.train()

            # Shuffle data
            indices = np.random.permutation(len(X_train))
            epoch_loss = 0.0

            for batch_idx in range(n_batches):
                start_idx = batch_idx * batch_size
                end_idx = min(start_idx + batch_size, len(X_train))
                batch_indices = indices[start_idx:end_idx]

                X_batch = X_train_t[batch_indices].to(self.device)
                y_batch = y_train_t[batch_indices].to(self.device)

                optimizer.zero_grad()

                # Forward pass
                encoded, _ = self.forward(X_batch)
                predictions = encoder_head(encoded)

                loss = criterion(predictions, y_batch)
                loss.backward()

                # Gradient clipping
                torch.nn.utils.clip_grad_norm_(all_params, max_norm=1.0)

                optimizer.step()
                epoch_loss += loss.item()

            epoch_loss /= n_batches
            history['train_loss'].append(epoch_loss)

            # Validation
            if X_val_t is not None:
                self.eval()
                encoder_head.eval()

                with torch.no_grad():
                    encoded_val, _ = self.forward(X_val_t)
                    val_preds = encoder_head(encoded_val)
                    val_loss = criterion(val_preds, y_val_t).item()

                history['val_loss'].append(val_loss)
                scheduler.step(val_loss)

                # Early stopping
                if val_loss < best_val_loss:
                    best_val_loss = val_loss
                    patience_counter = 0
                    best_state = {
                        'encoder': self.attention_encoder.state_dict(),
                        'head': encoder_head.state_dict()
                    }
                else:
                    patience_counter += 1

                if patience_counter >= early_stopping_patience:
                    logger.info(f"Early stopping at epoch {epoch + 1}")
                    break

                if (epoch + 1) % 10 == 0:
                    logger.info(
                        f"Epoch {epoch + 1}/{epochs} - "
                        f"Train Loss: {epoch_loss:.6f}, Val Loss: {val_loss:.6f}"
                    )
            else:
                if (epoch + 1) % 10 == 0:
                    logger.info(f"Epoch {epoch + 1}/{epochs} - Train Loss: {epoch_loss:.6f}")

        # Restore best model
        if best_state is not None:
            self.attention_encoder.load_state_dict(best_state['encoder'])

        self._encoder_trained = True
        logger.info(f"Encoder training complete. Best val loss: {best_val_loss:.6f}")

        return history

    def fit_xgboost(
        self,
        X_train: np.ndarray,
        y_train: np.ndarray,
        X_val: Optional[np.ndarray] = None,
        y_val: Optional[np.ndarray] = None
    ) -> Dict[str, Any]:
        """
        Train XGBoost head on encoded representations.

        Trains both direction classifier and magnitude regressor.

        Args:
            X_train: Training sequences (n_samples, seq_len, n_features)
            y_train: Training targets (n_samples,) - signed returns
            X_val: Validation sequences
            y_val: Validation targets

        Returns:
            Dictionary with training metrics
        """
        if not XGB_AVAILABLE:
            raise ImportError("XGBoost is required. Install with: pip install xgboost")

        logger.info("Training XGBoost head on encoded representations...")

        # Encode all training data
        self.eval()
        encoded_train = self.encode(X_train)

        if X_val is not None:
            encoded_val = self.encode(X_val)

        # Prepare targets
        direction_train = (y_train > 0).astype(int)
        magnitude_train = np.abs(y_train)

        # Initialize XGBoost models
        self._xgb_direction = xgb.XGBClassifier(
            n_estimators=self.config.xgb_n_estimators,
            max_depth=self.config.xgb_max_depth,
            learning_rate=self.config.xgb_learning_rate,
            subsample=self.config.xgb_subsample,
            colsample_bytree=self.config.xgb_colsample_bytree,
            reg_alpha=self.config.xgb_reg_alpha,
            reg_lambda=self.config.xgb_reg_lambda,
            min_child_weight=self.config.xgb_min_child_weight,
            gamma=self.config.xgb_gamma,
            objective='binary:logistic',
            eval_metric='logloss',
            use_label_encoder=False,
            random_state=42,
            n_jobs=-1
        )

        self._xgb_magnitude = xgb.XGBRegressor(
            n_estimators=self.config.xgb_n_estimators,
            max_depth=self.config.xgb_max_depth,
            learning_rate=self.config.xgb_learning_rate,
            subsample=self.config.xgb_subsample,
            colsample_bytree=self.config.xgb_colsample_bytree,
            reg_alpha=self.config.xgb_reg_alpha,
            reg_lambda=self.config.xgb_reg_lambda,
            min_child_weight=self.config.xgb_min_child_weight,
            gamma=self.config.xgb_gamma,
            objective='reg:squarederror',
            random_state=42,
            n_jobs=-1
        )

        # Train direction classifier
        if X_val is not None:
            direction_val = (y_val > 0).astype(int)
            self._xgb_direction.fit(
                encoded_train, direction_train,
                eval_set=[(encoded_val, direction_val)],
                verbose=False
            )
        else:
            self._xgb_direction.fit(encoded_train, direction_train)

        logger.info("Direction classifier trained")

        # Train magnitude regressor
        if X_val is not None:
            magnitude_val = np.abs(y_val)
            self._xgb_magnitude.fit(
                encoded_train, magnitude_train,
                eval_set=[(encoded_val, magnitude_val)],
                verbose=False
            )
        else:
            self._xgb_magnitude.fit(encoded_train, magnitude_train)

        logger.info("Magnitude regressor trained")

        self._is_trained = True

        # Compute training metrics
        metrics = self._compute_training_metrics(
            encoded_train, direction_train, magnitude_train,
            encoded_val if X_val is not None else None,
            direction_val if X_val is not None else None,
            magnitude_val if X_val is not None else None
        )

        return metrics

    def _compute_training_metrics(
        self,
        X_train: np.ndarray,
        y_dir_train: np.ndarray,
        y_mag_train: np.ndarray,
        X_val: Optional[np.ndarray] = None,
        y_dir_val: Optional[np.ndarray] = None,
        y_mag_val: Optional[np.ndarray] = None
    ) -> Dict[str, Any]:
        """Compute training metrics for XGBoost head."""
        from sklearn.metrics import accuracy_score, mean_squared_error, mean_absolute_error

        metrics = {}

        # Training metrics
        dir_pred_train = self._xgb_direction.predict(X_train)
        mag_pred_train = self._xgb_magnitude.predict(X_train)

        metrics['train_direction_accuracy'] = accuracy_score(y_dir_train, dir_pred_train)
        metrics['train_magnitude_mse'] = mean_squared_error(y_mag_train, mag_pred_train)
        metrics['train_magnitude_mae'] = mean_absolute_error(y_mag_train, mag_pred_train)

        # Validation metrics
        if X_val is not None:
            dir_pred_val = self._xgb_direction.predict(X_val)
            mag_pred_val = self._xgb_magnitude.predict(X_val)

            metrics['val_direction_accuracy'] = accuracy_score(y_dir_val, dir_pred_val)
            metrics['val_magnitude_mse'] = mean_squared_error(y_mag_val, mag_pred_val)
            metrics['val_magnitude_mae'] = mean_absolute_error(y_mag_val, mag_pred_val)

        logger.info(f"Training metrics: {metrics}")
        return metrics

    def predict(
        self,
        x: Union[np.ndarray, torch.Tensor],
        return_attention: bool = False
    ) -> Union[PVAPrediction, List[PVAPrediction]]:
        """
        Make prediction for input sequence(s).

        Args:
            x: Input sequence(s) (seq_len, features) or (batch, seq_len, features)
            return_attention: Whether to include attention weights

        Returns:
            PVAPrediction or list of PVAPrediction for batch input
        """
        if not self._is_trained:
            raise RuntimeError("Model not trained. Call fit_encoder() and fit_xgboost() first.")

        self.eval()

        # Handle input dimensions
        single_input = False
        if isinstance(x, np.ndarray):
            if x.ndim == 2:
                x = x[np.newaxis, ...]
                single_input = True
            x = torch.from_numpy(x).float()
        elif x.dim() == 2:
            x = x.unsqueeze(0)
            single_input = True

        x = x.to(self.device)

        with torch.no_grad():
            # Get encoded representation
            encoded, attentions = self.forward(x, return_attention=return_attention)
            encoded_np = encoded.cpu().numpy()

            # Get attention weights if requested
            attention_weights = None
            if return_attention and attentions:
                attention_weights = attentions[-1].cpu().numpy()

        # XGBoost predictions
        direction_probs = self._xgb_direction.predict_proba(encoded_np)[:, 1]
        direction = 2 * direction_probs - 1  # Map [0,1] to [-1,1]

        magnitude = self._xgb_magnitude.predict(encoded_np)

        # Compute confidence (calibrated)
        confidence = self._compute_confidence(direction_probs, magnitude)

        # Raw prediction (signed return)
        raw_pred = direction * magnitude

        # Create predictions
        predictions = []
        for i in range(len(encoded_np)):
            pred = PVAPrediction(
                direction=float(direction[i]),
                magnitude=float(magnitude[i]),
                confidence=float(confidence[i]),
                raw_prediction=float(raw_pred[i]),
                attention_weights=attention_weights[i] if attention_weights is not None else None
            )
            predictions.append(pred)

        if single_input:
            return predictions[0]
        return predictions

    def _compute_confidence(
        self,
        direction_probs: np.ndarray,
        magnitude: np.ndarray
    ) -> np.ndarray:
        """
        Compute calibrated confidence score.

        Combines direction certainty and magnitude information.

        Args:
            direction_probs: Probability of upward movement [0,1]
            magnitude: Predicted magnitude

        Returns:
            Confidence scores [0,1]
        """
        # Direction certainty (how far from 0.5)
        direction_certainty = np.abs(direction_probs - 0.5) * 2

        # Temperature-scaled softmax for calibration
        temp = self.config.confidence_temperature
        scaled_certainty = np.exp(direction_certainty / temp)
        scaled_certainty = scaled_certainty / (1 + scaled_certainty)

        # Magnitude component (higher magnitude = higher confidence, bounded)
        magnitude_factor = np.tanh(magnitude * 100)  # Scale for typical return magnitudes

        # Combined confidence
        confidence = 0.7 * scaled_certainty + 0.3 * magnitude_factor
        confidence = np.clip(confidence, 0.0, 1.0)

        return confidence

    def get_attention_scores(
        self,
        x: Union[np.ndarray, torch.Tensor],
        layer_idx: int = -1
    ) -> np.ndarray:
        """
        Get attention scores for interpretability analysis.

        Args:
            x: Input sequence(s)
            layer_idx: Which layer's attention to return (-1 = last)

        Returns:
            Attention scores of shape (batch, n_heads, seq_len, seq_len)
        """
        self.eval()

        if isinstance(x, np.ndarray):
            x = torch.from_numpy(x).float()

        if x.dim() == 2:
            x = x.unsqueeze(0)

        x = x.to(self.device)

        with torch.no_grad():
            _, attentions = self.forward(x, return_attention=True)

        if not attentions:
            raise RuntimeError("No attention weights available")

        return attentions[layer_idx].cpu().numpy()

    def save(self, path: str) -> None:
        """
        Save model to disk.

        Args:
            path: Directory path to save model
        """
        save_dir = Path(path)
        save_dir.mkdir(parents=True, exist_ok=True)

        # Save config
        config_path = save_dir / "config.json"
        with open(config_path, 'w') as f:
            json.dump(self.config.to_dict(), f, indent=2)

        # Save encoder state
        encoder_path = save_dir / "encoder.pt"
        torch.save(self.attention_encoder.state_dict(), encoder_path)

        # Save XGBoost models
        if self._xgb_direction is not None:
            direction_path = save_dir / "xgb_direction.json"
            self._xgb_direction.save_model(str(direction_path))

        if self._xgb_magnitude is not None:
            magnitude_path = save_dir / "xgb_magnitude.json"
            self._xgb_magnitude.save_model(str(magnitude_path))

        # Save training state
        state_path = save_dir / "training_state.json"
        with open(state_path, 'w') as f:
            json.dump({
                'is_trained': self._is_trained,
                'encoder_trained': self._encoder_trained
            }, f)

        logger.info(f"Model saved to {save_dir}")

    @classmethod
    def load(cls, path: str, device: str = "cpu") -> 'PVAModel':
        """
        Load model from disk.

        Args:
            path: Directory path containing saved model
            device: Device to load model to

        Returns:
            Loaded PVAModel instance
        """
        load_dir = Path(path)

        # Load config
        config_path = load_dir / "config.json"
        with open(config_path, 'r') as f:
            config_dict = json.load(f)

        config_dict['device'] = device
        config = PVAConfig.from_dict(config_dict)

        # Create model
        model = cls(config)

        # Load encoder state
        encoder_path = load_dir / "encoder.pt"
        model.attention_encoder.load_state_dict(
            torch.load(encoder_path, map_location=device)
        )

        # Load XGBoost models
        direction_path = load_dir / "xgb_direction.json"
        if direction_path.exists():
            model._xgb_direction = xgb.XGBClassifier()
            model._xgb_direction.load_model(str(direction_path))

        magnitude_path = load_dir / "xgb_magnitude.json"
        if magnitude_path.exists():
            model._xgb_magnitude = xgb.XGBRegressor()
            model._xgb_magnitude.load_model(str(magnitude_path))

        # Load training state
        state_path = load_dir / "training_state.json"
        if state_path.exists():
            with open(state_path, 'r') as f:
                state = json.load(f)
                model._is_trained = state.get('is_trained', False)
                model._encoder_trained = state.get('encoder_trained', False)

        logger.info(f"Model loaded from {load_dir}")
        return model


def create_pva_model(
    input_features: int = 30,
    sequence_length: int = 100,
    device: str = "cpu"
) -> PVAModel:
    """
    Convenience function to create a PVA model with default configuration.

    Args:
        input_features: Number of input features per timestep
        sequence_length: Length of input sequences
        device: Device to use ('cpu' or 'cuda')

    Returns:
        Initialized PVAModel
    """
    config = PVAConfig(
        input_features=input_features,
        sequence_length=sequence_length,
        device=device
    )
    return PVAModel(config)


if __name__ == "__main__":
    # Test the PVA model
    print("Testing PVAModel...")

    # Create sample data
    np.random.seed(42)
    n_samples = 1000
    seq_length = 100
    n_features = 30

    X = np.random.randn(n_samples, seq_length, n_features).astype(np.float32)
    y = np.random.randn(n_samples).astype(np.float32) * 0.01  # Small returns

    # Split data
    split_idx = int(n_samples * 0.8)
    X_train, X_val = X[:split_idx], X[split_idx:]
    y_train, y_val = y[:split_idx], y[split_idx:]

    print(f"\nData shapes:")
    print(f"  X_train: {X_train.shape}")
    print(f"  X_val: {X_val.shape}")

    # Create model
    config = PVAConfig(
        input_features=n_features,
        sequence_length=seq_length,
        device="cpu",
        n_layers=2,  # Smaller for testing
        d_model=64
    )
    model = PVAModel(config)

    # Train encoder (short training for test)
    print("\nTraining encoder...")
    history = model.fit_encoder(
        X_train, y_train, X_val, y_val,
        epochs=5, batch_size=32
    )
    print(f"  Final train loss: {history['train_loss'][-1]:.6f}")

    # Train XGBoost head
    print("\nTraining XGBoost head...")
    metrics = model.fit_xgboost(X_train, y_train, X_val, y_val)
    print(f"  Direction accuracy: {metrics.get('val_direction_accuracy', 'N/A'):.4f}")
    print(f"  Magnitude MAE: {metrics.get('val_magnitude_mae', 'N/A'):.6f}")

    # Make predictions
    print("\nMaking predictions...")
    predictions = model.predict(X_val[:5])
    for i, pred in enumerate(predictions):
        print(f"  Sample {i}: dir={pred.direction:.3f}, mag={pred.magnitude:.6f}, conf={pred.confidence:.3f}")

    # Get attention scores
    print("\nGetting attention scores...")
    attention = model.get_attention_scores(X_val[:1])
    print(f"  Attention shape: {attention.shape}")

    # Test save/load
    print("\nTesting save/load...")
    save_path = "/tmp/pva_model_test"
    model.save(save_path)

    loaded_model = PVAModel.load(save_path)
    loaded_pred = loaded_model.predict(X_val[0])
    print(f"  Loaded model prediction: dir={loaded_pred.direction:.3f}")

    print("\nTest complete!")
