import cv2
import sys
import os
import numpy as np
import matplotlib.pyplot as plt
import glob
from pathlib import Path
import shutil

import argparse

argp = argparse.ArgumentParser("blood_mapgen")

argp.add_argument("in_file")
argp.add_argument("out_dir")
args = argp.parse_args()

input_file = Path(args.in_file)
if not input_file.exists() or input_file.is_dir():
    print(f"Input file {args.in_file} doesn't exist, or is not a file")
    sys.exit(1)

output_dir = Path(args.out_dir)
if output_dir.exists() and not output_dir.is_dir():
    print(f"Out directory {args.out_dir} is not a directory")
    sys.exit(1)

# Ensure output dir exists
output_dir.mkdir(parents=True, exist_ok=True)

def process_file(p):
    dist_out_name = (output_dir / Path(p).name).with_name(Path(p).stem + "_dist.png")

    # Load image with the blood normals
    img_normal = cv2.imread(p,cv2.IMREAD_UNCHANGED)

    (b,g,r, a) = cv2.split(img_normal)

    # Threshold the image to convert to a binary image 
    ret, thresh = cv2.threshold(r, 5,  
                                255, cv2.THRESH_BINARY)
    # Perform the distance transform
    dist = cv2.distanceTransform(thresh,  
                                cv2.DIST_L2, 0) 

    # Normalize the distance transform 
    normalized = cv2.normalize(dist, None, 0,  
                            255, cv2.NORM_MINMAX, cv2.CV_8U)

    cv2.imwrite(dist_out_name.__fspath__(), normalized)

    # Generate a normal map from the distance map
    zx = cv2.Sobel(dist, cv2.CV_64F, 1, 0, ksize=7, scale=1.0)     
    zy = cv2.Sobel(dist, cv2.CV_64F, 0, 1, ksize=7, scale=1.0)

    zxyn = cv2.merge([zx, zy])
    zxyn = cv2.normalize(zxyn, None, 0,  
                            255, cv2.NORM_MINMAX, cv2.CV_8U)
    (r, g) = cv2.split(zxyn)
    zeros = np.zeros_like(r)

    magic = cv2.merge([zeros, g, r])
    normal_border_out_name = (output_dir / Path(p).name).with_name(Path(p).stem + "_normal_border.png")
    
    cv2.imwrite(normal_border_out_name.__fspath__(), magic)

process_file(input_file.__fspath__())