"""
Training utilities for CNN models
"""

import numpy as np
from tensorflow import keras
from typing import Optional, Dict, Any, List
from kwslib.models.cnn import build_cnn_model
from .callbacks import MLFlowCallback


class CNNTrainer:
    """Helper class for training CNN models"""
    
    def __init__(self, api_client=None, use_mlflow: bool = False, mlflow_run_id: Optional[str] = None):
        """
        Initialize CNN trainer.
        
        Args:
            api_client: Optional API client for MLFlow logging
            use_mlflow: Whether to use MLFlow logging
            mlflow_run_id: MLFlow run ID if using MLFlow
        """
        self.api = api_client
        self.use_mlflow = use_mlflow
        self.mlflow_run_id = mlflow_run_id
        self.model = None
    
    def build_model(
        self,
        input_shape: tuple,
        num_classes: int,
        learning_rate: float = 0.001
    ) -> keras.Model:
        """
        Build CNN model.
        
        Args:
            input_shape: Input shape (excluding batch dimension)
            num_classes: Number of classes
            learning_rate: Learning rate for optimizer
        
        Returns:
            Compiled Keras model
        """
        print(f"\n=== Building CNN model ===")
        print(f"Input shape: {input_shape}")
        print(f"Number of classes: {num_classes}")
        
        # Log hyperparameters to MLFlow if enabled
        if self.use_mlflow and self.mlflow_run_id and self.api:
            try:
                self.api.mlflow.log_params(
                    run_id=self.mlflow_run_id,
                    params={
                        "learning_rate": str(learning_rate),
                        "batch_size": "32",
                        "epochs": "100",
                        "num_classes": str(num_classes),
                        "input_shape": str(input_shape),
                        "optimizer": "Adam",
                        "loss": "sparse_categorical_crossentropy",
                    }
                )
                print("Logged hyperparameters to MLFlow")
            except Exception as e:
                print(f"Warning: Failed to log params to MLFlow: {e}")
        
        model = build_cnn_model(input_shape, num_classes)
        model.summary()
        self.model = model
        return model
    
    def train(
        self,
        X_train: np.ndarray,
        y_train: np.ndarray,
        batch_size: int = 32,
        epochs: int = 100,
        validation_split: float = 0.2,
        callbacks: Optional[List[keras.callbacks.Callback]] = None
    ) -> keras.callbacks.History:
        """
        Train the model.
        
        Args:
            X_train: Training features
            y_train: Training labels
            batch_size: Batch size
            epochs: Number of epochs
            validation_split: Validation split ratio
            callbacks: Optional list of callbacks
        
        Returns:
            Training history
        """
        print("\n=== Training model ===")
        
        # Default callbacks
        default_callbacks = [
            keras.callbacks.EarlyStopping(
                monitor='val_loss',
                patience=10,
                restore_best_weights=True
            ),
            keras.callbacks.ReduceLROnPlateau(
                monitor='val_loss',
                factor=0.5,
                patience=5,
                min_lr=1e-7
            ),
            keras.callbacks.ModelCheckpoint(
                'best_model.h5',
                monitor='val_loss',
                save_best_only=True
            )
        ]
        
        # Add MLFlow callback if enabled
        if self.use_mlflow and self.mlflow_run_id and self.api:
            default_callbacks.append(
                MLFlowCallback(self.api.mlflow, self.mlflow_run_id)
            )
        
        # Combine with user-provided callbacks
        if callbacks:
            all_callbacks = default_callbacks + callbacks
        else:
            all_callbacks = default_callbacks
        
        if self.model is None:
            raise ValueError("Model not built. Call build_model() first.")
        
        history = self.model.fit(
            X_train, y_train,
            batch_size=batch_size,
            epochs=epochs,
            validation_split=validation_split,
            callbacks=all_callbacks,
            verbose=1
        )
        
        return history
