import json
from typing import Any, Dict, List

import tensorflow as tf
from tensorflow import keras
import base64
import io
import os
import numpy as np
from PIL import Image



class PreTrainedPipeline():
    def __init__(self, path: str):
        # load the model
        self.model = keras.models.load_model(os.path.join(path, "tf_model.h5"))

    def __call__(self, inputs: "Image.Image")-> List[Dict[str, Any]]:

        # convert img to numpy array, resize and normalize to make the prediction
        img = np.array(inputs)

        im = tf.image.resize(img, (128, 128))
        im = tf.cast(im, tf.float32) / 255.0
        pred_mask = self.model.predict(im[tf.newaxis, ...])
        
        # take the best performing class for each pixel
        # the output of argmax looks like this [[1, 2, 0], ...]
        pred_mask_arg = tf.argmax(pred_mask, axis=-1)

        labels = []
        
        # convert the prediction mask into binary masks for each class
        binary_masks = {}
        mask_codes = {}
        
        # when we take tf.argmax() over pred_mask, it becomes a tensor object
        # the shape becomes TensorShape object, looking like this TensorShape([128]) 
        # we need to take get shape, convert to list and take the best one
        
        rows = pred_mask_arg[0][1].get_shape().as_list()[0]
        cols = pred_mask_arg[0][2].get_shape().as_list()[0]
        
        for cls in range(pred_mask.shape[-1]):

            binary_masks[f"mask_{cls}"] = np.zeros(shape = (pred_mask.shape[1], pred_mask.shape[2])) #create masks for each class
            
            for row in range(rows):

                for col in range(cols):

                    if pred_mask_arg[0][row][col] == cls:
                        
                        binary_masks[f"mask_{cls}"][row][col] = 1
                    else:
                        binary_masks[f"mask_{cls}"][row][col] = 0

            mask = binary_masks[f"mask_{cls}"]
            mask *= 255
            img = Image.fromarray(mask.astype(np.int8), mode="L")
               
            # we need to make it readable for the widget
            with io.BytesIO() as out:
                img.save(out, format="PNG")
                png_string = out.getvalue()
                mask = base64.b64encode(png_string).decode("utf-8")

            mask_codes[f"mask_{cls}"] = mask
    

            # widget needs the below format, for each class we return label and mask string
            labels.append({
                "label": f"LABEL_{cls}",
                "mask": mask_codes[f"mask_{cls}"],
                "score": 1.0,
            })
        return labels