"""Example 6: Evaluator — TfMLflowEpochLogger & evaluate_model_tf.

Demonstrates how to use the evaluator module for:

1. compute_retrieval_metrics — Low-level metric computation from logits.
2. TfMLflowEpochLogger  — A Keras callback that logs training/validation
   metrics to MLflow at each epoch.
3. evaluate_model_tf     — A standalone evaluation loop that computes
   retrieval metrics (Recall@K, MRR, NDCG) and optionally logs to MLflow.

The example creates a simple Two-Tower model that produces a similarity
matrix (logits) from in-batch negatives, which is the expected input
format for the evaluator.

Prerequisites:
    pip install "mlflow-utils[tensorflow]"

Run:
    python examples/06_evaluator.py
"""

import numpy as np
import tensorflow as tf
from dgtw_mlflow_utils import (
    NumpyMLflowLogger,
    TfMLflowEpochLogger,
    compute_retrieval_metrics,
    compute_retrieval_metrics_numpy,
    evaluate_model_numpy,
    evaluate_model_tf,
)

# ======================================================================
# 1) Build a simple Two-Tower model (for demonstration)
# ======================================================================


class SimpleTwoTower(tf.keras.Model):
    """A minimal Two-Tower model that produces a (batch_size, batch_size) similarity matrix from query and item embeddings.

    This is the format expected by evaluate_model_tf() and
    compute_retrieval_metrics().
    """

    def __init__(self, vocab_size: int = 100, embedding_dim: int = 32):
        """Initialize the SimpleTwoTower model.

        Args:
            vocab_size: Size of the vocabulary for embeddings. Defaults to 100.
            embedding_dim: Dimension of the embedding vectors. Defaults to 32.
        """
        super().__init__()
        self.query_embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
        self.item_embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
        self.query_dense = tf.keras.layers.Dense(embedding_dim, activation="relu")
        self.item_dense = tf.keras.layers.Dense(embedding_dim, activation="relu")

    def call(self, inputs, training=False):
        """Args: inputs: dict with "query_id" and "item_id" tensors.

        Returns: logits: (batch_size, batch_size) similarity matrix.
        """
        if isinstance(inputs, dict):
            query_ids = inputs["query_id"]
            item_ids = inputs["item_id"]
        else:
            query_ids, item_ids = inputs

        query_emb = self.query_dense(self.query_embedding(query_ids))
        item_emb = self.item_dense(self.item_embedding(item_ids))

        # Dot product → (batch_size, batch_size) similarity matrix
        logits = tf.matmul(query_emb, item_emb, transpose_b=True)
        return logits

    def train_step(self, data):
        """Custom train step for Two-Tower with in-batch negatives."""
        with tf.GradientTape() as tape:
            logits = self(data, training=True)
            batch_size = tf.shape(logits)[0]
            labels = tf.eye(batch_size)
            loss = tf.keras.losses.categorical_crossentropy(labels, logits, from_logits=True)
            loss = tf.reduce_mean(loss)

        gradients = tape.gradient(loss, self.trainable_variables)
        self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
        return {"loss": loss}

    def test_step(self, data):
        """Custom test step for Two-Tower with in-batch negatives."""
        logits = self(data, training=False)
        batch_size = tf.shape(logits)[0]
        labels = tf.eye(batch_size)
        loss = tf.keras.losses.categorical_crossentropy(labels, logits, from_logits=True)
        loss = tf.reduce_mean(loss)
        return {"loss": loss}


def create_synthetic_dataset(
    num_samples: int = 500,
    vocab_size: int = 100,
    batch_size: int = 16,
) -> tf.data.Dataset:
    """Create a synthetic dataset of (query_id, item_id) pairs."""
    query_ids = np.random.randint(0, vocab_size, size=num_samples)
    item_ids = np.random.randint(0, vocab_size, size=num_samples)

    dataset = tf.data.Dataset.from_tensor_slices(
        {
            "query_id": query_ids,
            "item_id": item_ids,
        }
    )
    dataset = dataset.batch(batch_size, drop_remainder=True)
    dataset = dataset.prefetch(tf.data.AUTOTUNE)
    return dataset


# ======================================================================
# PART A: Using compute_retrieval_metrics directly
# ======================================================================


def example_compute_metrics():
    """Directly compute retrieval metrics from a logits matrix."""
    print("=" * 60)
    print("  Part A: compute_retrieval_metrics (low-level)")
    print("=" * 60)

    # Simulate a (8, 8) logits matrix from a batch of 8 items
    # Ground truth: item i is the correct match for query i (diagonal)
    batch_size = 8
    logits = tf.random.normal((batch_size, batch_size))

    # Make the diagonal scores higher so the model "gets it right"
    logits = logits + tf.eye(batch_size) * 3.0

    metrics = compute_retrieval_metrics(logits, k=3)

    print(f"  Recall@3 : {metrics['recall_3'].numpy():.4f}")
    print(f"  MRR@3    : {metrics['mrr_3'].numpy():.4f}")
    print(f"  NDCG@3   : {metrics['ndcg_3'].numpy():.4f}")

    # Try with worse predictions (no diagonal boost)
    bad_logits = tf.random.normal((batch_size, batch_size))
    bad_metrics = compute_retrieval_metrics(bad_logits, k=3)

    print()
    print("  With random (bad) predictions:")
    print(f"  Recall@3 : {bad_metrics['recall_3'].numpy():.4f}")
    print(f"  MRR@3    : {bad_metrics['mrr_3'].numpy():.4f}")
    print(f"  NDCG@3   : {bad_metrics['ndcg_3'].numpy():.4f}")
    print()


# ======================================================================
# PART B: Using TfMLflowEpochLogger during training
# ======================================================================


def example_epoch_logger():
    """Train a model with TfMLflowEpochLogger callback."""
    print("=" * 60)
    print("  Part B: TfMLflowEpochLogger (Keras callback)")
    print("=" * 60)

    # Create model and data
    model = SimpleTwoTower(vocab_size=100, embedding_dim=32)
    train_ds = create_synthetic_dataset(num_samples=320, batch_size=16)
    val_ds = create_synthetic_dataset(num_samples=80, batch_size=16)

    # Compile with optimizer only — loss is handled in train_step
    model.compile(optimizer="adam")

    # Create the callback
    # Set enable_mlflow=False for this demo (no MLflow server needed)
    # In production, set enable_mlflow=True and provide experiment_name
    callback = TfMLflowEpochLogger(
        enable_mlflow=False,  # ← set True when MLflow is running
        experiment_name="TAT_Recommendation_Research",
        run_name="Train_TwoTower",
        # Any extra kwargs are logged as MLflow params:
        log_params={
            "embedding_dim": 32,
            "optimizer": "adam",
            "batch_size": 16,
        },
    )

    print("\n  Training with TfMLflowEpochLogger...\n")

    # The callback handles the full MLflow lifecycle:
    #   on_train_begin → starts MLflow run + logs params
    #   on_epoch_end   → logs training loss/metrics per epoch
    #   on_test_end    → logs validation metrics
    #   on_train_end   → ends MLflow run
    model.fit(
        train_ds,
        validation_data=val_ds,
        epochs=3,
        callbacks=[callback],
        verbose=1,
    )

    print("\n  Training complete!")
    print("  → When enable_mlflow=True, all epoch metrics are logged to MLflow.\n")

    return model  # return the trained model for Part C


# ======================================================================
# PART C: Using evaluate_model_tf for standalone evaluation
# ======================================================================


def example_evaluate_model_tf(model=None):
    """Run standalone evaluation with evaluate_model_tf()."""
    print("=" * 60)
    print("  Part C: evaluate_model_tf (standalone evaluation)")
    print("=" * 60)

    # Use the provided model or create a new one
    if model is None:
        model = SimpleTwoTower(vocab_size=100, embedding_dim=32)
        # Warm up the model with one forward pass
        dummy = create_synthetic_dataset(num_samples=16, batch_size=16)
        for batch in dummy:
            model(batch, training=False)
            break

    # Create evaluation dataset
    eval_ds = create_synthetic_dataset(num_samples=160, batch_size=16)

    print("\n  Running evaluate_model_tf()...\n")

    # evaluate_model_tf processes the dataset batch-by-batch,
    # computes Recall@K, MRR, NDCG, Loss per batch, and averages them.
    # Set enable_mlflow=False for this demo.
    metrics = evaluate_model_tf(
        model=model,
        dataset=eval_ds,
        k=5,  # Recall@5
        enable_mlflow=False,  # ← set True when MLflow is running
        experiment_name="TAT_Recommendation_Research",
        run_name="Final_Evaluation",
        # Extra params logged to MLflow:
        log_params={
            "dataset_size": 160,
            "model_version": "v1",
        },
    )

    print("\n  Returned metrics dict:")
    for key, val in metrics.items():
        print(f"    {key}: {val:.4f}")

    print("\n  → When enable_mlflow=True, these metrics are logged to MLflow.\n")


# ======================================================================
# PART D: Using compute_retrieval_metrics_numpy directly
# ======================================================================


def example_compute_metrics_numpy():
    """Directly compute retrieval metrics from a numpy logits matrix."""
    print("=" * 60)
    print("  Part D: compute_retrieval_metrics_numpy (low-level)")
    print("=" * 60)

    # Simulate a (8, 8) logits matrix from a batch of 8 items
    batch_size = 8
    logits = np.random.normal(size=(batch_size, batch_size))

    # Make the diagonal scores higher so the model "gets it right"
    logits = logits + np.eye(batch_size) * 3.0

    metrics = compute_retrieval_metrics_numpy(logits, k=3)

    print(f"  Recall@3 : {metrics['recall_3']:.4f}")
    print(f"  MRR@3    : {metrics['mrr_3']:.4f}")
    print(f"  NDCG@3   : {metrics['ndcg_3']:.4f}")

    # Try with worse predictions (no diagonal boost)
    bad_logits = np.random.normal(size=(batch_size, batch_size))
    bad_metrics = compute_retrieval_metrics_numpy(bad_logits, k=3)

    print()
    print("  With random (bad) predictions (NumPy):")
    print(f"  Recall@3 : {bad_metrics['recall_3']:.4f}")
    print(f"  MRR@3    : {bad_metrics['mrr_3']:.4f}")
    print(f"  NDCG@3   : {bad_metrics['ndcg_3']:.4f}")
    print()


# ======================================================================
# PART E: Using evaluate_model_numpy for standalone evaluation
# ======================================================================


def example_evaluate_model_numpy():
    """Run standalone evaluation with evaluate_model_numpy()."""
    print("=" * 60)
    print("  Part E: evaluate_model_numpy (standalone NumPy evaluation)")
    print("=" * 60)

    # Create evaluation dataset
    eval_ds = create_synthetic_dataset(num_samples=160, batch_size=16)

    # Convert the tf.data.Dataset to an iterable of dicts for the example
    def dataset_generator():
        for batch in eval_ds:
            yield {
                "query_id": batch["query_id"].numpy(),
                "item_id": batch["item_id"].numpy(),
            }

    # Create a mock predict function (normally this would be your ML model inference)
    def mock_predict_fn(batch):
        # Return a random similarity matrix (batch_size, batch_size)
        # We add some signal on the diagonal to simulate a slightly trained model
        b_size = len(batch["query_id"])
        logits = np.random.normal(size=(b_size, b_size))
        logits += np.eye(b_size) * 2.0
        return logits

    print("\n  Running evaluate_model_numpy()...\n")

    # evaluate_model_numpy processes the dataset batch-by-batch
    # Set enable_mlflow=False for this demo.
    metrics = evaluate_model_numpy(
        predict_fn=mock_predict_fn,
        dataset=dataset_generator(),
        k=5,  # Recall@5
        enable_mlflow=False,  # ← set True when MLflow is running
        experiment_name="TAT_Recommendation_Research",
        run_name="Final_Evaluation_Numpy",
        # Extra params logged to MLflow:
        log_params={
            "dataset_size": 160,
            "model_version": "mock_v1",
        },
    )

    print("\n  Returned metrics dict (NumPy):")
    for key, val in metrics.items():
        print(f"    {key}: {val:.4f}")

    print("\n  → When enable_mlflow=True, these metrics are logged to MLflow.\n")


# ======================================================================
# PART F: Using NumpyMLflowLogger during a custom training loop
# ======================================================================


def example_numpy_epoch_logger():
    """Simulate a custom training loop with NumpyMLflowLogger."""
    print("=" * 60)
    print("  Part F: NumpyMLflowLogger (custom training loop)")
    print("=" * 60)

    # Create the logger
    logger = NumpyMLflowLogger(
        enable_mlflow=False,  # ← set True when MLflow is running
        experiment_name="TAT_Recommendation_Research",
        run_name="Train_TwoTower_Numpy",
        log_params={
            "embedding_dim": 32,
            "optimizer": "adam",
            "batch_size": 16,
        },
    )

    print("\n  Training with NumpyMLflowLogger...\n")

    logger.on_train_begin()

    # Simulate 3 epochs of training
    for epoch in range(3):
        # ... your custom training code here ...
        # Simulate some loss and metrics
        metrics = {
            "loss": 2.5 - (epoch * 0.5),
            "recall_5": 0.4 + (epoch * 0.1),
        }
        logger.on_epoch_end(epoch, metrics)

    logger.on_train_end()

    print("\n  Training complete!")
    print("  → When enable_mlflow=True, all epoch metrics are logged to MLflow.\n")


# ======================================================================
# Main
# ======================================================================


def main():
    """Run all evaluator examples demonstrating metrics computation and logging."""
    print()
    print("  dgtw_mlflow_utils Evaluator Examples")
    print("  ================================\n")

    # Part A: Low-level metric computation
    example_compute_metrics()

    # Part B: Keras training with epoch logging
    trained_model = example_epoch_logger()

    # Part C: Standalone evaluation (using the trained model)
    example_evaluate_model_tf(model=trained_model)

    # Part D: Low-level metric computation (NumPy)
    example_compute_metrics_numpy()

    # Part E: Standalone evaluation (NumPy)
    example_evaluate_model_numpy()

    # Part F: Custom training loop (NumPy)
    example_numpy_epoch_logger()

    print("=" * 60)
    print("  All done! ✓")
    print("=" * 60)
    print()
    print("  To enable MLflow tracking, set enable_mlflow=True and")
    print("  ensure your MLflow server is running.")
    print()


if __name__ == "__main__":
    main()
