import itertools

import detectron2.data.transforms as T
from detectron2.config import LazyCall as L
from detectron2.data import (
    build_detection_test_loader,
    build_detection_train_loader,
    get_detection_dataset_dicts,
)
from detectron2.evaluation import COCOEvaluator
from detrex.data import DetrDatasetMapper
from detrex.data.datasets import register_custom_ovd_instances
from omegaconf import OmegaConf
import json

dataloader = OmegaConf.create()

# AerialMaritimeDrone

dataset_root_dir = "path_for_ArTaxOr/ArTaxOr/"
dataset_ano_dir = dataset_root_dir + "annotations/"
dataset_train_dir = dataset_root_dir + "train/"
dataset_valid_dir = dataset_root_dir + "test/"
train_ano_file = dataset_ano_dir + "train.json"
valid_ano_file = dataset_ano_dir + "test.json"


# Case 0: If you want to define it by yourself, you can change it on ovdino/detrex/data/datasets/custom_ovd.py, you need to uncomment the code (ovdino/detrex/data/datasets/__init__.py L21) first.
# Case 1 (Recommend): If you follow the coco format, you need uncomment and change the following code.
# # 1. Define custom_meta_info, just a example, you need to change it to your own.

with open(train_ano_file, "r") as f:
    data = json.load(f)

# 提取 categories
categories = data["categories"]

# 提取 thing_classes（按 id 排序）
thing_classes = [c["name"] for c in sorted(categories, key=lambda x: x["id"])]

# 构造 id 映射：原始 id → 连续 id（Detectron2 要求连续）
thing_dataset_id_to_contiguous_id = {
    c["id"]: i for i, c in enumerate(sorted(categories, key=lambda x: x["id"]))
}

import random
random.seed(42)  # 固定种子，保证每次一样
thing_colors = [
    [random.randint(0, 255)/255, random.randint(0, 255)/255, random.randint(0, 255)/255]
    for _ in thing_classes
]

meta_info = {
    "thing_dataset_id_to_contiguous_id": thing_dataset_id_to_contiguous_id,
    "thing_classes": thing_classes,
    "thing_colors" : thing_colors
}

# # 2. Register custom train dataset.
register_custom_ovd_instances(
    "ArTaxOr_train",  # dataset_name
    meta_info,
    train_ano_file,  # annotations_json_file
    dataset_train_dir,  # image_root
    7,  # number_of_classes, default: 2. You also need to change model.num_classes in the ovdino/projects/ovdino/configs/ovdino_swin_tiny224_bert_base_ft_custom_24ep.py#L37.
    "identity",  # template, default: full
)
# # 3. Register custom val dataset.
register_custom_ovd_instances(
    "ArTaxOr_val",  # dataset_name
    meta_info,
    valid_ano_file,  # annotations_json_file
    dataset_valid_dir,  # image_root
    7,  # number_of_classes, default: 2. You also need to change model.num_classes in the ovdino/projects/ovdino/configs/ovdino_swin_tiny224_bert_base_ft_custom_24ep.py#L37.
    "identity",  # template, default: full
)

dataloader.train = L(build_detection_train_loader)(
    dataset=L(get_detection_dataset_dicts)(names="ArTaxOr_train"),
    mapper=L(DetrDatasetMapper)(
        augmentation=[
            L(T.RandomFlip)(),
            L(T.ResizeShortestEdge)(
                short_edge_length=(
                    480,
                    512,
                    544,
                    576,
                    608,
                    640,
                    672,
                    704,
                    736,
                    768,
                    800,
                ),
                max_size=1333,
                sample_style="choice",
            ),
        ],
        augmentation_with_crop=[
            L(T.RandomFlip)(),
            L(T.ResizeShortestEdge)(
                short_edge_length=(400, 500, 600),
                sample_style="choice",
            ),
            L(T.RandomCrop)(
                crop_type="absolute_range",
                crop_size=(384, 600),
            ),
            L(T.ResizeShortestEdge)(
                short_edge_length=(
                    480,
                    512,
                    544,
                    576,
                    608,
                    640,
                    672,
                    704,
                    736,
                    768,
                    800,
                ),
                max_size=1333,
                sample_style="choice",
            ),
        ],
        is_train=True,
        mask_on=False,
        img_format="RGB",
    ),
    total_batch_size=16,
    num_workers=4,
)

dataloader.test = L(build_detection_test_loader)(
    dataset=L(get_detection_dataset_dicts)(
        names="ArTaxOr_val", filter_empty=False
    ),
    mapper=L(DetrDatasetMapper)(
        augmentation=[
            L(T.ResizeShortestEdge)(
                short_edge_length=800,
                max_size=1333,
            ),
        ],
        augmentation_with_crop=None,
        is_train=False,
        mask_on=False,
        img_format="RGB",
    ),
    num_workers=4,
)

dataloader.evaluator = L(COCOEvaluator)(
    dataset_name="${..test.dataset.names}",
)
