import cv2
import sys
import argparse
import os
import glob
import time
import numpy as np
import tqdm

from PIL import Image
from pano_lsd_align import (
    panoEdgeDetection, 
    rotatePanorama
)

if __name__ == '__main__':
    # disable OpenCV3's non thread safe OpenCL option
    cv2.ocl.setUseOpenCL(False)

    parser = argparse.ArgumentParser()
    parser.add_argument('--s2d3d_root', required=True)
    parser.add_argument('--out_root', required=True)
    parser.add_argument('--qError', default=0.7, type=float)
    parser.add_argument('--refineIter', default=3, type=int)
    args = parser.parse_args()

    for split in tqdm.tqdm(['train', 'test'], desc="Splits:"):
        #ext = "*.jpg"
        ext = "*_rgb.png"
        filenames =glob.glob(os.path.join(args.s2d3d_root, split, ext))
        for filename in tqdm.tqdm(filenames, desc=f"{split} split:"):
            img_ori = np.array(Image.open(filename).resize((1024, 512)))
            olines, vp, views, edges, panoEdge, score, angle = panoEdgeDetection(img_ori,
                                     refineIter=args.refineIter)
            # panoEdge = (panoEdge > 0)
            # print('Vanishing point:')
            # for v in vp[2::-1]:
            #     print('%.6f %.6f %.6f' % tuple(v))
            # edg = rotatePanorama(panoEdge.astype(np.float64), vp[2::-1])

            depth = cv2.imread(filename.replace('_rgb', '_depth'), cv2.IMREAD_ANYDEPTH)
            depth = cv2.resize(depth, (1024, 512), interpolation=cv2.INTER_NEAREST)
            depth = depth.astype(np.float32) * (1.0 / 512.0)

            img = rotatePanorama(img_ori / 255.0, vp[2::-1], order=1)
            depth = rotatePanorama(depth[:, :, np.newaxis], vp[2::-1], order=0)

            # one = img.copy() * 0.5
            # one[(edg > 0.5).sum(-1) > 0] = 0
            # one[edg[..., 0] > 0.5, 0] = 1
            # one[edg[..., 1] > 0.5, 1] = 1
            # one[edg[..., 2] > 0.5, 2] = 1
            # Image.fromarray((edg * 255).astype(np.uint8)).save('%s_edg.png' % args.o_prefix)
            Image.fromarray((img * 255).astype(np.uint8)).save(
                os.path.join(args.out_root, split, os.path.basename(filename).replace("_rgb", "_color"))
            )
            cv2.imwrite(os.path.join(args.out_root, split,
                os.path.basename(filename).replace("_rgb", "_depth").replace('png', 'exr')
            ), depth.astype(np.float32).squeeze())
            # Image.fromarray((one * 255).astype(np.uint8)).save('%s_one.png' % args.o_prefix)
