#!/usr/bin/env python

import os.path as osp
from glob import glob

from PIL import Image
from tqdm import tqdm

from common import (get_default_parser, add_crop_options, crop_patches,
                    create_file_list, copy_dataset, create_label_list,
                    get_path_tuples)

# According to the official doc(https://github.com/CAPTAIN-WHU/iSAID_Devkit), 
# the files should be organized as follows:
# 
# iSAID
# ├── test
# │   └── images
# │       ├── P0006.png
# │       └── ...
# │       └── P0009.png
# ├── train
# │   └── images
# │       ├── P0002_instance_color_RGB.png
# │       ├── P0002_instance_id_RGB.png
# │       ├── P0002.png
# │       ├── ...
# │       ├── P0010_instance_color_RGB.png
# │       ├── P0010_instance_id_RGB.png
# │       └── P0010.png
# └── val
#     └── images
#         ├── P0003_instance_color_RGB.png
#         ├── P0003_instance_id_RGB.png
#         ├── P0003.png
#         ├── ...
#         ├── P0004_instance_color_RGB.png
#         ├── P0004_instance_id_RGB.png
#         └── P0004.png

CLASSES = ('background', 'ship', 'storage_tank', 'baseball_diamond',
           'tennis_court', 'basketball_court', 'ground_track_field', 'bridge',
           'large_vehicle', 'small_vehicle', 'helicopter', 'swimming_pool',
           'roundabout', 'soccer_ball_field', 'plane', 'harbor')
# Refer to https://github.com/Z-Zheng/FarSeg/blob/master/data/isaid.py
COLOR_MAP = [[0, 0, 0], [0, 0, 63], [0, 191, 127], [0, 63, 0], [0, 63, 127],
             [0, 63, 191], [0, 63, 255], [0, 127, 63], [0, 127, 127],
             [0, 0, 127], [0, 0, 191], [0, 0, 255], [0, 63, 63], [0, 127, 191],
             [0, 127, 255], [0, 100, 155]]
SUBSETS = ('train', 'val')
SUBDIR = 'images'
FILE_LIST_PATTERN = "{subset}.txt"
LABEL_LIST_NAME = "labels.txt"
URL = ""


def flatten(nested_list):
    flattened_list = []
    for ele in nested_list:
        if isinstance(ele, list):
            flattened_list.extend(flatten(ele))
        else:
            flattened_list.append(ele)
    return flattened_list


def rgb2mask(rgb):
    palette = flatten(COLOR_MAP)
    # Pad with zero
    palette = palette + [0] * (256 * 3 - len(palette))
    ref = Image.new(mode='P', size=(1, 1))
    ref.putpalette(palette)
    mask = rgb.quantize(palette=ref, dither=0)
    return mask


if __name__ == '__main__':
    parser = get_default_parser()
    parser.add_argument(
        '--crop_size', type=int, help="Size of cropped patches.", default=800)
    parser.add_argument(
        '--crop_stride',
        type=int,
        help="Stride of sliding windows when cropping patches. `crop_size` will be used only if `crop_size` is not None.",
        default=600)
    args = parser.parse_args()

    out_dir = osp.join(args.out_dataset_dir,
                       osp.basename(osp.normpath(args.in_dataset_dir)))

    assert args.crop_size is not None
    # According to https://github.com/CAPTAIN-WHU/iSAID_Devkit/blob/master/preprocess/split.py
    # Set keep_last=True
    crop_patches(
        args.crop_size,
        args.crop_stride,
        data_dir=args.in_dataset_dir,
        out_dir=out_dir,
        subsets=SUBSETS,
        subdirs=(SUBDIR, ),
        glob_pattern='*.png',
        max_workers=8,
        keep_last=True)

    for subset in SUBSETS:
        path_tuples = []
        print(f"Processing {subset} labels...")
        for im_subdir in tqdm(glob(osp.join(out_dir, subset, SUBDIR, "*/"))):
            im_name = osp.basename(im_subdir[:-1])  # Strip trailing '/'
            if '_' in im_name:
                # Do not process labels
                continue
            mask_subdir = osp.join(out_dir, subset, SUBDIR,
                                   im_name + '_instance_color_RGB')
            for mask_path in glob(osp.join(mask_subdir, '*.png')):
                # Convert RGB files to mask files (pseudo color)
                rgb = Image.open(mask_path).convert('RGB')
                mask = rgb2mask(rgb)
                # Write to the original location
                mask.save(mask_path)
            path_tuples.extend(
                get_path_tuples(
                    im_subdir,
                    mask_subdir,
                    glob_pattern='*.png',
                    data_dir=args.out_dataset_dir))
        path_tuples.sort()

        file_list = osp.join(
            args.out_dataset_dir, FILE_LIST_PATTERN.format(subset=subset))
        create_file_list(file_list, path_tuples)
        print(f"Write file list to {file_list}.")

    label_list = osp.join(args.out_dataset_dir, LABEL_LIST_NAME)
    create_label_list(label_list, CLASSES)
    print(f"Write label list to {label_list}.")