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])
    ssize = (600, 888)
    rgb_image_s = cv2.resize(rgb_image, ssize)
    full_image[274 : 274 + ssize[1], 105 : 105 + ssize[0], :] = rgb_image_s

    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, 120)
        x = 345 + (i // 2) * 200
        y = 130 + (i % 2) * 1070
        full_image[y : y + 120, x : x + 120, 0] = marker
        full_image[y : y + 120, x : x + 120, 1] = marker
        full_image[y : y + 120, x : x + 120, 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,
        "Tongue Color Checker",
        (210, 1380),
        cv2.FONT_HERSHEY_TRIPLEX,
        1.0,
        (0, 0, 0),
        1,
        cv2.LINE_AA,
    )
    cv2.putText(
        full_image,
        "Chiba University",
        (250, 1420),
        cv2.FONT_HERSHEY_TRIPLEX,
        1.0,
        (0, 0, 0),
        1,
        cv2.LINE_AA,
    )
    cv2.putText(
        full_image,
        date_str,
        (300, 1460),
        cv2.FONT_HERSHEY_TRIPLEX,
        1.0,
        (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)
    np.savetxt(f"TCC_RGB{filename[3:-4]}.csv", rgb, fmt="%d", delimiter=",")

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


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