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")

    n_bits = 4

    model_dev = compile_torch_model(model_dev, X_train, rounding_threshold_bits=6, p_error=0.1)

    # Accuracy in simulation
    accs = test_with_concrete(
        model_dev,
        test_dataloader,
        use_sim=True,
    )

    print(f"Simulated FHE execution for {n_bits} bit model_devwork accuracy: {accs:.2f}")

    # 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)


# BEGIN: insert your ML task here
# Typically
import time

import numpy as np
import torch
import torch.utils
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from tqdm import tqdm

from concrete.ml.torch.compile import compile_torch_model

X, y = load_digits(return_X_y=True)

# The sklearn Digits data-set, though it contains digit images, keeps these images in vectors
# so we need to reshape them to 2D first. The images are 8x8 px in size and monochrome
X = np.expand_dims(X.reshape((-1, 8, 8)), 1)

X_train, X_test, Y_train, Y_test = train_test_split(
    X, y, test_size=0.25, shuffle=True, random_state=42
)


class TinyCNN(nn.Module):
    """A very small CNN to classify the sklearn digits data-set."""

    def __init__(self, n_classes) -> None:
        """Construct the CNN with a configurable number of classes."""
        super().__init__()

        # This model_devwork has a total complexity of 1216 MAC
        self.conv1 = nn.Conv2d(1, 8, 3, stride=1, padding=0)
        self.conv2 = nn.Conv2d(8, 16, 3, stride=2, padding=0)
        self.conv3 = nn.Conv2d(16, 32, 2, stride=1, padding=0)
        self.fc1 = nn.Linear(32, n_classes)

    def forward(self, x):
        """Run inference on the tiny CNN, apply the decision layer on the reshaped conv output."""
        x = self.conv1(x)
        x = torch.relu(x)
        x = self.conv2(x)
        x = torch.relu(x)
        x = self.conv3(x)
        x = torch.relu(x)
        x = x.flatten(1)
        x = self.fc1(x)
        return x


torch.manual_seed(42)


def train_one_epoch(model_dev, optimizer, train_loader):
    # Cross Entropy loss for classification when not using a softmax layer in the model_devwork
    loss = nn.CrossEntropyLoss()

    model_dev.train()
    avg_loss = 0
    for data, target in train_loader:
        optimizer.zero_grad()
        output = model_dev(data)
        loss_model_dev = loss(output, target.long())
        loss_model_dev.backward()
        optimizer.step()
        avg_loss += loss_model_dev.item()

    return avg_loss / len(train_loader)


def test_torch(model_dev, test_loader):
    """Test the model_devwork: measure accuracy on the test set."""

    # Freeze normalization layers
    model_dev.eval()

    all_y_pred = np.zeros((len(test_loader)), dtype=np.int64)
    all_targets = np.zeros((len(test_loader)), dtype=np.int64)

    # Iterate over the batches
    idx = 0
    for data, target in test_loader:
        # Accumulate the ground truth labels
        endidx = idx + target.shape[0]
        all_targets[idx:endidx] = target.numpy()

        # Run forward and get the predicted class id
        output = model_dev(data).argmax(1).detach().numpy()
        all_y_pred[idx:endidx] = output

        idx += target.shape[0]

    # Print out the accuracy as a percentage
    n_correct = np.sum(all_targets == all_y_pred)
    print(
        f"Test accuracy for fp32 weights and activations: "
        f"{n_correct / len(test_loader) * 100:.2f}%"
    )


def test_with_concrete(quantized_module, test_loader, use_sim):
    """Test a neural model_devwork that is quantized and compiled with Concrete ML."""

    # Casting the inputs into int64 is recommended
    all_y_pred = np.zeros((len(test_loader)), dtype=np.int64)
    all_targets = np.zeros((len(test_loader)), dtype=np.int64)

    # Iterate over the test batches and accumulate predictions and ground truth labels in a vector
    idx = 0
    for data, target in tqdm(test_loader):
        data = data.numpy()
        target = target.numpy()

        fhe_mode = "simulate" if use_sim else "execute"

        # Quantize the inputs and cast to appropriate data type
        y_pred = quantized_module.forward(data, fhe=fhe_mode)

        endidx = idx + target.shape[0]

        # Accumulate the ground truth labels
        all_targets[idx:endidx] = target

        # Get the predicted class id and accumulate the predictions
        y_pred = np.argmax(y_pred, axis=1)
        all_y_pred[idx:endidx] = y_pred

        # Update the index
        idx += target.shape[0]

    # Compute and report results
    n_correct = np.sum(all_targets == all_y_pred)

    return n_correct / len(test_loader)


# Create the tiny CNN with 10 output classes
N_EPOCHS = 50

# Create a train data loader
train_dataset = TensorDataset(torch.Tensor(X_train), torch.Tensor(Y_train))
train_dataloader = DataLoader(train_dataset, batch_size=64)

# Create a test data loader to supply batches for model_devwork evaluation (test)
test_dataset = TensorDataset(torch.Tensor(X_test), torch.Tensor(Y_test))
test_dataloader = DataLoader(test_dataset)

# Train the model_devwork with Adam, output the test set accuracy every epoch
model_dev = TinyCNN(10)
losses_bits = []
optimizer = torch.optim.Adam(model_dev.parameters())
for _ in tqdm(range(N_EPOCHS), desc="Training"):
    losses_bits.append(train_one_epoch(model_dev, optimizer, train_dataloader))

# END: insert your ML task here

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