from model import SUPPORTED_SUMM_MODELS
from model.base_model import SummModel
from model.single_doc import LexRankModel

from dataset.st_dataset import SummDataset
from dataset.non_huggingface_datasets import ScisummnetDataset

from typing import List, Tuple


def get_lxr_train_set(dataset: SummDataset, size: int = 100) -> List[str]:

    """
    return some dummy summarization examples, in the format of a list of sources
    """
    subset = []
    for i in range(size):
        subset.append(next(iter(dataset.train_set)))

    src = list(
        map(
            lambda x: " ".join(x.source)
            if dataset.is_dialogue_based or dataset.is_multi_document
            else x.source[0]
            if isinstance(dataset, ScisummnetDataset)
            else x.source,
            subset,
        )
    )

    return src


def assemble_model_pipeline(
    dataset: SummDataset, model_list: List[SummModel] = SUPPORTED_SUMM_MODELS
) -> List[Tuple[SummModel, str]]:

    """
    Return initialized list of all model pipelines that match the summarization task of given dataset.

    :param SummDataset `dataset`: Dataset to retrieve model pipelines for.
    :param List[SummModel] `model_list`: List of candidate model classes (uninitialized). Defaults to `model.SUPPORTED_SUMM_MODELS`.
    :returns List of tuples, where each tuple contains an initialized model and the name of that model as `(model, name)`.
    """

    dataset = dataset if isinstance(dataset, SummDataset) else dataset()

    single_doc_model_list = list(
        filter(
            lambda model_cls: not (
                model_cls.is_dialogue_based
                or model_cls.is_query_based
                or model_cls.is_multi_document
            ),
            model_list,
        )
    )
    single_doc_model_instances = [
        model_cls(get_lxr_train_set(dataset))
        if model_cls == LexRankModel
        else model_cls()
        for model_cls in single_doc_model_list
    ]

    multi_doc_model_list = list(
        filter(lambda model_cls: model_cls.is_multi_document, model_list)
    )

    query_based_model_list = list(
        filter(lambda model_cls: model_cls.is_query_based, model_list)
    )

    dialogue_based_model_list = list(
        filter(lambda model_cls: model_cls.is_dialogue_based, model_list)
    )
    dialogue_based_model_instances = (
        [model_cls() for model_cls in dialogue_based_model_list]
        if dataset.is_dialogue_based
        else []
    )

    matching_models = []
    if dataset.is_query_based:
        if dataset.is_dialogue_based:
            for query_model_cls in query_based_model_list:
                for dialogue_model in dialogue_based_model_list:
                    full_query_dialogue_model = query_model_cls(
                        model_backend=dialogue_model
                    )
                    matching_models.append(
                        (
                            full_query_dialogue_model,
                            f"{query_model_cls.model_name} ({dialogue_model.model_name})",
                        )
                    )
        else:
            for query_model_cls in query_based_model_list:
                for single_doc_model in single_doc_model_list:
                    full_query_model = (
                        query_model_cls(
                            model_backend=single_doc_model,
                            data=get_lxr_train_set(dataset),
                        )
                        if single_doc_model == LexRankModel
                        else query_model_cls(model_backend=single_doc_model)
                    )
                    matching_models.append(
                        (
                            full_query_model,
                            f"{query_model_cls.model_name} ({single_doc_model.model_name})",
                        )
                    )
        return matching_models

    if dataset.is_multi_document:
        for multi_doc_model_cls in multi_doc_model_list:
            for single_doc_model in single_doc_model_list:
                full_multi_doc_model = (
                    multi_doc_model_cls(
                        model_backend=single_doc_model, data=get_lxr_train_set(dataset)
                    )
                    if single_doc_model == LexRankModel
                    else multi_doc_model_cls(model_backend=single_doc_model)
                )
                matching_models.append(
                    (
                        full_multi_doc_model,
                        f"{multi_doc_model_cls.model_name} ({single_doc_model.model_name})",
                    )
                )
        return matching_models

    if dataset.is_dialogue_based:
        return list(
            map(
                lambda db_model: (db_model, db_model.model_name),
                dialogue_based_model_instances,
            )
        )

    return list(
        map(lambda s_model: (s_model, s_model.model_name), single_doc_model_instances)
    )