"""Create Comparison views with different experiments"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/project/comparison.ipynb.

# %% auto 0
__all__ = ['logger']

# %% ../../nbs/project/comparison.ipynb 3
import typing as t
import logging

from fastcore.utils import patch
from tqdm import tqdm

from .core import Project
from ..model.notion_model import NotionModel
import ragas_experimental.model.notion_typing as nmt
from ..experiment import Experiment
from ..dataset import Dataset

# %% ../../nbs/project/comparison.ipynb 4
logger = logging.getLogger(__name__)

# %% ../../nbs/project/comparison.ipynb 5
# utility function to check if a model has a title property and get the name of the title property
@t.overload
def _get_title_property(
    model: NotionModel | t.Type[NotionModel], raise_exception: t.Literal[True] = True
) -> str: ...
@t.overload
def _get_title_property(
    model: NotionModel | t.Type[NotionModel], raise_exception: t.Literal[False] = False
) -> t.Optional[str]: ...
def _get_title_property(
    model: NotionModel | t.Type[NotionModel], raise_exception: bool = True
) -> t.Optional[str]:
    has_title = False
    for field in model._fields.keys():
        if isinstance(model._fields[field], nmt.Title):
            has_title = True
            title_property = field
            return title_property

    if not has_title:
        if raise_exception:
            raise ValueError("Model has no title property")
        else:
            return None

# %% ../../nbs/project/comparison.ipynb 8
def _validate_experiments(experiments: t.Sequence[Experiment]):
    # validate we have more than 2 experiments
    if len(experiments) < 2:
        raise ValueError("We need at least 2 experiments to compare")

    # validate that all experiments are of the same model
    top_exp = experiments[0]
    title_property = _get_title_property(top_exp.model)
    for exp in experiments:
        if not isinstance(exp, Experiment):
            raise ValueError("All experiments must be of type Experiment")
        if top_exp != exp.model:
            logger.warning(
                f"Experiments have different models: {top_exp.model} and {exp.model}"
            )
        if title_property != _get_title_property(exp.model):
            raise ValueError("All experiments must have the same title property.")

# %% ../../nbs/project/comparison.ipynb 12
def _model_to_dict(model: NotionModel) -> dict:
    # drop ID filed
    data = {}
    for field_name in model._fields.keys():
        if isinstance(model._fields[field_name], nmt.ID):
            continue
        data[field_name] = model.__getattribute__(field_name)
    return data

# %% ../../nbs/project/comparison.ipynb 14
def _combine_experiments(experiments: t.Sequence[Experiment]):
    """Group experiment rows by their title property value."""
    if not experiments:
        return []

    title_property: str = _get_title_property(experiments[0].model)

    # Create a dictionary to group rows by title value
    grouped_by_title = {}

    # Process each experiment
    for exp in experiments:
        for row in exp:
            title_value = getattr(row, title_property)

            # Create key if it doesn't exist
            if title_value not in grouped_by_title:
                grouped_by_title[title_value] = []

            # Add this row to the appropriate group
            row_dict = _model_to_dict(row)
            row_dict["experiment_name"] = exp.name
            grouped_by_title[title_value].append(row_dict)

    # Convert dictionary to list and add id_str
    result = []
    for i, (_, rows) in enumerate(grouped_by_title.items()):
        for row in rows:
            row["id_str"] = str(i)
        result.append(rows)

    return result

# %% ../../nbs/project/comparison.ipynb 17
@patch
def compare_experiments(
    self: Project,
    *experiments: Experiment,
):
    _validate_experiments(experiments)

    # create a combined Model with all the fields of the experiments
    class CombinedModel(NotionModel):
        id_str: str = nmt.Text()
        experiment_name: str = nmt.Text()

    for exp in experiments:
        for field in exp.model._fields.keys():
            if field not in CombinedModel._fields:
                CombinedModel._fields[field] = exp.model._fields[field]

    # create a new database with the combined model
    properties = {}
    for field in CombinedModel._fields.keys():
        properties.update(CombinedModel._fields[field]._to_notion_property())
    comparison_database_id = self._notion_backend.create_new_database(
        parent_page_id=self.comparisons_page_id,
        title=f"{' and '.join([exp.name for exp in experiments])}",
        properties=properties,
    )

    # make sure all experiments are synced to upstream
    for exp in experiments:
        exp.load()

    # group together by title property
    grouped_experiments = _combine_experiments(experiments)

    # append these to database
    for grouped_row in tqdm(grouped_experiments, desc="Uploading to Notion"):
        for row in grouped_row:
            combined_model_instance = CombinedModel(**row)
            self._notion_backend.create_page_in_database(
                database_id=comparison_database_id,
                properties=combined_model_instance.to_notion()["properties"],
            )
    # Get the URL for the created database
    # The format for Notion URLs is: https://www.notion.so/{database_id}
    notion_url = f"https://www.notion.so/{comparison_database_id.replace('-', '')}"

    return notion_url
