from pathlib import Path
from typing import Any, Dict, List, Union

import numpy as np

from polimorfo.datasets import CocoDataset, InstanceCoco, SemanticCocoDataset


def rename_annotation_keys(ann: Dict[str, Union[str, int, list, np.ndarray]]) -> Dict:
    """Renames keys for a Coco annotation in order to feed add_annotation method properly.
    Args:
        ann (Dict[str, Union[str, int, list, np.ndarray]]): Coco annotation
    Returns:
        Dict: same as input annotation with keys renamed.
    """
    key_map = {
        "category_id": "cat_id",
        "segmentation": "segmentation",
        "area": "area",
        "iscrowd": "is_crowd",
        "score": "score",
        "bbox": "bbox",
    }
    assert set(list(ann.keys()) + ["score"]) >= (
        set(key_map.keys())
    ), f"{set(ann.keys())}, {set(key_map.keys())}"
    return {v: ann.get(k, None) for k, v in key_map.items()}


def get_img_meta_by_name(ds, image_name):
    try:
        return [
            img_meta
            for img_meta in ds.imgs.values()
            if img_meta["file_name"] == image_name
        ][0]
    except:
        log.warning("image not present in dataset.")
        return None


def merge_datasets(
    datasets: List[Union[SemanticCocoDataset, CocoDataset, InstanceCoco]],
    dest_path: Union[str, Path] = None,
) -> CocoDataset:

    # compute category to cat_idx map for each dataset
    cat_to_idx = [
        {v["name"]: k for k, v in ds_item.cats.items()} for ds_item in datasets
    ]
    category_names = [list(ctx.keys()) for ctx in cat_to_idx]
    flat_category_name = sorted(
        [item for sublist in category_names for item in sublist]
    )

    # collect all category names
    merged_cat_names = set(flat_category_name)

    merged_idx_to_cat = dict(enumerate(sorted(merged_cat_names), start=1))
    merged_cat_to_idx = {v: k for k, v in merged_idx_to_cat.items()}

    # map categories for each dataset to merged one
    cat_map = [{v: merged_cat_to_idx[k] for k, v in ctx.items()} for ctx in cat_to_idx]

    # initializize merged dataset
    out_ds = CocoDataset(dest_path)
    for _, cat_name in sorted(merged_idx_to_cat.items()):
        out_ds.add_category(cat_name, "thing")

    for ds_item, cat_map_item in zip(datasets, cat_map):
        img_names = {a["file_name"]: a["id"] for a in out_ds.imgs.values()}
        # add images
        for img_meta in ds_item.imgs.values():
            # remark: if datasets are not disjoint the might bo overlap here to take care of, i.e. possible multiple definitions of images
            if img_meta["file_name"] not in img_names:
                img_idx = out_ds.add_image(**img_meta)
                img_names[img_meta["file_name"]] = img_idx
            else:
                img_idx = img_names[img_meta["file_name"]]

            orig_img_idx = img_meta["id"]
            anns = ds_item.get_annotations(img_idx=orig_img_idx)
            # add anns
            for ann in anns:
                mapped_ann = rename_annotation_keys(ann)
                mapped_ann["img_id"] = img_idx  # new image idx
                mapped_ann["cat_id"] = cat_map_item[
                    mapped_ann["cat_id"]
                ]  # new category index
                out_ds.add_annotation(**mapped_ann)

    out_ds.reindex()
    return out_ds
