"""Example 4: Upload a TensorFlow/Keras Model with Evaluation Callback.

Trains a simple Keras model on a synthetic dataset, logs metrics per epoch
using TfMLflowEpochLogger, and then uploads the trained model to MLflow.

Prerequisites:
    pip install mlflow-utils[tensorflow]

Run:
    python examples/04_upload_tensorflow.py
"""

import numpy as np
import tensorflow as tf
from dgtw_mlflow_utils import (
    ModelUploader,
    ModelVariant,
    RunAction,
    Stage,
    TfMLflowEpochLogger,
    UploadConfig,
)


def create_simple_model(input_dim: int = 10, output_dim: int = 3):
    """Create a simple Keras classification model."""
    model = tf.keras.Sequential(
        [
            tf.keras.layers.Dense(64, activation="relu", input_shape=(input_dim,)),
            tf.keras.layers.Dropout(0.3),
            tf.keras.layers.Dense(32, activation="relu"),
            tf.keras.layers.Dense(output_dim, activation="softmax"),
        ]
    )
    model.compile(
        optimizer="adam",
        loss="sparse_categorical_crossentropy",
        metrics=["accuracy"],
    )
    return model


def main():
    """Main function to execute the example."""
    # ------------------------------------------------------------------
    # 1) Create synthetic dataset
    # ------------------------------------------------------------------
    print("Generating synthetic dataset...")
    np.random.seed(42)
    X_train = np.random.randn(1000, 10).astype(np.float32)
    y_train = np.random.randint(0, 3, size=1000)
    X_test = np.random.randn(200, 10).astype(np.float32)
    y_test = np.random.randint(0, 3, size=200)

    # ------------------------------------------------------------------
    # 2) Train with TfMLflowEpochLogger callback
    # ------------------------------------------------------------------
    print("Training model with MLflow epoch logging...\n")

    model = create_simple_model(input_dim=10, output_dim=3)

    # The callback will log training metrics to MLflow at each epoch
    # Set enable_mlflow=True to activate tracking (False for local-only)
    mlflow_callback = TfMLflowEpochLogger(
        enable_mlflow=False,  # set True when MLflow server is running
        experiment_name="TAT_Classification_Research",
        run_name="Train_SimpleNN",
        # Additional params are logged automatically:
        log_params={
            "model_type": "Sequential",
            "optimizer": "adam",
            "epochs": 10,
        },
    )

    model.fit(
        X_train,
        y_train,
        validation_data=(X_test, y_test),
        epochs=10,
        batch_size=32,
        callbacks=[mlflow_callback],
        verbose=1,
    )

    # Evaluate
    loss, accuracy = model.evaluate(X_test, y_test, verbose=0)
    print(f"\nTest Loss: {loss:.4f}")
    print(f"Test Accuracy: {accuracy:.4f}")

    # ------------------------------------------------------------------
    # 3) Upload the trained model to MLflow
    # ------------------------------------------------------------------
    config = UploadConfig(
        architecture="SimpleNN",  # custom architecture string
        variant=ModelVariant.BASELINE,
        component="Classification",
        stage=Stage.RESEARCH,
        use_case="SyntheticDemo",
        action=RunAction.UPLOAD,
        tracking_uri="http://127.0.0.1:5001",
        params={
            "test_loss": round(loss, 4),
            "test_accuracy": round(accuracy, 4),
            "epochs": 10,
            "batch_size": 32,
        },
        tags={
            "framework": "tensorflow",
            "developer": "ai-team",
        },
    )

    uploader = ModelUploader(config)
    print()
    print(uploader.summary())
    print()

    result = uploader.upload_tensorflow(tf_model=model)

    print("Upload successful!")
    print(f"  Run ID             : {result['run_id']}")
    print(f"  Model URI          : {result['model_uri']}")
    print(f"  Registered Model   : {result['registered_model_name']}")


if __name__ == "__main__":
    main()
