"""
.. currentmodule:: neuralk

.. _example_on_premise:

Using Seldon with On-Premise Server
====================================

The :class:`Seldon` model provides a scikit-learn compatible interface
for using Neuralk's In-Context Learning model with an on-premise or self-hosted
Seldon server. Use the ``host`` parameter to specify your server URL.

.. note::

    For this example to run, the environment variable ``HOST`` must be set
    with your Seldon server URL.
"""

# %%
# Simple example on toy data
# --------------------------
#
# We start by using Seldon with a host parameter on simple data that needs no preprocessing.
#
# Generate simple data:

# %%
import os

import numpy as np
from sklearn.datasets import make_classification
from sklearn.decomposition import PCA
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

from neuralk import Seldon, SeldonClassifier

HOST = os.environ.get("HOST", "http://my.seldon.server")

X, y = make_classification(n_samples=200, n_features=10, n_classes=2, random_state=42)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# Ensure data is in the correct format (float32 for features, int64 for labels)
X_train = X_train.astype(np.float32)
X_test = X_test.astype(np.float32)
y_train = y_train.astype(np.int64)

print(f"{X_train.shape=} {y_train.shape=} {X_test.shape=} {y_test.shape=}")

# %%
# Now we apply Neuralk's Seldon model with a host parameter.
#
# .. note::
#
#    The host URL is read from the ``HOST`` environment variable.
#    If your server requires authentication, you can pass it via the
#    ``default_headers`` parameter.

# %%

# Initialize the model with your Seldon server URL from HOST env var
model = Seldon(
    host=HOST,
    model="seldon-small",
    timeout_s=300,
)

# Note: nothing actually happens during fit() -- in-context learning models are
# pretrained but require no fitting on our specific dataset. The fit method
# only stores the training data.
model = model.fit(X_train, y_train)

# Make predictions
predictions = model.predict(X_test)
probabilities = model.predict_proba(X_test)

accuracy = accuracy_score(y_test, predictions)
print(f"Accuracy: {accuracy}")
print(f"Predictions shape: {predictions.shape}")
print(f"Probabilities shape: {probabilities.shape}")

# %%
# Working with authentication
# ---------------------------
#
# If your Seldon server requires authentication, you can pass authentication
# headers via the ``default_headers`` parameter:

# %%
# Example with authentication headers
model_with_auth = Seldon(
    host=HOST,
    model="seldon-small",
    default_headers={"Authorization": "Bearer your-token"},
)

# %%
# Advanced configuration
# ----------------------
#
# Seldon supports various configuration options for
# fine-tuning the connection and request behavior:

# %%
# Example with advanced configuration
model_advanced = Seldon(
    host=HOST,
    dataset_name="my-dataset",
    model="seldon-large",  # Use a different model
    timeout_s=600,  # Longer timeout
    metadata={"source": "example", "version": "1.0"},
    user="user123",
    api_version="v1",
)

# %%
# Integration with scikit-learn pipelines
# ----------------------------------------
#
# Seldon can be integrated into scikit-learn pipelines:

# %%


# Create a pipeline with preprocessing and Seldon
pipeline = make_pipeline(
    StandardScaler(),
    PCA(n_components=5),
    SeldonClassifier(
        host=HOST,
        model="seldon-small",
    ),
)

# Fit and predict
pipeline.fit(X_train, y_train)
pipeline_predictions = pipeline.predict(X_test)
pipeline_accuracy = accuracy_score(y_test, pipeline_predictions)
print(f"Pipeline accuracy: {pipeline_accuracy}")