"""
Standalone seasonal data loader for ForestPersons dataset.
This file provides dataset registration functions without modifying the original croptrain module.
"""

import os
import numpy as np
from detectron2.data import DatasetCatalog, MetadataCatalog
from detectron2.structures import BoxMode

dataset_root = "/mnt/home2/annonymous/ForestPersons_v3"


def load_forestpersons_coco_json(json_file, image_root, dataset_name):
    """
    Load a COCO-format JSON file and add required keys for crop training.
    Similar to detectron2's load_coco_json but adds custom keys required by the dataset mapper.
    """
    from pycocotools import coco as coco_api
    
    coco_api = coco_api.COCO(json_file)
    
    # sort indices for reproducible results
    img_ids = sorted(coco_api.imgs.keys())
    imgs = coco_api.loadImgs(img_ids)
    anns = [coco_api.imgToAnns[img_id] for img_id in img_ids]
    
    imgs_anns = list(zip(imgs, anns))
    
    dataset_dicts = []
    ann_keys = ["iscrowd", "bbox", "category_id"]
    
    for (img_dict, anno_dict_list) in imgs_anns:
        record = {}
        record["file_name"] = os.path.join(image_root, img_dict["file_name"])
        record["height"] = img_dict["height"]
        record["width"] = img_dict["width"]
        record["image_id"] = img_dict["id"]
        
        # Add required keys for crop training
        record["full_image"] = True
        record["two_stage_crop"] = False  
        record["crop_area"] = np.array([-1, -1, -1, -1], dtype=np.float32)
        
        objs = []
        for anno in anno_dict_list:
            assert anno["image_id"] == record["image_id"]
            
            obj = {key: anno[key] for key in ann_keys if key in anno}
            
            if "bbox" in obj and len(obj["bbox"]) == 0:
                continue
                
            # Convert bbox format
            obj['bbox'] = BoxMode.convert(obj['bbox'], BoxMode.XYWH_ABS, BoxMode.XYXY_ABS)    
            obj["bbox_mode"] = BoxMode.XYXY_ABS
            
            objs.append(obj)
            
        record["annotations"] = objs
        dataset_dicts.append(record)
        
    return dataset_dicts


def register_forestpersons_summer(cfg=None):
    """
    Register the ForestPersons Summer subset dataset with Detectron2.
    
    Args:
        cfg: Configuration object (unused but kept for consistency)
    """
    # Register training and validation datasets for summer subset using custom loader
    DatasetCatalog.register(
        "missing_person_summer_train", 
        lambda: load_forestpersons_coco_json(
            f"{dataset_root}/annotations/train_summer.json", 
            f"{dataset_root}/images",
            "missing_person_summer_train"
        )
    )
    DatasetCatalog.register(
        "missing_person_summer_test", 
        lambda: load_forestpersons_coco_json(
            f"{dataset_root}/annotations/val_summer.json", 
            f"{dataset_root}/images",
            "missing_person_summer_test"
        )
    )
    
    # Set metadata
    MetadataCatalog.get("missing_person_summer_train").set(
        thing_classes=["person"],
        thing_dataset_id_to_contiguous_id={1: 0},
        evaluator_type="coco"
    )
    MetadataCatalog.get("missing_person_summer_test").set(
        thing_classes=["person"],
        thing_dataset_id_to_contiguous_id={1: 0},
        evaluator_type="coco"
    )


def register_forestpersons_winter(cfg=None):
    """
    Register the ForestPersons Winter subset dataset with Detectron2.
    
    Args:
        cfg: Configuration object (unused but kept for consistency)
    """
    # Register training and validation datasets for winter subset using custom loader
    DatasetCatalog.register(
        "missing_person_winter_train", 
        lambda: load_forestpersons_coco_json(
            f"{dataset_root}/annotations/train_winter.json", 
            f"{dataset_root}/images",
            "missing_person_winter_train"
        )
    )
    DatasetCatalog.register(
        "missing_person_winter_test", 
        lambda: load_forestpersons_coco_json(
            f"{dataset_root}/annotations/val_winter.json", 
            f"{dataset_root}/images",
            "missing_person_winter_test"
        )
    )
    
    # Set metadata
    MetadataCatalog.get("missing_person_winter_train").set(
        thing_classes=["person"],
        thing_dataset_id_to_contiguous_id={1: 0},
        evaluator_type="coco"
    )
    MetadataCatalog.get("missing_person_winter_test").set(
        thing_classes=["person"],
        thing_dataset_id_to_contiguous_id={1: 0},
        evaluator_type="coco"
    )
    
    print("ForestPersons Winter dataset registered successfully!")


def register_forestpersons_standing(cfg=None):
    """
    Register the ForestPersons Standing subset dataset with Detectron2.
    
    Args:
        cfg: Configuration object (unused but kept for consistency)
    """
    # Register training and validation datasets for standing subset using custom loader
    DatasetCatalog.register(
        "missing_person_standing_train", 
        lambda: load_forestpersons_coco_json(
            f"{dataset_root}/annotations/train_standing_all.json", 
            f"{dataset_root}/images",
            "missing_person_standing_train"
        )
    )
    DatasetCatalog.register(
        "missing_person_standing_test", 
        lambda: load_forestpersons_coco_json(
            f"{dataset_root}/annotations/val_standing_all.json", 
            f"{dataset_root}/images",
            "missing_person_standing_test"
        )
    )
    
    # Set metadata
    MetadataCatalog.get("missing_person_standing_train").set(
        thing_classes=["person"],
        thing_dataset_id_to_contiguous_id={1: 0},
        evaluator_type="coco"
    )
    MetadataCatalog.get("missing_person_standing_test").set(
        thing_classes=["person"],
        thing_dataset_id_to_contiguous_id={1: 0},
        evaluator_type="coco"
    )
    
    print("ForestPersons Standing dataset registered successfully!")


def register_forestpersons_combined(cfg=None):
    """
    Register both seasonal subsets together for comparison purposes.
    
    Args:
        cfg: Configuration object (unused but kept for consistency)
    """
    register_forestpersons_summer(cfg)
    register_forestpersons_winter(cfg)
    print("Both ForestPersons seasonal datasets registered successfully!")

def register_forestpersons_original(cfg=None):
    DatasetCatalog.register(
        "missing_person_original_train",
        lambda: load_forestpersons_coco_json(
            f"{dataset_root}/annotations/train.json", 
            f"{dataset_root}/images",
            "missing_person_original_train"
        )
    )
    DatasetCatalog.register(
        "missing_person_original_val",
        lambda: load_forestpersons_coco_json(
            f"{dataset_root}/annotations/val.json",
            f"{dataset_root}/images",
            "missing_person_original_val"
        )
    )
    DatasetCatalog.register(
        "missing_person_original_test",
        lambda: load_forestpersons_coco_json(
            f"{dataset_root}/annotations/test.json",
            f"{dataset_root}/images",
            "missing_person_original_test"
        )
    )

    MetadataCatalog.get("missing_person_original_train").set(
        thing_classes=["person"],
        thing_dataset_id_to_contiguous_id={1: 0},
        evaluator_type="coco"
    )
    MetadataCatalog.get("missing_person_original_val").set(
        thing_classes=["person"],
        thing_dataset_id_to_contiguous_id={1: 0},
        evaluator_type="coco"
    )
    MetadataCatalog.get("missing_person_original_test").set(
        thing_classes=["person"],
        thing_dataset_id_to_contiguous_id={1: 0},
        evaluator_type="coco"
    )
if __name__ == "__main__":
    # Test the registration functions
    print("Testing dataset registration...")
    register_forestpersons_combined()
    register_forestpersons_standing()
    print("All tests passed!")