import logging
from concurrent.futures import ProcessPoolExecutor, as_completed

import numpy as np

from pegasus.embedding_functions import MultiModalEmbeddingFunction


#logging
logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s', level=logging.INFO)
logger = logging.getLogger(__name__)


def optimized_embedding_function(modality, data):
    #creates and applies a MultiModalEmbeddingFunction for a specfic modality 

    """
    inputs:

        modality: A string representing the modality in lower case 'text' 'vision' 'audio'
        data: A numpy array representing the data'
    
    Returns:
        the embeddings generated by MultiModalEmbeddingFunction
    
    """
    try:     
        return MultiModalEmbeddingFunction(modality)(data)
    except Exception as e:
        logger.error(f"Failed to generate embeddings: {str(e)}")
        raise

class Pegasus:
    """
    Pegasus is the main multi-modal embedding class

    Inputs:
        modality: A string representing the modality => "text' 'audio'
        multi_process: A boolean indicating if multiprocessing will be enabled
        n_processes: An integer indicating that the number of processes to use
    """
    def __init__(
            self, 
            modality, 
            multi_process=False, 
            n_processes=1, 
            hosted=False
        ):

        if not isinstance(modality, str) or modality not in {"text", "audio", "vision", "sensor", "heatmap"}:
            logger.error(f"Invalid modality: {modality}")
            raise ValueError("Invalid modality")

        if not isinstance(multi_process, bool):
            logger.error(f"Invalid multi_process value: {multi_process}")
            raise ValueError("multi_process should be a boolean")

        if not isinstance(n_processes, int) or n_processes < 1:
            logger.error(f"Invalid n_processes value: {n_processes}")
            raise ValueError("n_processes should be a positive integer")
        
        self.modality = modality
        self.multi_process = multi_process and n_processes > 1
        self.n_processes = n_processes
        self.hosted = False

    def _embed_data(self, data):
        """
        Embeds the data using MultiModalEmbeddingFunction 

        Args:
            data: a numpy array representing the data

        Returns:
            The embeddings generated by MultiModalEmbeddingFunction
        
        """
        if self.modality not in {"text", "audio", "vision", "sensor", "heatmap"}:
            raise ValueError("Invalid modality")
        return optimized_embedding_function(self.modality, data)
    
    def embed_data(self, data):
        """
        Embeds the data using MultiModalEmbeddingFunction

        if multiprocessing is enabled, the data is split and processed in parallel 

        Inputs:
            data: a numpy array or a list representing the data

        Returns:
            the embeddings generated by the MultiModalEmbeddingFunction
        
        """
        if not isinstance(data, np.ndarray):
            try:
                data = np.array(data)
            except Exception as e:
                logger.error(f"Failed to convert data to numpy array: {str(e)}")
                raise
        
        if not self.multi_process:
            return self._embed_data(data)
        
        if self.multi_process:
            try:
                with ProcessPoolExecutor(max_workers=self.n_processes) as executor:
                    future_to_data = {executor.submit(self._embed_data, d): d for d in data}
                    return {
                        future_to_data[future]: future.result() for future in as_completed(future_to_data)
                    }
            except Exception as e:
                logger.error(f"Failed to embed data in parallel: {str(e)}")
                raise
        else:
            return self._embed_data(data)