"""Training orchestration and utilities."""

import os
from pathlib import Path
from typing import Optional, Dict
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm
import json

from .config import TrainingConfig
from .evaluator import evaluate_model


class EarlyStopping:
    """Early stopping to stop training when validation loss doesn't improve."""

    def __init__(self, patience: int = 10, min_delta: float = 1e-4):
        """Initialize early stopping.

        Args:
            patience: Number of epochs to wait before stopping
            min_delta: Minimum change to qualify as improvement
        """
        self.patience = patience
        self.min_delta = min_delta
        self.counter = 0
        self.best_loss = None
        self.should_stop = False

    def __call__(self, val_loss: float) -> bool:
        """Check if training should stop.

        Args:
            val_loss: Current validation loss

        Returns:
            True if training should stop
        """
        if self.best_loss is None:
            self.best_loss = val_loss
        elif val_loss > self.best_loss - self.min_delta:
            self.counter += 1
            if self.counter >= self.patience:
                self.should_stop = True
        else:
            self.best_loss = val_loss
            self.counter = 0

        return self.should_stop


class Trainer:
    """Model trainer with support for MPS/CUDA/CPU."""

    def __init__(
        self,
        model: nn.Module,
        config: TrainingConfig,
        task: str = 'classification',
        pos_weight: float = None  # 添加类别权重参数
    ):
        """Initialize trainer.

        Args:
            model: PyTorch model to train
            config: Training configuration
            task: Task type ('classification' or 'regression')
            pos_weight: Positive class weight for imbalanced data
        """
        self.model = model
        self.config = config
        self.task = task
        self.pos_weight = pos_weight

        # Set device
        self.device = config.get_device()
        self.model = self.model.to(self.device)
        print(f"Training on device: {self.device}")

        # Loss function (with class weight if provided)
        if hasattr(model, 'get_loss_function'):
            self.criterion = model.get_loss_function(pos_weight=pos_weight, device=self.device)
        else:
            if task == 'classification':
                if pos_weight is not None:
                    # 使用float32以兼容MPS
                    pos_weight_tensor = torch.tensor([pos_weight], dtype=torch.float32).to(self.device)
                    self.criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight_tensor)
                else:
                    self.criterion = nn.BCELoss()
            else:
                self.criterion = nn.MSELoss()

        if pos_weight is not None:
            print(f"Using weighted loss - positive class weight: {pos_weight:.2f}")

        # Optimizer
        self.optimizer = torch.optim.AdamW(
            model.parameters(),
            lr=config.learning_rate,
            weight_decay=config.weight_decay
        )

        # Learning rate scheduler
        if config.lr_scheduler == 'plateau':
            self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
                self.optimizer,
                mode='min',
                factor=config.lr_factor,
                patience=config.lr_patience,
                verbose=True
            )
        elif config.lr_scheduler == 'cosine':
            self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
                self.optimizer,
                T_max=config.max_epochs
            )
        else:
            self.scheduler = None

        # Early stopping
        self.early_stopping = EarlyStopping(
            patience=config.early_stopping_patience,
            min_delta=config.early_stopping_min_delta
        )

        # Training history
        self.history = {
            'train_loss': [],
            'val_loss': [],
            'learning_rates': []
        }

        # Best model tracking
        self.best_val_loss = float('inf')
        self.best_model_path = None

    def train_epoch(self, train_loader: DataLoader) -> float:
        """Train for one epoch.

        Args:
            train_loader: Training data loader

        Returns:
            Average training loss
        """
        self.model.train()
        total_loss = 0.0
        num_batches = 0

        pbar = tqdm(train_loader, desc='Training')
        for batch_X, batch_y in pbar:
            batch_X = batch_X.to(self.device)
            batch_y = batch_y.to(self.device)

            # Forward pass
            self.optimizer.zero_grad()
            outputs = self.model(batch_X).squeeze()

            # Handle different output shapes
            if self.task == 'classification':
                batch_y = batch_y.float()

            loss = self.criterion(outputs, batch_y)

            # Backward pass
            loss.backward()
            self.optimizer.step()

            # Track loss
            total_loss += loss.item()
            num_batches += 1

            # Update progress bar
            pbar.set_postfix({'loss': f'{loss.item():.4f}'})

        return total_loss / num_batches

    def validate(self, val_loader: DataLoader) -> float:
        """Validate the model.

        Args:
            val_loader: Validation data loader

        Returns:
            Average validation loss
        """
        self.model.eval()
        total_loss = 0.0
        num_batches = 0

        with torch.no_grad():
            for batch_X, batch_y in val_loader:
                batch_X = batch_X.to(self.device)
                batch_y = batch_y.to(self.device)

                outputs = self.model(batch_X).squeeze()

                if self.task == 'classification':
                    batch_y = batch_y.float()

                loss = self.criterion(outputs, batch_y)

                total_loss += loss.item()
                num_batches += 1

        return total_loss / num_batches

    def fit(
        self,
        train_loader: DataLoader,
        val_loader: DataLoader,
        save_dir: Optional[str] = None
    ) -> Dict:
        """Train the model.

        Args:
            train_loader: Training data loader
            val_loader: Validation data loader
            save_dir: Directory to save checkpoints

        Returns:
            Training history
        """
        if save_dir is None:
            save_dir = self.config.checkpoint_dir

        save_dir = Path(save_dir)
        save_dir.mkdir(parents=True, exist_ok=True)

        print(f"Starting training for {self.config.max_epochs} epochs...")
        print(f"Model parameters: {self.model.count_parameters()}")

        for epoch in range(self.config.max_epochs):
            # Train
            train_loss = self.train_epoch(train_loader)

            # Validate
            val_loss = self.validate(val_loader)

            # Update history
            self.history['train_loss'].append(train_loss)
            self.history['val_loss'].append(val_loss)
            current_lr = self.optimizer.param_groups[0]['lr']
            self.history['learning_rates'].append(current_lr)

            # Print progress
            print(f"Epoch {epoch + 1}/{self.config.max_epochs}")
            print(f"  Train loss: {train_loss:.4f}")
            print(f"  Val loss: {val_loss:.4f}")
            print(f"  LR: {current_lr:.6f}")

            # Save best model
            if val_loss < self.best_val_loss:
                self.best_val_loss = val_loss
                if self.config.save_best:
                    self.best_model_path = save_dir / 'best_model.pt'
                    self.model.save(str(self.best_model_path))
                    print(f"  Saved best model (val_loss: {val_loss:.4f})")

            # Learning rate scheduling
            if self.scheduler is not None:
                if self.config.lr_scheduler == 'plateau':
                    self.scheduler.step(val_loss)
                else:
                    self.scheduler.step()

            # Early stopping
            if self.early_stopping(val_loss):
                print(f"Early stopping at epoch {epoch + 1}")
                break

        # Save final model
        final_model_path = save_dir / 'final_model.pt'
        self.model.save(str(final_model_path))

        # Save training history
        history_path = save_dir / 'training_history.json'
        with open(history_path, 'w') as f:
            json.dump(self.history, f, indent=2)

        print(f"\nTraining completed!")
        print(f"Best validation loss: {self.best_val_loss:.4f}")
        print(f"Best model saved to: {self.best_model_path}")

        return self.history

    def evaluate(self, test_loader: DataLoader) -> Dict[str, float]:
        """Evaluate the model on test data.

        Args:
            test_loader: Test data loader

        Returns:
            Dictionary of evaluation metrics
        """
        # Load best model if available
        if self.best_model_path is not None and self.best_model_path.exists():
            print(f"Loading best model from {self.best_model_path}")
            checkpoint = torch.load(self.best_model_path, map_location=self.device)
            self.model.load_state_dict(checkpoint['model_state_dict'])

        # Evaluate
        metrics = evaluate_model(self.model, test_loader, self.task, self.device)

        print("\nTest Metrics:")
        for key, value in metrics.items():
            print(f"  {key}: {value:.4f}")

        return metrics


if __name__ == "__main__":
    # Test the trainer
    from ..models.dlinear.model import DLinear, DLinearConfig
    from ..data.fetcher import BinanceDataFetcher
    from ..data.preprocessor import CryptoPreprocessor, split_data
    from ..data.dataset import create_dataloaders

    # Fetch data
    print("Fetching data...")
    fetcher = BinanceDataFetcher()
    df = fetcher.fetch_ohlcv('BTC', '1h')

    # Preprocess
    print("\nPreprocessing...")
    preprocessor = CryptoPreprocessor(lookback=168)
    X, y, _ = preprocessor.preprocess_pipeline(df, task='classification')

    # Split
    train_data, val_data, test_data = split_data(X, y)

    # Create dataloaders
    train_loader, val_loader, test_loader = create_dataloaders(
        train_data, val_data, test_data,
        batch_size=32
    )

    # Create model
    print("\nCreating model...")
    config = DLinearConfig(
        input_dim=X.shape[-1],
        lookback=168,
        task='classification'
    )
    model = DLinear(config)

    # Train
    train_config = TrainingConfig(
        batch_size=32,
        max_epochs=5,  # Short test
        early_stopping_patience=3
    )

    trainer = Trainer(model, train_config, task='classification')
    history = trainer.fit(train_loader, val_loader, save_dir='/tmp/test_training')

    # Evaluate
    metrics = trainer.evaluate(test_loader)