from typing import Union
from collections import defaultdict

from data import load_jsonl


all_experiments = load_jsonl('experiments.jsonl')

experiment_ids_by_product_id = defaultdict(list)
for experiment in all_experiments:
    product_id = experiment.get('product_id')
    if product_id is not None:
        experiment_ids_by_product_id[product_id].append(experiment["id"])

all_products = list(experiment_ids_by_product_id.keys())

def get_experiment_by_id(experiment_id: int) -> Union[dict, None]:
    for experiment in all_experiments:
        if experiment['id'] == experiment_id:
            return experiment
    return None


def filter_experiments(
    experiments: Union[list[dict], None] = None,
    product_id: Union[str, list[str], None] = None,
    experiment_id: Union[int, list[int], None] = None,
) -> list[dict]:
    filtered_experiments = experiments
    if filtered_experiments is None:
        filtered_experiments = all_experiments
    if product_id is not None:
        if isinstance(product_id, str):
            product_id = {product_id}
        filtered_experiments = [exp for exp in filtered_experiments if exp.get('product_id') in product_id]
    if experiment_id is not None:
        if isinstance(experiment_id, int):
            experiment_id = {experiment_id}
        filtered_experiments = [exp for exp in filtered_experiments if exp['id'] in experiment_id]
    return filtered_experiments


def get_experiment(
    experiments: Union[list[dict], None] = None,
    product_id: Union[str, list[str], None] = None,
    experiment_id: Union[int, None] = None,
) -> Union[dict, None]:
    filtered_experiments = filter_experiments(
        experiments=experiments,
        product_id=product_id,
        experiment_id=experiment_id
    )
    if filtered_experiments:
        return filtered_experiments[0]
    else:
        return None