"""
Creates color+normal datasets based on annotation file.
The datasets can be used to train MLP, CycleGAN, Pix2Pix models.
"""
import os

import cv2
import hydra
import imageio
import numpy as np
import torch

from digit_depth.dataio.create_csv import (combine_csv, create_pixel_csv,
                                           create_train_test_csv)
from digit_depth.dataio.data_loader import data_loader
from digit_depth.dataio.generate_sphere_gt_normals import generate_sphere_gt_normals
from digit_depth.third_party import data_utils

base_path = os.path.abspath(os.path.dirname(os.path.dirname(__file__)))


@hydra.main(config_path=f"{base_path}/config", config_name="rgb_to_normal.yaml", version_base=None)
def main(cfg):
    normal_dataloader, normal_dataset = data_loader(
        dir_dataset=os.path.join(base_path, "images"), params=cfg.dataloader
    )
    dirs = [
        f"{base_path}/datasets/A/imgs",
        f"{base_path}/datasets/B/imgs",
        f"{base_path}/datasets/A/csv",
        f"{base_path}/datasets/B/csv",
        f"{base_path}/datasets/train_test_split",
    ]
    for dir in dirs:
        print(f"Creating directory: {dir}")
        os.makedirs(f"{dir}", exist_ok=True)
    # iterate over images
    img_idx = 0
    radius_bearing = np.int32(0.5 * 6.0 * cfg.mm_to_pixel)
    while img_idx < len(normal_dataset):
        # read img + annotations
        data = normal_dataset[img_idx]
        if cfg.dataloader.annot_flag:
            img, annot = data
            if annot.shape[0] == 0:
                img_idx = img_idx + 1
                continue
        else:
            img = data

        # get annotation circle params
        if cfg.dataloader.annot_flag:
            annot_np = annot.cpu().detach().numpy()
            center_y, center_x, radius_annot = (
                annot_np[0][1],
                annot_np[0][0],
                annot_np[0][2],
            )
        else:
            center_y, center_x, radius_annot = 0, 0, 0

        img_color_np = (img.permute(2, 1, 0).cpu().detach().numpy())  # (3,320,240) -> (240,320,3)

        # apply foreground mask
        fg_mask = np.zeros(img_color_np.shape[:2], dtype="uint8")
        fg_mask = cv2.circle(fg_mask, (center_x, center_y), radius_annot, 255, -1)

        # 1. rgb -> normal (generate gt surface normals)
        img_mask = cv2.bitwise_and(img_color_np, img_color_np, mask=fg_mask)
        img_normal_np = generate_sphere_gt_normals(
            img_mask, center_x, center_y, radius=radius_bearing
        )

        # 2. downsample and convert to NumPy: (320,240,3) -> (160,120,3)
        img_normal_np = data_utils.interpolate_img(
            img=torch.tensor(img_normal_np).permute(2, 0, 1), rows=160, cols=120)
        img_normal_np = img_normal_np.permute(1, 2, 0).cpu().detach().numpy()
        img_color_ds = data_utils.interpolate_img(
            img=torch.tensor(img_color_np).permute(2, 0, 1), rows=160, cols=120)
        img_color_np = img_color_ds.permute(1, 2, 0).cpu().detach().numpy()

        # 3. save csv files for color and normal images

        if cfg.dataset.save_dataset:
            imageio.imwrite(
                f"{dirs[0]}/{img_idx:04d}.png", (img_color_np * 255).astype(np.uint8)
            )
            imageio.imwrite(f"{dirs[1]}/{img_idx:04d}.png", (img_normal_np*255).astype(np.uint8))
            print(f"Saved image {img_idx:04d}")
        img_idx += 1

    # post-process CSV files and create train/test split
    create_pixel_csv( img_dir=dirs[0], save_dir=dirs[2], img_type="color")
    create_pixel_csv(img_dir=dirs[1], save_dir=dirs[3], img_type="normal")
    combine_csv(dirs[2],img_type="color")
    combine_csv(dirs[3],img_type="normal")
    create_train_test_csv(color_path=f'{dirs[2]}/combined.csv',
                          normal_path=f'{dirs[3]}/combined.csv',
                          save_dir=dirs[4])


if __name__ == "__main__":
    main()