import tensorflow as tf
import coremltools as ct
import numpy as np
import PIL
from huggingface_hub import hf_hub_download
from huggingface_hub import snapshot_download
import os

# Helper class to extract features from one model, and then feed those features into a classification head
# Because coremltools will only perform inference on OSX, an alternative tensorflow inference pipeline uses
# a tensorflow feature extractor and feeds the features into a dynamically created keras model based on the coreml classification head.
class CoreMLPipeline:    
    def __init__(self, config, auth_key, use_tf):
        self.config = config
        self.use_tf = use_tf
        if use_tf:
            extractor_path = snapshot_download(repo_id=config["tf_extractor_repoid"], use_auth_token  = auth_key)
        else:
            extractor_path = hf_hub_download(repo_id=config["coreml_extractor_repoid"],
                                             filename=config["coreml_extractor_path"], use_auth_token  = auth_key)

        classifier_path = hf_hub_download(repo_id=config["coreml_classifier_repoid"], filename=config["coreml_classifier_path"],
                                          use_auth_token = auth_key)

        print(f"Loading extractor...{extractor_path}")
        if use_tf:
            self.extractor = tf.saved_model.load(os.path.join(extractor_path, config["tf_extractor_path"]))
        else:
            self.extractor = ct.models.MLModel(extractor_path)

        print(f"Loading classifier...{classifier_path}")
        self.classifier = ct.models.MLModel(classifier_path)

        if use_tf:
            self.make_keras_model()

    #Only MacOS can run inference on CoreML models. Convert it to tensorflow to match the tf feature extractor            
    def make_keras_model(self):
        spec = self.classifier.get_spec()
        nnClassifier = spec.neuralNetworkClassifier
        labels = nnClassifier.stringClassLabels.vector
        input = tf.keras.Input(shape = (1280))
        activation = "sigmoid" if len(labels) == 1 else "softmax"
        x = tf.keras.layers.Dense(len(labels), activation = activation)(input)
        model = tf.keras.Model(input,x, trainable = False)
        weights = np.array(nnClassifier.layers[0].innerProduct.weights.floatValue)
        weights = weights.reshape((len(labels),1280))
        weights = weights.T
        bias = np.array(nnClassifier.layers[0].innerProduct.bias.floatValue)    
        model.set_weights([weights,bias])    
        self.tf_model = model
        self.labels = labels
            
    def classify(self,resized):
        if self.use_tf:
            image = tf.image.convert_image_dtype(resized, tf.float32)
            image = tf.expand_dims(image, 0)
            features = self.extractor.signatures['serving_default'](image)
            input = {"input_1":features["output_1"]}
            output = self.tf_model.predict(input)
            results = {}
            for i,label in enumerate(self.labels):
                results[label] = output[i]
        else:
            features = self.extractor.predict({"image":resized})
            features = features["Identity"]
            output = self.classifier.predict({"features":features[0]})
            results = output["Identity"]
        return results