import os.path as osp
from glob import glob

import cv2
import numpy as np
from tqdm import tqdm


def make_dataset_list():
    imgs_path_list = glob("./*.png")

    dataset_list = []
    for path in imgs_path_list:
        img_num = int(osp.basename(path)[:-4])
        target_path = osp.join(osp.dirname(path), "{:08}.png".format(img_num + 1))
        if osp.exists(target_path):
            dataset_list.append({"seq_0": path, "seq_1": target_path})

    return dataset_list


dataset_list = make_dataset_list()
print(f"samples: {len(dataset_list)}")

mae_list = []
for pair in tqdm(dataset_list):
    seq_0 = cv2.imread(pair["seq_0"])
    seq_1 = cv2.imread(pair["seq_1"])

    img_num = int(osp.basename(pair["seq_0"])[:-4])
    cv2.imwrite(f"./diff/diff_{img_num}_src.jpg", np.clip(np.abs(seq_0 - seq_1), 0, 255).astype(np.uint8))

    cv2.imshow("diff1", np.clip(np.abs(seq_0 - seq_1), 0, 255).astype(np.uint8))

    # medianフィルタ
    # seq_0 = cv2.medianBlur(seq_0, 3)
    # seq_1 = cv2.medianBlur(seq_1, 3)

    # Gaussianフィルタ
    # seq_0 = cv2.GaussianBlur(seq_0, (3, 3), 0)
    # seq_1 = cv2.GaussianBlur(seq_1, (3, 3), 0)

    # 移動平均フィルタ
    # seq_0 = cv2.blur(seq_0, (3, 3))
    # seq_1 = cv2.blur(seq_1, (3, 3))

    # バイラテラルフィルタ
    # seq_0 = cv2.bilateralFilter(seq_0, 3, 20, 20)
    # seq_1 = cv2.bilateralFilter(seq_1, 3, 20, 20)

    # ノンローカルミーンフィルタ
    seq_0 = cv2.fastNlMeansDenoisingColored(seq_0, None, 10, 10, 7, 21)
    seq_1 = cv2.fastNlMeansDenoisingColored(seq_1, None, 10, 10, 7, 21)

    cv2.imshow("seq_0", seq_0)
    cv2.imshow("seq_1", seq_1)

    cv2.imwrite(f"./diff/diff_{img_num}_denoise.jpg", np.clip(np.abs(seq_0 - seq_1), 0, 255).astype(np.uint8))
    cv2.imshow("diff2", np.clip(np.abs(seq_0 - seq_1), 0, 255).astype(np.uint8))

    mae_list.append(np.mean(np.abs(seq_0 - seq_1)))
    cv2.waitKey(3000)

print(f"ave_mae: {sum(mae_list) / len(mae_list)}")