import shutil
import sys
from pathlib import Path

from concrete.ml.deployment import FHEModelDev
from concrete.ml.deployment import FHEModelClient


def compile_and_make_it_deployable(model_dev, X_train):

    path_to_model = Path("compiled_model")

    # Compile into FHE
    model_dev.compile(X_train)

    # Saving the model
    shutil.rmtree(path_to_model, ignore_errors=True)
    fhemodel_dev = FHEModelDev(path_to_model, model_dev)
    fhemodel_dev.save(via_mlir=True)

    # To see the size of the key
    fhemodel_client = FHEModelClient(path_to_model)

    # Generate the keys
    fhemodel_client.generate_private_and_evaluation_keys()
    evaluation_keys = fhemodel_client.get_serialized_evaluation_keys()

    print(f"Your keys will be {sys.getsizeof(evaluation_keys) / 1024 / 1024}-megabytes long")

    # Check accuracy with p_error
    y_pred_simulated = model_dev.predict(X_test, fhe="simulate")
    simulated_accuracy = accuracy_score(Y_test, y_pred_simulated)

    print(f"Concrete average precision score (simulate): {simulated_accuracy:0.2f}")


# BEGIN: insert your ML task here
# Typically
from concrete.ml.sklearn import NeuralNetClassifier
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split
from sklearn.datasets import load_iris
from torch import nn

# Get iris data-set
X, y = load_iris(return_X_y=True)

# Split into train and test
X_train, X_test, Y_train, Y_test = train_test_split(X, y, test_size=0.25, random_state=42)

# Scikit-Learn and Concrete ML neural networks only handle float32 input values
X_train, X_test = X_train.astype("float32"), X_test.astype("float32")

params = {
    "module__n_layers": 3,
    "module__activation_function": nn.ReLU,
    "max_epochs": 1000,
    "verbose": 0,
}

model_dev = NeuralNetClassifier(**params)

model_dev = model_dev.fit(X=X_train, y=Y_train)

# Evaluate the Concrete ML model in the clear
y_pred_simulated = model_dev.predict(X_test)

simulated_accuracy = accuracy_score(Y_test, y_pred_simulated)
print(f"The test accuracy of the trained Concrete ML simulated model is {simulated_accuracy:.2f}")

# END: insert your ML task here

compile_and_make_it_deployable(model_dev, X_train)
print("Your model is ready to be deployable.")