"""
This script works on synthetic data, all HDR frames with a fixed exposure.
It removes 10 frames from the training set to use it for colmap calibration. 
This subset is identified with the "colmap" suffix, where both left and right have the same exposure.
The rest of the frames are used for two-exposure training (e1 and e2).

The masks are generated by the script lantern_scripts/combine_left_right.py, where 0 means unknown, and 1 means known.

Quantization and gamma-tonemapping are used on all images to simulate a camera.

input: left, right in HDR

output (all in LDR):
left_colmap, left_colmap_mask, 
right_colmap, right_colmap_mask, 
left_e1, left_e1_mask,
right_e2, right_e2_mask
"""
import argparse
import os

import cv2
import numpy as np
from tqdm import tqdm

from lantern_scripts.combine_left_right import apply_the_exposure_get_weights, make_weights_binary
from lantern_scripts.tonemap import tonemap
from PIL import Image

CALIB_FRAMES = [0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20]
# all the rest is used for two-exposures

def process_one_image(pixels, out_dir, name, has_exposure, exposure_value, zmin, zmax):
    if has_exposure:
        e1_ldr, e1_hdr, weights_e1 = apply_the_exposure_get_weights(pixels, exposure_value)
        weights_e1_unit = make_weights_binary(e1_ldr, weights_e1, zmin, zmax)
        mask_e1 = weights_e1_unit[:, :, 0] & weights_e1_unit[:, :, 1] & weights_e1_unit[:, :, 2]
        e1_mask_addr = os.path.join(out_dir, f'{name}_mask.png')
        # e1_mask_addr = os.path.join(args.out_dir, pano_name + '_e1.png')
        Image.fromarray(mask_e1).save(e1_mask_addr)
    else:
        raise NotImplementedError
    # out_addr = os.path.join(args.out_dir, pano_name + '_e1.exr')

    e1_ldr_tonemapped = tonemap(e1_ldr)
    out_addr = os.path.join(out_dir, f'{name}.png')
    os.makedirs(os.path.dirname(out_addr), exist_ok=True)

    cv2.imwrite(out_addr, e1_ldr_tonemapped) 

def process_images(images, out_dir, prefix, has_exposure, exposure_value, zmin, zmax):
    os.makedirs(out_dir)
    print("Processing images in: ", out_dir)
    images = sorted(images)
    for pano_addr in tqdm(images):
        # original_pano = EnvironmentMap(pano_addr, 'latlong')
        # pixels = original_pano.data[:, :, 0:3]
        original_pano = np.array(cv2.imread(str(pano_addr), cv2.IMREAD_UNCHANGED)).astype("float32")
        pixels = original_pano[:, :, 0:3]

        pano_file = os.path.basename(pano_addr)
        pano_name =  pano_file[:-4]
        name = prefix + pano_name

        process_one_image(pixels, out_dir, name, has_exposure, exposure_value, zmin, zmax)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--left_dir', type=str, default='/mnt/workspace/share_w_Junming/Dual_Cameras/left/')
    parser.add_argument('--right_dir', type=str, default='/mnt/workspace/share_w_Junming/Dual_Cameras/right/')
    parser.add_argument('--out_dir', type=str, default='/mnt/workspace/share_w_Junming/Dual_Cameras/lantern-HDR/')
    parser.add_argument('--has_exposure', action='store_true')
    parser.add_argument('--exposure1', type=float, default=1.0)
    parser.add_argument('--exposure2', type=float, default=0.009)
    args = parser.parse_args()
    
    assert os.path.isdir(args.left_dir + '/HDR_Normal/'), 'The data folder for left set is missing!'
    assert os.path.isdir(args.right_dir + '/HDR_Normal/'), 'The data folder for right set is missing!'

    assert not os.path.isdir(args.out_dir), 'The output folder exists!! Take care of it first!'
    


    paths_left_all = args.left_dir + '/HDR_Normal/'
    images_left_all = [os.path.join(paths_left_all, f) for f in os.listdir(paths_left_all) if os.path.isfile(os.path.join(paths_left_all, f)) and f.endswith('.exr')]
    # images_left_all.sort()
    # sort images named like: r_00.exr, r_01.exr, ..., r_99.exr, r_100.exr, r_101.exr, ..., r_199.exr, r_200.exr, ...
    # in order
    images_left_all.sort(key=lambda f: int(''.join(filter(str.isdigit, f))))
    images_left_calib = [images_left_all[i] for i in CALIB_FRAMES]
    images_left_2exp = [images_left_all[i] for i in range(len(images_left_all)) if i not in CALIB_FRAMES]

    paths_right_all = args.right_dir + '/HDR_Normal/'
    images_right_all = [os.path.join(paths_right_all, f) for f in os.listdir(paths_right_all) if os.path.isfile(os.path.join(paths_right_all, f)) and f.endswith('.exr')]
    # images_right_all.sort()
    # sort images named like: r_00.exr, r_01.exr, ..., r_99.exr, r_100.exr, r_101.exr, ..., r_199.exr, r_200.exr, ...
    # in order
    images_right_all.sort(key=lambda f: int(''.join(filter(str.isdigit, f))))
    images_right_calib = [images_right_all[i] for i in CALIB_FRAMES]
    images_right_2exp = [images_right_all[i] for i in range(len(images_right_all)) if i not in CALIB_FRAMES]

    process_images(images_left_calib, args.out_dir + '/left_colmap_baseline/', 'lhs_', args.has_exposure, args.exposure1, 0.0, 0.95)
    process_images(images_right_calib, args.out_dir + '/right_colmap_baseline/', 'rhs_', args.has_exposure, args.exposure1, 0.0, 0.95)

    process_images(images_left_2exp, args.out_dir + '/left_e1/', 'lhs_', args.has_exposure, args.exposure1, 0.0, 0.95)
    process_images(images_right_2exp, args.out_dir + '/right_e2/', 'rhs_', args.has_exposure, args.exposure2, 0.2, 1.0)

    print("Done!")

if __name__ == "__main__":
    main()