#!/usr/bin/env python3

# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.

import csv
import hashlib
import os
import shutil

import numpy as np
from PIL import Image

SIZE = 32

np.random.seed(0)

PATHS = {
    'images': [
        'tiles/Site1/Site1_RGB_0_0_0_4000_4000.png',
        'tiles/Site2/Site2_RGB_0_0_0_4000_4000.png',
    ],
    'annotation': 'mapping/final_dataset.csv',
}


def create_annotation(path: str, img_paths: list[str]) -> None:
    cols = ['img_path', 'xmin', 'ymin', 'xmax', 'ymax', 'group', 'AGB']
    data = []
    for img_path in img_paths:
        data.append(
            [os.path.basename(img_path), 0, 0, SIZE / 2, SIZE / 2, 'banana', 6.75]
        )
        data.append(
            [os.path.basename(img_path), SIZE / 2, SIZE / 2, SIZE, SIZE, 'cacao', 6.75]
        )

    with open(path, 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(cols)
        writer.writerows(data)


def create_img(path: str) -> None:
    Z = np.random.rand(SIZE, SIZE, 3) * 255
    img = Image.fromarray(Z.astype('uint8')).convert('RGB')
    img.save(path)


if __name__ == '__main__':
    data_root = 'reforesTree'

    # remove old data
    if os.path.isdir(data_root):
        shutil.rmtree(data_root)

    # create imagery
    for path in PATHS['images']:
        os.makedirs(os.path.join(data_root, os.path.dirname(path)), exist_ok=True)
        create_img(os.path.join(data_root, path))

    # create annotations
    os.makedirs(
        os.path.join(data_root, os.path.dirname(PATHS['annotation'])), exist_ok=True
    )
    create_annotation(os.path.join(data_root, PATHS['annotation']), PATHS['images'])

    # compress data
    shutil.make_archive(data_root, 'zip', data_root)

    # Compute checksums
    with open(data_root + '.zip', 'rb') as f:
        md5 = hashlib.md5(f.read()).hexdigest()
        print(f'{data_root}: {md5}')
