import os
import json
import shutil
import argparse
import random
import pandas as pd

FOLD_CNT = 10


def generate_fold(args):
    df = pd.DataFrame(columns=['img_id', 'fold'])
    json_file_path = args.ann_file
    assert os.path.exists(json_file_path), f'json file not found: {json_file_path}'
    with open(json_file_path, 'r') as f:
        json_data = json.load(f)
    
    # reproducable random shuffle
    random.seed(1)
    random.shuffle(json_data['images'])
    fold_size = len(json_data['images']) // FOLD_CNT
    for image in json_data['images']:
        df.loc[len(df), df.columns] = [image['id'], min(len(df) // fold_size, FOLD_CNT - 1)]
    df.to_csv(args.fold_file, index=False)


def generate_fold_app(args):
    df = pd.DataFrame(columns=['img_id', 'fold'])
    json_file_path = args.ann_file
    assert os.path.exists(json_file_path), f'json file not found: {json_file_path}'
    with open(json_file_path, 'r') as f:
        json_data = json.load(f)

    apps = set()

    for image in json_data['images']:
        appid = str(image['id'])[:-3]
        apps.add(appid)
    
    apps = list(apps)
    bel_app = {}
    
    # reproducable random shuffle
    apps = sorted(apps)
    random.seed(1)
    random.shuffle(apps)
    fold_size = len(apps) // FOLD_CNT
    for app in apps:
        bel_app[app] = min(len(bel_app) // fold_size, FOLD_CNT - 1)
    for image in json_data['images']:
        appid = str(image['id'])[:-3]
        df.loc[len(df), df.columns] = [image['id'], bel_app[appid]]
    
    df = df.sort_values(by='fold')
    df.to_csv(args.fold_file, index=False)


def split(args):
    json_file_path = args.ann_file
    image_path = args.img_dir if args.img_dir and os.path.exists(args.img_dir) else None
    fold_file = args.fold_file
    output_path = args.output_path
    
    assert os.path.exists(json_file_path), f'json file not found: {json_file_path}'
    # assert os.path.exists(image_path), f'image path not found: {image_path}'
    assert os.path.exists(fold_file), f'fold file not found: {fold_file}'

    with open(json_file_path, 'r') as f:
        json_data = json.load(f)
    
    df = pd.read_csv(fold_file)
    
    train_dataset = {'images': list(), 'categories': json_data['categories'], 'annotations': list()}
    val_dataset = {'images': list(), 'categories': json_data['categories'], 'annotations': list()}
    test_dataset = {'images': list(), 'categories': json_data['categories'], 'annotations': list()}

    train_folds = [int(fold) for fold in args.train_folds.split(',')] if args.train_folds else []
    val_folds = [int(fold) for fold in args.val_folds.split(',')] if args.val_folds else []
    test_folds = [int(fold) for fold in args.test_folds.split(',')] if args.test_folds else []

    print(f'train folds: {train_folds}')
    print(f'val folds: {val_folds}')
    print(f'test folds: {test_folds}')
    
    train_imgid = set()
    val_imgid = set()
    test_imgid = set()

    for _, row in df.iterrows():
        if row['fold'] in train_folds:
            train_imgid.add(row['img_id'])
        elif row['fold'] in val_folds:
            val_imgid.add(row['img_id'])
        elif row['fold'] in test_folds:
            test_imgid.add(row['img_id'])
        else:
            raise ValueError(f'fold not found: {row["fold"]}')

    for image in json_data['images']:
        if image['id'] in train_imgid:
            train_dataset['images'].append(image)
        elif image['id'] in val_imgid:
            val_dataset['images'].append(image)
        elif image['id'] in test_imgid:
            test_dataset['images'].append(image)
        else:
            raise ValueError(f'image id not found: {image["id"]}')
 
    for annotation in json_data['annotations']:
        if annotation['image_id'] in train_imgid:
            train_dataset['annotations'].append(annotation)
        elif annotation['image_id'] in val_imgid:
            val_dataset['annotations'].append(annotation)
        elif annotation['image_id'] in test_imgid:
            test_dataset['annotations'].append(annotation)
        else:
            raise ValueError(f'annotation image_id not found: {annotation["image_id"]}')
    
    print(f'train dataset: {len(train_dataset["images"])} images, {len(train_dataset["annotations"])} annotations')
    print(f'val dataset: {len(val_dataset["images"])} images, {len(val_dataset["annotations"])} annotations')
    print(f'test dataset: {len(test_dataset["images"])} images, {len(test_dataset["annotations"])} annotations')
    
    print(os.path.abspath(output_path))
    train_image_path = os.path.join(output_path, 'images', 'instances_train2017')
    val_image_path = os.path.join(output_path, 'images', 'instances_val2017')
    test_image_path = os.path.join(output_path, 'images', 'instances_test2017')
    os.makedirs(output_path, exist_ok=True)
    os.makedirs(train_image_path, exist_ok=True)
    os.makedirs(val_image_path, exist_ok=True)
    os.makedirs(test_image_path, exist_ok=True)
    os.makedirs(os.path.join(output_path, 'annotations'), exist_ok=True)
    
    with open(os.path.join(output_path, 'annotations', 'instances_train2017.json'), 'w') as f:
        json.dump(train_dataset, f, indent=4)
    with open(os.path.join(output_path, 'annotations', 'instances_val2017.json'), 'w') as f:
        json.dump(val_dataset, f, indent=4)
    with open(os.path.join(output_path, 'annotations', 'instances_test2017.json'), 'w') as f:
        json.dump(test_dataset, f, indent=4)

    if image_path:
        for image in train_dataset['images']:
            shutil.copy(os.path.join(image_path, image['file_name']), train_image_path)
        for image in val_dataset['images']:
            shutil.copy(os.path.join(image_path, image['file_name']), val_image_path)
        for image in test_dataset['images']:
            shutil.copy(os.path.join(image_path, image['file_name']), test_image_path)


def main(args):
    if args.gen:
       generate_fold_app(args)
    split(args)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--ann_file', type=str, required=True)
    parser.add_argument('--img_dir', type=str, required=True)
    parser.add_argument('--output_path', type=str, default='./coco_split')
    parser.add_argument('--train_folds', type=str, default=None)
    parser.add_argument('--val_folds', type=str, default=None)
    parser.add_argument('--test_folds', type=str, default=None)
    parser.add_argument('--gen', action='store_true')
    parser.add_argument('--fold_file', type=str, required=True)
    args = parser.parse_args()
    main(args)