"""
Copyright (c) : Calixte Mayoraz 2024
https://gitlab.com/calixtemayoraz
"""
import os
import tensorflow as tf
import logging
import numpy as np
from .preprocess import preprocess


class MusicSquidModel:
    MODEL_PATH = os.path.abspath(os.path.join(__file__, '..', 'model'))

    def __init__(self):
        self._model: tf.keras.Model = tf.keras.models.load_model(MusicSquidModel.MODEL_PATH)
        logging.info("Loaded model %s", self.model_version)

    @property
    def model_version(self) -> str:
        return self._model.name

    @property
    def output_shape(self) -> tuple:
        return self._model.output_shape

    def embed(self, filename) -> 'np.ndarray|tuple[np.ndarray,np.ndarray]':
        """
        Returns an embedding for the input file.

        Since the model was trained on 30s segments of tracks,
        we compute an embedding for all 30s segments of the input data and
        return a median vector for the output embeddings.


        Parameters
        ----------
        filename: str

        Returns
        -------
        np.ndarray|tuple[np.ndarray,np.ndarray]:
            the embedded track or the embedded track and class probabilities
        """
        return np.median(self._model(preprocess(filename)), axis=0)