import cv2
import numpy as np
import os

# Read shield with alpha channel
def readAlpha(path):
    src = cv2.imread(path, cv2.IMREAD_UNCHANGED)
    src = cv2.cvtColor(src, cv2.COLOR_BGRA2RGBA)
    # b, g, r, alpha = cv2.split(src)
    # dst1 = cv2.merge((b, g, r))
    # alpha = cv2.cvtColor(alpha, cv2.COLOR_GRAY2BGR)
    # dst1_64 = dst1.astype(np.float64)
    # alpha_64 = alpha.astype(np.float64)
    # alpha_64 *= 1./255.
    # dst2_64 = np.multiply(dst1_64, alpha_64)
    # dst2 = dst2_64.astype(np.uint8)
    return src

# Read shield with alpha channel and points
def readShield(path, step_board=0):
    pict = readAlpha(path)
    h_pict = pict.shape[0]
    w_pict = pict.shape[1]
    pts_pict = np.array(
        [[step_board, step_board], [w_pict - step_board, step_board], [w_pict - step_board, h_pict - step_board],
         [step_board, h_pict - step_board]])
    return pict, pts_pict

def ReadShields(shield_album):
    picts_paths = []
    f = open(shield_album)
    for g in f:
        if g[-1] == '\n':
            picts_paths.append(g[:-1])
        else:
            picts_paths.append(g)
    f.close()
    # print(picts_paths)
    picts = []
    pts_picts = []
    for shield_path in picts_paths:
        pict, pts_pict = readShield(shield_path)
        picts.append(pict)
        pts_picts.append(pts_pict)
    return picts, pts_picts

def ReadShieldsAni(shield_album):
    picts_paths = []
    # filelist = os.listdir(shield_album)
    # for g in filelist:
    #     picts_paths.append(shield_album + "/" + g)
    # # print(picts_paths)
    picts = []
    pts_picts = []
    for i in range(0,55):
        str_i = str(i).zfill(4)
        pict, pts_pict = readShield(shield_album + "/frame_"+str_i+".png")
        picts.append(pict)
        pts_picts.append(pts_pict)
    return picts, pts_picts

# Overaling transformed shield on frame
def overlaying(pict, img, pts_pict, pts_frame, part=1.0):
    # Homography
    # h, status = cv2.findHomography(pts_pict, pts_frame)

    
    
#     x1 = int(pts_frame[0][1])
#     if x1 < 0:
#         x1 = 0
#     x2 = int(pts_frame[2][1])
#     y1 = int(pts_frame[0][0])
#     if y1 < 0:
#         y1 = 0
#     y2 = int(pts_frame[2][0])

    (y1, x1), (y2, x2) = pts_frame.min(axis=0), pts_frame.max(axis=0)
    if (x2 - x1) > 50 and (y2 - y1) > 50:
        try:
            area = img[x1:x2, y1:y2]
            # rows, cols, channels = area.shape

            b, g, r, alpha = cv2.split(pict)
            pict_rgb = cv2.merge((r, g, b))

            h, status = cv2.findHomography(pts_pict, pts_frame)
            imgWarp = cv2.warpPerspective(pict_rgb, h, (img.shape[1], img.shape[0]))
            alphaWarp = cv2.warpPerspective(alpha, h, (img.shape[1], img.shape[0]))
            pict_area = imgWarp[x1:x2, y1:y2]
            alpha_area = alphaWarp[x1:x2, y1:y2]

            # pict = cv2.resize(pict, (cols, rows))

            znam = 1./255.
            alpha_area = cv2.cvtColor(alpha_area, cv2.COLOR_GRAY2BGR)
            pict_32 = pict_area.astype(np.float32)
            area_32 = area.astype(np.float32)
            alpha_32 = alpha_area.astype(np.float32)
            alpha_32 *= znam
            if part != 1.:
                alpha_32 *= part
            alpha_nega_32 = 1. - alpha_32
            dst1_32 = np.multiply(pict_32, alpha_32)
            dst2_32 = np.multiply(area_32, alpha_nega_32)
            dst_32 = dst1_32 + dst2_32
            img[x1:x2, y1:y2] = dst_32.astype(np.uint8)

            # Warp transform shield
            # imgWarp = cv2.warpPerspective(pict_rgb, h, (img.shape[1], img.shape[0]))
            # alphaWarp = cv2.warpPerspective(alpha, h, (img.shape[1], img.shape[0]))

            # Placing sheild on frame
            # imgAug = cv2.addWeighted(img, 1.0, imgWarp, 0.5, 0.0)
        except:
            pass

    return img

def applistShields(album_paths):
    picts_paths = []
    picts_names = []
    f = open(album_paths)
    for g in f:
        if g[-1] == '\n':
            picts_paths.append(g[:-1])
        else:
            picts_paths.append(g)
    f.close()
    for path in picts_paths:
        name = path.split('.')[0].split("/")[1]
        picts_names.append(name)
    return picts_paths, picts_names