import cv2
import sys
import numpy as np
import re

if len(sys.argv) < 2:
    print("Please provide the input files.")
    sys.exit(1)

basecolor_file = None
normal_file = None
rma_file = None

# Loop through the command-line arguments
for arg in sys.argv[1:]:
    if '_alb.' in arg.lower():
        basecolor_file = arg
    elif '_nrm.' in arg.lower():
        normal_file = arg
    elif '_rma.' in arg.lower():
        rma_file = arg

# Check if all files are provided
if not all([basecolor_file, normal_file, rma_file]):
    print("Please provide all input files: base_color, normal, roughness, ambient_occlusion, and metallic.")
    sys.exit(1)

basecolor = cv2.imread(basecolor_file, cv2.IMREAD_COLOR)
normal = cv2.imread(normal_file, cv2.IMREAD_COLOR)
orm = cv2.imread(rma_file, cv2.IMREAD_COLOR)

basecolorb, basecolorg, basecolorr = basecolor.T
normalb, normalg, normalr = normal.T
aob, metallicb, roughnessb, = orm.T

BCR = np.array([basecolorb, basecolorg, basecolorr, roughnessb])
NMO = np.array([metallicb, normalg, normalr, aob])

result = re.search(r'TX_(.*)_', basecolor_file)
output_dir = re.search(r'^(.+)\\.+$', basecolor_file).group(1)

cv2.imwrite(output_dir + f"\\{result.group(1)}_BCR.png", BCR.T)
cv2.imwrite(output_dir + f"\\{result.group(1)}_NMO.png", NMO.T)