"Client-server interface custom implementation for filter models."

import json
import concrete.numpy as cnp
from filters import Filter

from concrete.ml.deployment import FHEModelClient
from concrete.ml.version import __version__ as CML_VERSION


class CustomFHEClient(FHEModelClient):
    """Client interface to encrypt and decrypt FHE data associated to a Filter."""

    def load(self):
        """Load the parameters along with the FHE specs."""

        # Load the client
        self.client = cnp.Client.load(self.path_dir / "client.zip", self.key_dir)

        # Load the filter's parameters from the json file
        with (self.path_dir / "serialized_processing.json").open("r", encoding="utf-8") as f:
            serialized_processing = json.load(f)

        # Make sure the version in serialized_model is the same as CML_VERSION
        assert serialized_processing["cml_version"] == CML_VERSION, (
            f"The version of Concrete ML library ({CML_VERSION}) is different "
            f"from the one used to save the model ({serialized_processing['cml_version']}). "
            "Please update to the proper Concrete ML version.",
        )

        # Initialize the filter model using its filter name
        filter_name = serialized_processing["model_post_processing_params"]["filter_name"]
        self.model = Filter(filter_name)

        return self.model