import csv
import datetime

import cv2
import numpy as np

checker_size = 150
margin = 30


# 舌カラーチャートの生成
def make_tcc(no_str=""):
    # Load Lab values from the CSV file
    with open(f"TCC_Lab{no_str}.csv", "r") as csvfile:
        reader = csv.reader(csvfile)
        data = []
        for row in reader:
            data.append([float(x) for x in row])
    # print("Lab values")
    # print(np.array(data))

    cols = 4
    rows = 6
    image_w = checker_size * cols + margin * (cols + 1)
    image_h = checker_size * rows + margin * (rows + 1)
    print("Image size: %dx%d" % (image_w, image_h))

    lab_image = np.zeros((image_h, image_w, 3), dtype=np.float32)
    for c in range(cols):
        for r in range(rows):
            x = image_w - (c + 1) * (checker_size + margin)
            y = margin + r * (checker_size + margin)
            idx = c * rows + r
            lab_image[y : y + checker_size, x : x + checker_size, 0] = data[idx][0]
            lab_image[y : y + checker_size, x : x + checker_size, 1] = data[idx][1]
            lab_image[y : y + checker_size, x : x + checker_size, 2] = data[idx][2]

    rgb_image = (cv2.cvtColor(lab_image, cv2.COLOR_LAB2BGR) * 255).astype(np.uint8)
    # print(rgb_image.shape)
    # print(rgb_image.dtype)
    # print(np.max(rgb_image[:,:,0]))
    # print(np.max(rgb_image[:,:,1]))
    # print(np.max(rgb_image[:,:,2]))

    output_image(rgb_image.copy(), 4, 6, len(data), f"tcc{no_str}.tif", True)

    # ここからTIASカラーチャート用画像の生成

    full_image = np.full((1545, 810, 3), dtype=np.uint8, fill_value=[255, 255, 255])
    full_image[210 : 210 + image_h, 30 : 30 + image_w, :] = rgb_image

    aruco = cv2.aruco
    dictionary = aruco.getPredefinedDictionary(aruco.DICT_4X4_50)
    for i in range(2):
        marker_id = 40 + i
        marker = aruco.generateImageMarker(dictionary, marker_id, 150)
        x = 330 + (i // 2) * 200
        y = 30 + (i % 2) * 1335
        full_image[y : y + 150, x : x + 150, 0] = marker
        full_image[y : y + 150, x : x + 150, 1] = marker
        full_image[y : y + 150, x : x + 150, 2] = marker
        x = 50 + (i // 2) * 700
        y = 60 + (i % 2) * 1420
        cv2.putText(
            full_image,
            "id:%d" % marker_id,
            (x, y),
            cv2.FONT_HERSHEY_TRIPLEX,
            1,
            (0, 0, 0),
            1,
            cv2.LINE_AA,
        )

    cut_color = (200, 200, 200)
    points = [
        (300, 0),
        (300, 180),
        (0, 180),
        (0, 1350),
        (330, 1544),
        (480, 1544),
        (809, 1350),
        (809, 180),
        (510, 180),
        (510, 0),
        (330, 0),
    ]
    for i in range(len(points)):
        cv2.line(full_image, points[i], points[(i + 1) % len(points)], cut_color, 2)

    date_str = datetime.datetime.now().strftime("%Y.%m.%d")
    cv2.putText(
        full_image,
        f"Tongue Color Checker, Chiba University  {date_str}",
        (45, 1345),
        cv2.FONT_HERSHEY_TRIPLEX,
        0.8,
        (0, 0, 0),
        1,
        cv2.LINE_AA,
    )

    cv2.putText(
        full_image,
        "54x103 mm",
        (600, 1520),
        cv2.FONT_HERSHEY_TRIPLEX,
        0.9,
        (0, 0, 0),
        1,
        cv2.LINE_AA,
    )

    cv2.imwrite(f"tcc{no_str}full.tif", full_image)


# 指定RGB値とその近似色の色票を生成
def make_test(bc, gc, rc, step):
    rgb = []
    for bb in range(bc - step, bc + step + 1, step):
        for gg in range(gc - step, gc + step + 1, step):
            for rr in range(rc - step, rc + step + 1, step):
                rgb.append([bb, gg, rr])
    # print(np.array(rgb))

    cols = 4
    rows = 7
    image_w = checker_size * cols + margin * (cols + 1)
    image_h = checker_size * rows + margin * (rows + 1)
    print("Image size: %dx%d" % (image_w, image_h))

    image = np.zeros((image_h, image_w, 3), dtype=np.uint8)
    for c in range(cols):
        for r in range(rows):
            x = image_w - (c + 1) * (checker_size + margin)
            y = margin + r * (checker_size + margin)
            idx = c * rows + r
            if idx < len(rgb):
                image[y : y + checker_size, x : x + checker_size, 0] = rgb[idx][0]
                image[y : y + checker_size, x : x + checker_size, 1] = rgb[idx][1]
                image[y : y + checker_size, x : x + checker_size, 2] = rgb[idx][2]

    output_image(
        image, cols, rows, len(rgb), "test%03d_%03d_%03d_%d.png" % (bc, gc, rc, step)
    )


# 色票画像を出力
def output_image(rgb_image, cols, rows, num, filename, output_values=False):
    rgb = np.zeros((num, 3), dtype=np.uint8)
    for c in range(cols):
        for r in range(rows):
            idx = c * rows + r
            if idx < num:
                x = rgb_image.shape[1] - (c + 1) * (checker_size + margin)
                y = margin + r * (checker_size + margin)
                val = rgb_image[y + 1, x + 1, :]
                rgb[idx, :] = val
                if output_values:
                    cv2.putText(
                        rgb_image,
                        "b%3d, g%3d, r%3d" % (val[0], val[1], val[2]),
                        (x, y - 5),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        0.5,
                        (255, 255, 255),
                        1,
                        cv2.LINE_AA,
                    )

    print("BGR values")
    print(rgb)

    # cv2.imshow("Image", rgb_image)
    # cv2.waitKey(0)
    cv2.imwrite(filename, rgb_image)


# メイン
if __name__ == "__main__":
    make_tcc("6")  # チャートを作るLab値ファイル TCC_Lab#.csv #の数値を指定
    # make_test(92, 110, 181, 5)