from .synthetic_generator import synthetic_generator_config
from .binary_classifier import binary_classifer_config
from .rag_scoring import rag_scoring_config
from .ues_idp import ues_idp_config
from .kilt_filter import KILT_dataset_process
from .superglue_filter import superGlue
from .label_filter import filter_tsv_by_label
from .prompts import context_relevance_system_prompt, answer_relevance_system_prompt, answer_faithfulness_system_prompt
from typing import List

class ARES:
    """
    ARES class encapsulates configurations for various components of the ARES system.
    It allows for the setup and management of configurations for synthetic query generation,
    binary classification, RAG scoring, and UES/IDP processes.
    """
    
    # A placeholder object used to denote default values in the configuration specification
    REQUIRED_BUT_HAS_DEFAULT = object()

    # Configuration specification dictionary defining the parameters for each component
    config_spec = {
        "ues_idp": {
            "in_domain_prompts_dataset": (str, None),  # Required parameter with no default value
            "unlabeled_evaluation_set": (str, None),  # Required parameter with no default value
            "context_relevance_system_prompt": (str,context_relevance_system_prompt),  # Optional with default
            "answer_relevance_system_prompt": (str, answer_relevance_system_prompt),  # Optional with default
            "answer_faithfulness_system_prompt": (str, answer_faithfulness_system_prompt),  # Optional with default
            "debug_mode": (bool, False),  # Optional with default
            "documents": (int, "None"),  # Optional with default
            "model_choice": (str, "gpt-3.5-turbo-1106"),  # Optional with default
            "request_delay": (int, 0),  # Optional with default
            "vllm": (bool, False),  # Optional with default
            "host_url": (str, "None")  # Optional with default
        },

        "synthetic_query_generator": {
            "document_filepaths": (list, None),  # Required parameter with no default value
            "few_shot_prompt_filename": (list, None),  # Required parameter with no default value
            "synthetic_queries_filenames": (list, None),  # Required parameter with no default value
            "documents_sampled": (int, None),  # Required parameter with no default value
            "model_choice": (str, "google/flan-t5-xxl"),  # Optional with default
            "api_model": (bool, False), # Optional with default
            "clean_documents": (bool, False),  # Optional with default
            "regenerate_synth_questions": (bool, True),  # Optional with default
            "percentiles": (list, [0.05, 0.25, 0.5, 0.95]),  # Optional with default
            "question_temperatures": (list, [2.0, 1.5, 1.0, 0.5, 0.0]),  # Optional with default
            "regenerate_answers": (bool, True),  # Optional with default
            "number_of_negatives_added_ratio": (float, 0.5),  # Optional with default
            "lower_bound_for_negatives": (int, 5),  # Optional with default
            "number_of_contradictory_answers_added_ratio": (float, 0.67),  # Optional with default
            "number_of_positives_added_ratio": (float, 0.0),  # Optional with default
            "regenerate_embeddings": (float, True),  # Optional with default
            "synthetic_query_prompt": (str, "You are an expert question-answering system. Generate only one question based on the provided document. Ensure the question is answerable within the context of the document. Do not generate multiple questions. Do not provide labels, headers, or additional text. Only return a single, clear question. Generating more than one question will be considered incorrect output.\n\n"),
            "synthetic_valid_answer_prompt": (str, "You are an expert question-answering system. You must create an answer for the provided question. The answer must be answerable within the context of the document. Return only the answer, nothing else.\n\n"),
            "synthetic_contradictory_answer_prompt": (str, "Create an answer for the given question that contradicts the provided document. You should create false information that disagrees with what exists within the content of the document.  Return only the false answer, without any labels or additional text.\n\n")
        },
        
        "classifier_model": {
            "training_dataset": (list, None),  # Required parameter with no default value
            "validation_set": (list, None),  # Required parameter with no default value
            "label_column": (list, None),  # Required parameter with no default value
            "num_epochs": (int, None),  # Required parameter with no default value
            "patience_value": (int, None),  # Required parameter with no default value
            "learning_rate": (float, None),  # Required parameter with no default value
            "training_dataset_path": (str, "None"),  # Optional with default
            "validation_dataset_path": (str, "None"),  # Optional with default
            "model_choice": (str, "microsoft/deberta-v3-large"),  # Optional with default
            "validation_set_scoring": (bool, True),  # Optional with default
            "assigned_batch_size": (int, REQUIRED_BUT_HAS_DEFAULT, 1),  # Default is 1
            "gradient_accumulation_multiplier": (int, REQUIRED_BUT_HAS_DEFAULT, 32),  # Default is 32
            "number_of_runs": (int, 1),  # Optional with default
            "num_warmup_steps": (int, 100),  # Optional with default
            "training_row_limit": (int, -1),  # Optional with default
            "validation_row_limit": (int, -1)  # Optional with default
        },

        "ppi": {
            "evaluation_datasets": (list, None),  # Required parameter with no default value
            "labels": (list, None),  # Required parameter with no default value
            "checkpoints": (list, None),  # Required parameter with no default value
            "few_shot_examples_filepath": (str, "None"), # Optional with default
            "gold_label_paths": (list, ["None"]),  # Optional with default
            "rag_type": (str, "question_answering"),  # Optional with default
            "model_choice": (str, "microsoft/deberta-v3-large"),  # Optional with default
            "llm_judge": (str, "None"),  # Optional with default
            "assigned_batch_size": (int, 1),  # Optional with default
            "number_of_labels": (int, 2),  # Optional with default
            "alpha": (int, 0.05),  # Optional with default
            "num_trials": (int, 1000),  # Optional with default
            "vllm": (bool, False),  # Optional with default
            "host_url": (str, "http://0.0.0.0:8000/v1"),  # Optional with default
            "request_delay": (int, 0),  # Optional with default
            "debug_mode": (bool, False),  # Optional with default
            "machine_label_llm_model": (str, "None"),  # Optional with default
            "gold_machine_label_path": (str, "None"),  # Optional with default
            "prediction_filepaths": (list, ["None"])  # Optional with default
        }
    }

    def __init__(self, synthetic_query_generator={}, ues_idp={}, classifier_model={}, ppi={}):
        """
        Initializes the ARES class with configurations for different components.

        Args:
            synthetic_query_generator (dict): Configuration for the synthetic query generator.
            ues_idp (dict): Configuration for UES IDP.
            classifier_model (dict): Configuration for the classifier model.
            ppi (dict): Configuration for PPI (Protein-Protein Interaction).

        Each configuration dictionary is passed to the `prepare_config` method to validate
        and prepare the final configuration based on the default settings defined in `config_spec`.
        """
        self.synthetic_query_generator_config = self.prepare_config("synthetic_query_generator", synthetic_query_generator)
        self.classifier_model_config = self.prepare_config("classifier_model", classifier_model)
        self.ppi_config = self.prepare_config("ppi", ppi)
        self.ues_idp_config = self.prepare_config("ues_idp", ues_idp)

    def generate_synthetic_data(self):
        """
        Generates synthetic data using the synthetic query generator configuration.
        If the configuration is empty, the generation process is skipped.
        """
        if not self.synthetic_query_generator_config:
            print("Skipping synthetic generator configuration due to missing parameters.")
        else:
            synthetic_generator_config(**self.synthetic_query_generator_config)

    def train_classifier(self):
        """
        Trains the classifier using the classifier model configuration.
        If the configuration is empty, the training process is skipped.
        """
        if not self.classifier_model_config:
            print("Skipping binary classifier configuration due to missing parameters.")
        else:
            binary_classifer_config(**self.classifier_model_config)

    def evaluate_RAG(self):
        """
        Evaluates the RAG using the PPI configuration.
        If the configuration is empty, the evaluation process is skipped.
        """
        if not self.ppi_config:
            print("Skipping RAG evaluation configuration due to no parameters")
        else:
            return rag_scoring_config(**self.ppi_config)

    def ues_idp(self):
        """
        Executes the UES IDP configuration.
        """
        return ues_idp_config(**self.ues_idp_config)

    def KILT_dataset(self, dataset_name):
        """
        Processes the KILT dataset with the specified dataset name.
        """
        KILT_dataset_process(dataset_name)

    def superGlue_dataset(self, dataset_name):
        """
        Processes the SuperGlue dataset with the specified dataset name.
        """
        superGlue(dataset_name)
    
    def filter_tsv_by_label(self, tsv_file, label, output_file=None):
        """
        Filters the TSV file by the specified label and saves the filtered data to a new TSV file.
        """
        filter_tsv_by_label(tsv_file, label, output_file)

    def run(self):
        """
        Executes the configurations for synthetic data generation, classifier training, and RAG evaluation.
        If any configuration is empty, the corresponding process is skipped.
        """
        self.generate_synthetic_data()
        self.train_classifier()
        self.evaluate_RAG()

        if not self.ues_idp_config:
            print(f"Skipping UES and IDP since either unlabeled evaluation set or in domain prompts not provided")
            exit()
        else:
            model_choice = self.ues_idp_config["model_choice"]
            ues_idp_config(**self.ues_idp_config)

    def prepare_config(self, component_name, user_config):
        """
        Prepares and validates the configuration for a given component based on the default settings.

        Args:
            component_name (str): The name of the component for which the configuration is being prepared.
            user_config (dict): The user-provided configuration dictionary.

        Returns:
            dict: The prepared and validated configuration dictionary.

        Raises:
            TypeError: If the provided value type does not match the expected type.
            ValueError: If a required parameter is missing and no default is specified.
        """
        if not user_config:
            return {}
        component = self.config_spec[component_name]
        prepared_config = {}
        for param, config in component.items():
            expected_type, default, default_value = config if len(config) == 3 else (*config, None)
            if param in user_config:
                value = user_config[param]
                if not isinstance(value, expected_type):
                    raise TypeError(f"Parameter '{param}' for {component_name} is expected to be of type {expected_type.__name__}, received {type(value).__name__} instead.")
                prepared_config[param] = value
            elif default is self.REQUIRED_BUT_HAS_DEFAULT:
                # If the parameter is required but not provided, use the specified default value.
                if param not in user_config:
                    print(f"\nWarning: '{param}' not provided for {component_name}, using default value {default_value}.")
                prepared_config[param] = default_value
            elif default is not None:
                prepared_config[param] = default
            else:
                raise ValueError(f"Missing required parameter '{param}' for {component_name}.")
        return prepared_config