<filename>query_formulation/models/__init__.py
from abc import ABC, abstractmethod

import numpy as np
import tensorflow as tf
import json

class ModelInterface(ABC):
    def __init__(self, model_address: str):
        self._model = tf.keras.models.load_model(model_address)

    @abstractmethod
    def is_good(self, cuis, added_cui) -> bool:
        pass


class SimpleCNNModel(ModelInterface):

    def is_good(self, cuis_vectors, added_cui_vector) -> bool:
        query_vectors = np.append(cuis_vectors, added_cui_vector.reshape((1, added_cui_vector.shape[0]), ), axis=0)
        query_vectors = self._check_text_size(query_vectors)
        y = self._model.predict(query_vectors)
        return y.argmax(axis=1)[0] == 2

    def _check_text_size(self, query_vectors):
        text_size = query_vectors.shape[0]
        embedding_size = query_vectors.shape[1]
        model_text_size = self._model.layers[0].get_config()['batch_input_shape'][1]
        if text_size < model_text_size:
            padding_size = model_text_size - text_size
            query_vectors = np.pad(query_vectors, ((0, padding_size), (0, 0)), 'constant')
        elif text_size > model_text_size:
            self._model = change_model(self._model, new_input_shape=(None, text_size, embedding_size))
        query_vectors = query_vectors.reshape((1, *query_vectors.shape), )
        return query_vectors


def change_model(model, new_input_shape=(None, 21, 200), custom_objects=None):
    # replace input shape of first layer

    config = model.layers[0].get_config()
    config['batch_input_shape'] = new_input_shape
    model_json = json.loads(model.to_json())
    model_json['config']['layers'][0]['config'] = config
    # rebuild model architecture by exporting and importing via json
    new_model = tf.keras.models.model_from_json(json.dumps(model_json), custom_objects=custom_objects)


    # copy weights from old model to new one
    for layer in new_model.layers:
        try:
            layer.set_weights(model.get_layer(name=layer.name).get_weights())
        except:
            print("Could not transfer weights for layer {}".format(layer.name))

    return new_model


if __name__ == '__main__':
    model = tf.keras.models.load_model('simple_cnn_model')
    new_model = change_model(model, (None, 35, 200))
    new_model.summary()
