#!/usr/bin/env python3
# -*- coding: UTF-8 -*-
import cv2, os, sys, re

def Stitch_pic(Ori_image_path, Ori_label_path, Result_dir, img_pos, label_pos):
    
    # 读取两张没有stitch_pos通道的图片
    img1 = cv2.imread(Ori_image_path) # 底层图片
    img2 = cv2.imread(Ori_label_path) # 顶层图片
    
    Result_name = os.path.splitext(os.path.basename(Ori_image_path))[0]

    # 将img2调整为与img1大小相同
    if img1.shape != img2.shape:
        img2 = cv2.resize(img2, (img1.shape[1], img1.shape[0]))

    # 拼接图片
    if img_pos == 'up' and label_pos == 'down':
        result = cv2.vconcat([img1, img2])
    elif img_pos == 'down' and label_pos == 'up':
        result = cv2.vconcat([img2, img1])
    elif img_pos == 'left' and label_pos == 'right':
        result = cv2.hconcat([img1, img2])
    elif img_pos == 'right' and label_pos == 'left':
        result = cv2.hconcat([img2, img1])
    else:
        RED = '\033[91m'
        END = '\033[0m'
        print(RED + "The input of relative_pos is wrong, img_pos is " + img_pos + " label_pos is " + label_pos + END)
        os.exit()

    # 保存结果
    if not os.path.exists(Result_dir):
        os.makedirs(Result_dir)
    cv2.imwrite(os.path.join(Result_dir, Result_name+'.png'), result)
    print("堆叠图片写入地址：", os.path.join(Result_dir, Result_name+'.png'))


if __name__ == '__main__':
    Ori_image_path = sys.argv[1] # 背景所在路径
    Ori_label_path = sys.argv[2] # 上层图片所在路径
    Result_dir = sys.argv[3] # 结果所在目录
    relative_pos = str.lower(sys.argv[4])

    match_up = re.search(r'up', relative_pos)
    match_down = re.search(r'down', relative_pos)
    match_left = re.search(r'up', relative_pos)
    match_right = re.search(r'down', relative_pos)

    if match_up and match_down:
        pos_up = match_up.start()
        pos_down = match_down.start()
        if pos_down < pos_up :
            img_pos = "down"
            label_pos = "up"
        else:
            img_pos = "up"
            label_pos = "down"
    elif match_left and match_right:
        pos_left = match_up.start()
        pos_right = match_down.start()
        if pos_left < pos_right :
            img_pos = "left"
            label_pos = "right"
        else:
            img_pos = "left"
            label_pos = "right"
    else:
        print("Either 'up'/'down' 'left'/'right' is missing or in the text.")
        print("Set to default img_pos = up label_pos = down")
        img_pos = "up"
        label_pos = "down"

    # 进行对叠程序
    Stitch_pic(Ori_image_path, Ori_label_path, Result_dir, img_pos, label_pos)