"""
Keras integration for Armory SDK.

Provides automatic logging of training metrics to Armory.
"""

from typing import Any, Dict, Optional

import armory

# Lazy import
try:
    from tensorflow import keras

    HAS_KERAS = True
except ImportError:
    try:
        import keras

        HAS_KERAS = True
    except ImportError:
        HAS_KERAS = False
        keras = None


class ArmoryKerasCallback:
    """Keras callback that logs training metrics to Armory.

    Usage:
        import armory
        from armory.integrations import ArmoryKerasCallback

        with armory.start_session(project="my-project", goal="train cnn"):
            model.fit(
                x_train, y_train,
                callbacks=[ArmoryKerasCallback()],
            )
    """

    def __init__(
        self,
        log_batch_metrics: bool = False,
        batch_log_frequency: int = 100,
    ):
        """Initialize the callback.

        Args:
            log_batch_metrics: Whether to log metrics per batch
            batch_log_frequency: How often to log batch metrics
        """
        if not HAS_KERAS:
            raise ImportError(
                "keras or tensorflow is required for ArmoryKerasCallback. "
                "Install with: pip install tensorflow or pip install keras"
            )

        self.log_batch_metrics = log_batch_metrics
        self.batch_log_frequency = batch_log_frequency
        self._global_step = 0

    def set_params(self, params: Dict[str, Any]) -> None:
        """Called when training starts."""
        armory.log_params({
            "epochs": params.get("epochs"),
            "steps": params.get("steps"),
            "batch_size": params.get("batch_size"),
        })

    def set_model(self, model) -> None:
        """Called when model is set."""
        self.model = model

        # Log model summary
        try:
            total_params = model.count_params()
            armory.log_param("total_parameters", total_params)
        except Exception:
            pass

    def on_epoch_begin(self, epoch: int, logs: Optional[Dict] = None) -> None:
        """Called at the beginning of each epoch."""
        pass

    def on_epoch_end(self, epoch: int, logs: Optional[Dict] = None) -> None:
        """Called at the end of each epoch."""
        if logs:
            # Log all metrics
            metrics = {k: float(v) for k, v in logs.items() if v is not None}
            armory.log_metrics(metrics, step=epoch)

    def on_batch_end(self, batch: int, logs: Optional[Dict] = None) -> None:
        """Called at the end of each batch."""
        self._global_step += 1

        if self.log_batch_metrics and self._global_step % self.batch_log_frequency == 0:
            if logs:
                metrics = {
                    f"batch_{k}": float(v)
                    for k, v in logs.items()
                    if v is not None
                }
                armory.log_metrics(metrics, step=self._global_step)

    def on_train_begin(self, logs: Optional[Dict] = None) -> None:
        """Called at the beginning of training."""
        self._global_step = 0

    def on_train_end(self, logs: Optional[Dict] = None) -> None:
        """Called at the end of training."""
        armory.log_metric("total_batches", self._global_step)


# Make it work as both a class and a Keras callback
if HAS_KERAS:
    # Create a proper Keras callback by inheriting
    class _ArmoryKerasCallbackImpl(keras.callbacks.Callback, ArmoryKerasCallback):
        def __init__(self, *args, **kwargs):
            keras.callbacks.Callback.__init__(self)
            ArmoryKerasCallback.__init__(self, *args, **kwargs)

    # Replace the class with the proper implementation
    ArmoryKerasCallback = _ArmoryKerasCallbackImpl
