import numpy as np
from decimal import Decimal
from sentence_transformers import SentenceTransformer, util
from transformers import pipeline

class Swarm:
    """
    ToDo: Clean up the verbose outputs (prints)
    ToDo: Log instead of print? Or both?

    A multi-agent system for output validation and consensus using semantic similarity,
    paraphrase detection, and cosine similarity.

    This class implements the following:
    - Byzantine Fault Tolerance (BFT)-like validation.
    - Semantic matching to ensure compliance with predefined requirements.
    - Paraphrase detection to identify and group semantically similar outputs.
    - Consensus selection to determine the best output based on similarity.

    Parameters:
    - query (str): Input query for the agents.
    - llms (list): List of LLM configurations (provider, model, and API key).
    - clients (list): List of initialized clients.
    - state (list): Initial state memory for agents.
    - verbose (bool): Enable detailed logging.
    - sensitivity (int): Sensitivity factor for determining the number of agents.
    - minimum_bots (int): Minimum number of agents to instantiate.
    - maximum_bots (int): Maximum number of agents to instantiate.
    - confidence (float): Confidence level (0-1) for agent selection.
    - threshold (float): Similarity threshold for validation.
    - requirements (list): List of predefined requirements for output validation.
    - paraphrase_threshold (float): Similarity threshold for paraphrase detection.
    - model (str): Name of the SentenceTransformer model.
    - instructions (str): Instructions for the agents.

    Attributes:
    - bots (int): Calculated number of agents based on sensitivity and confidence.
    - paragraphs (list): Outputs generated by agents.
    """

    def __init__(
        self,
        query='',
        llms=None,
        clients=None,
        state=None,
        verbose=False,
        sensitivity=1,
        minimum_bots=1,
        maximum_bots=100,
        confidence=0.5,
        threshold=0.75,
        requirements=None,
        paraphrase_threshold=0.8,
        model='all-MiniLM-L6-v2',
        instructions='You are a helpful assistant.'
    ):
        self.llms = llms or []
        self.query = query
        self.state = state
        self.paragraphs = []
        self.clients = clients or []
        self.verbose = verbose
        self.threshold = threshold
        self.confidence = confidence
        self.sensitivity = sensitivity
        self.minimum_bots = minimum_bots
        self.maximum_bots = maximum_bots
        self.instructions = instructions
        self.requirements = requirements or []
        self.model = SentenceTransformer(model)
        self.paraphrase_threshold = paraphrase_threshold
        self.bots = int(
            min(
                self.maximum_bots,
                max(
                    self.minimum_bots,
                    Decimal(str(self.confidence)).as_integer_ratio()[1] * self.sensitivity
                )
            )
        )

        if self.verbose and not self.clients:
            print("\nBots:", self.bots)

    def check_initialization(self):
        """
        Validate essential attributes before running the Swarm.

        Returns:
            bool: True if all necessary attributes are initialized.
            
        ToDo: Perform the checks.
        """
        #return all(
        #    [
        #        self.model,
        #        self.requirements,
        #        self.instructions,
        #        self.threshold,
        #        self.llms,
        #        self.query,
        #        isinstance(self.state, (list, type(None))),
        #    ]
        #)
        return True

    def _create_paragraphs(self, llm, erase_query=False):
        """
        Generate output paragraphs from an LLM client.

        Args:
            llm: Initialized LLM client.
            erase_query (bool): Whether to remove the query from memory after execution.
        """
        if self.state is not None:
            llm.set_memory(self.state)
        self.paragraphs.append(llm.chat(q=self.query, erase_query=erase_query))

    def _create_client(self, llm_config):
        """
        Initialize an LLM client and generate output.

        Args:
            llm_config (dict): Configuration for the LLM client.
        """
        if len(self.clients) >= self.bots:
            return

        _llm = LLM(
            provider=llm_config['provider'],
            model=llm_config['model'],
            api_key=llm_config['key'],
            system_prompt=f"""{self.instructions} {self.requirements}"""
        )
        self._create_paragraphs(_llm)

        if self.verbose:
            print("\nParagraph created.")

        self.clients.append(_llm)

        if self.verbose:
            print("\nClient appended.")

    def create_clients(self):
        """
        Create LLM clients and distribute tasks among them.

        Returns:
            int: Total number of clients created.
        """
        self.clients = []
        counter = 0
        nbr_of_llms = len(self.llms)

        for _ in range(self.bots // nbr_of_llms):
            [self._create_client(x) for x in self.llms]
            counter += nbr_of_llms

        for _ in range(self.bots % nbr_of_llms):
            [self._create_client(x) for x in self.llms]
            counter += 1

            if len(self.clients) >= self.bots:
                break

        self.clients = self.clients[:self.bots]
        return counter

    def instantiate(self):
        """
        Ensure all prerequisites are met and initialize clients.

        Returns:
            bool: True if initialization is successful.
        """
        if self.check_initialization():
            if self.verbose:
                print("\nInitialization successful.")

            created_clients = self.create_clients()

            if self.verbose:
                print("\nClients created:", created_clients)

            return True

        return False

    def calculate_global_similarity(self, paragraphs, requirement_sentences):
        """
        Compute the global average similarity between outputs and requirements.

        Args:
            paragraphs (list): Generated outputs.
            requirement_sentences (list): Requirement sentences.

        Returns:
            float: Global average similarity score.
        """
        requirement_embeddings = self.model.encode(requirement_sentences)
        paragraph_embeddings = self.model.encode(paragraphs)

        all_similarities = []

        for paragraph_embedding in paragraph_embeddings:
            similarities = util.cos_sim(paragraph_embedding, requirement_embeddings)
            all_similarities.extend(similarities[0].tolist())

        return np.mean(all_similarities)

    def dynamic_threshold(self, global_average_similarity, threshold, adjustment_factor=0.8):
        """
        Adjust the threshold dynamically based on the global similarity.
    
        Parameters:
        - global_average_similarity (float): Average similarity score (0.0 to 1.0).
        - threshold (float): Initial threshold value (0.0 to 1.0).
        - adjustment_factor (float): Factor for adjusting the threshold.
    
        Returns:
        - float: Adjusted threshold value, clamped between 0.0 and 1.0.
    
        Raises:
        - ValueError: If input values are out of expected ranges.
        """
        # Input validation
        if not (0.0 <= global_average_similarity <= 1.0):
            raise ValueError("global_average_similarity must be between 0.0 and 1.0")
        if not (0.0 <= threshold <= 1.0):
            raise ValueError("threshold must be between 0.0 and 1.0")
        if not (0.0 <= adjustment_factor <= 1.0):
            raise ValueError("adjustment_factor must be between 0.0 and 1.0")
    
        # Calculate adjusted threshold
        adjusted_threshold = threshold - (adjustment_factor * (threshold - global_average_similarity))
    
        # Clamp the result between 0.0 and 1.0
        adjusted_threshold = max(0.0, min(1.0, adjusted_threshold))
    
        # Log for debugging
        if self.verbose:
            print(f"Dynamic Threshold Adjustment:")
            print(f"  Original Threshold: {threshold}")
            print(f"  Global Avg Similarity: {global_average_similarity}")
            print(f"  Adjustment Factor: {adjustment_factor}")
            print(f"  Adjusted Threshold: {adjusted_threshold}")
    
        return adjusted_threshold


    def classify_requirements(self, requirements):
        """
        Classify requirements into positive and negative based on sentiment.

        Args:
            requirements (list): List of requirement sentences.

        Returns:
            tuple: Positive and negative requirements.
        """
        classifier = pipeline("text-classification", model="lxyuan/distilbert-base-multilingual-cased-sentiments-student")

        positive_requirements = []
        negative_requirements = []

        for req in requirements:
            if self.verbose:
                print("\nClassify requirement:", req)

            result = classifier(req)[0]

            if self.verbose:
                print("\nResult:", result)

            label = result['label'].lower()

            if label == 'positive':
                positive_requirements.append(req)
            else:
                negative_requirements.append(req)

        return positive_requirements, negative_requirements

    def detect_paraphrases(self, compliant_paragraphs, compliant_embeddings, paraphrase_threshold):
        """
        Detect paraphrases among compliant paragraphs using cosine similarity.

        Args:
            compliant_paragraphs (list): Compliant paragraphs.
            compliant_embeddings (list): Embeddings of compliant paragraphs.
            paraphrase_threshold (float): Similarity threshold for paraphrases.

        Returns:
            list: Groups of paraphrases.
        """
        paraphrase_groups = []
        used_paragraphs = set()

        for i, paragraph_embedding in enumerate(compliant_embeddings):
            if i not in used_paragraphs:
                group = [compliant_paragraphs[i]]
                used_paragraphs.add(i)

                for j in range(i + 1, len(compliant_embeddings)):
                    similarity = util.cos_sim(paragraph_embedding, compliant_embeddings[j])
                    if similarity >= paraphrase_threshold:
                        group.append(compliant_paragraphs[j])
                        used_paragraphs.add(j)

                paraphrase_groups.append(group)

        return paraphrase_groups
    
    def get_consensus(self, paraphrase_groups, compliant_paragraphs, compliant_embeddings):
        """
        Determine the paragraph that is most similar to all other compliant paragraphs.
    
        This method calculates the consensus paragraph from paraphrase groups based on average cosine similarity. 
        If there are only two paraphrase groups, binary consensus is assumed, and the largest group is selected.
    
        Parameters:
        - paraphrase_groups (list of lists): Groups of paragraphs identified as paraphrases of each other.
        - compliant_paragraphs (list): List of all compliant paragraphs.
        - compliant_embeddings (list): Embeddings of the compliant paragraphs.
    
        Returns:
        - best_paragraph (str): The paragraph with the highest consensus.
        - highest_similarity (float): The similarity score of the selected paragraph.
        - group_size_of_best (int): The size of the paraphrase group containing the selected paragraph.
        """
    
        # Initialize variables to track the best consensus
        best_paragraph = None
        highest_similarity = -1
        group_size_of_best = 0
    
        try:
            # Step 1: Handle special case where there are exactly two paraphrase groups
            if len(paraphrase_groups) == 2:
                # Assume binary consensus and retain only the larger group
                paraphrase_groups = [max(paraphrase_groups, key=len)]
    
            # Step 2: Iterate over each paraphrase group to calculate average similarity
            for group in paraphrase_groups:
                if len(group) > 1:  # Only consider groups with more than one paragraph
                    # Encode all paragraphs in the group into embeddings
                    paragraph_embeddings = self.model.encode(group)
    
                    # Compute the average cosine similarity for each paragraph within the group
                    avg_similarities = [
                        util.cos_sim(paragraph_embeddings[i], paragraph_embeddings).mean().item()
                        for i in range(len(group))
                    ]
    
                    # Find the paragraph with the highest average similarity within the group
                    best_index = avg_similarities.index(max(avg_similarities))
                    if avg_similarities[best_index] > highest_similarity:
                        best_paragraph = group[best_index]
                        highest_similarity = avg_similarities[best_index]
                        group_size_of_best = len(group)
                else:
                    # Handle single-paragraph groups (no paraphrases in this group)
                    if highest_similarity == -1:  # Select the first single paragraph if no best found yet
                        best_paragraph = group[0]
                        group_size_of_best = len(group)
    
        except Exception as e:
            # Log the error for debugging purposes
            if self.verbose:
                print(f"Error during get_consensus: {e}")
            # Fallback in case of an error
            best_paragraph = "No consensus could be determined."
            highest_similarity = 0
            group_size_of_best = 0

        best_paragraph = best_paragraph or "No consensus could be determined."
        highest_similarity = highest_similarity if highest_similarity >= 0 else 0
        
        return best_paragraph, highest_similarity, group_size_of_best

    def run(self):
        """
        Execute the Swarm workflow and determine consensus output.

        Returns:
            str: Consensus paragraph.
        """
        consensus_paragraph = 'No consensus found.'

        if self.instantiate():
            if self.verbose:
                print("Class instantiated.")

            requirement_sentences = self.requirements

            global_average_similarity = self.calculate_global_similarity(self.paragraphs, requirement_sentences)

            if self.verbose:
                print("Global Average Similarity:", global_average_similarity)

            dynamic_threshold = self.dynamic_threshold(global_average_similarity, self.threshold)

            if self.verbose:
                print("Dynamic Threshold:", dynamic_threshold)

            positive_requirements, negative_requirements = self.classify_requirements(requirement_sentences)

            if self.verbose:
                print("Positive Requirements:", positive_requirements)
                print("Negative Requirements:", negative_requirements)

            paraphrase_groups = self.detect_paraphrases(
                self.paragraphs,
                [self.model.encode(paragraph) for paragraph in self.paragraphs],
                self.paraphrase_threshold
            )

            if self.verbose:
                print("\nParaphrase Groups:", paraphrase_groups)

            return consensus_paragraph

        return consensus_paragraph
