import openpyxl
import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression


def read_excel(filename: str, range_str: str, sheet_no: int = 0) -> np.array:
    """Excelから数値データ取得

    Args:
        filename (str): ファイル名
        range_str (str): セル範囲文字列 'A1:B2'

    Returns:
        np.array: セル値配列
    """
    wb = openpyxl.load_workbook(filename)
    sheet = wb[wb.sheetnames[sheet_no]]
    return np.array(
        [
            [cell.value for cell in sheet[range_str][i]]
            for i in range(len(sheet[range_str]))
        ]
    )


def extend_rgb(rgb_arr: np.array) -> np.array:
    """RGB値の高次項を拡張する

    Args:
        rgb_arr (np.array): RGB値の配列 1x3, Nx3

    Returns:
        np.array: 拡張した配列 1x11, Nx11
    """

    conv_2d = 0
    if rgb_arr.ndim == 1:
        rgb_arr = rgb_arr.reshape([1, len(rgb_arr)])
        conv_2d = 1
    extended = np.array(
        [
            [
                1,
                rgb[0],
                rgb[1],
                rgb[2],
                rgb[0] ** 2,
                rgb[1] ** 2,
                rgb[2] ** 2,
                rgb[0] * rgb[1],
                rgb[1] * rgb[2],
                rgb[2] * rgb[0],
                rgb[0] * rgb[1] * rgb[2],
            ]
            for rgb in rgb_arr
        ]
    )
    if conv_2d == 1:
        return extended[0]
    else:
        return extended


def lab_func(t: float) -> float:
    """XYZ->LAB変換に使用する暗部補正付き関数

    Args:
        t (float): 値

    Returns:
        float: 変換結果
    """
    if t > 0.008856:
        return t ** (1.0 / 3.0)
    else:
        return 7.787 * t + 16.0 / 116.0


def xyz_to_lab(xyz_arr: np.array, white_xyz: np.array) -> np.array:
    """XYZからCIELABへ変換

    Args:
        xyz_arr (np.array): XYZ三刺激値の配列 1x3, Nx3
        white_xyz (np.array): 白色点の三刺激値

    Returns:
        np.array: L*a*b*の配列 1x3, Nx3
    """

    conv_2d = 0
    if xyz_arr.ndim == 1:
        xyz_arr = xyz_arr.reshape([1, len(xyz_arr)])
        conv_2d = 1
    lab_arr = np.array(
        [
            [
                116.0 * lab_func(xyz[1] / white_xyz[1]) - 16.0,
                500.0
                * (lab_func(xyz[0] / white_xyz[0]) - lab_func(xyz[1] / white_xyz[1])),
                200.0
                * (lab_func(xyz[1] / white_xyz[1]) - lab_func(xyz[2] / white_xyz[2])),
            ]
            for xyz in xyz_arr
        ]
    )
    if conv_2d == 1:
        return lab_arr[0]
    else:
        return lab_arr


if __name__ == "__main__":
    # --------------------------------------------------------------
    # Excelからデータ取得
    light_sp = read_excel("Spectrum.xlsx", "B2:AF2")[0]
    cmf_sp_arr = read_excel("Spectrum.xlsx", "B3:AF5")
    filter_sp_arr = read_excel("Spectrum.xlsx", "B6:AF8")
    tongue_sp_arr = read_excel("Spectrum.xlsx", "B9:AF16")

    # tcc_name = "handy_tcc"
    # tcc_name = "takano_tcc"
    tcc_name = "sugawara_tcc_mat8485"
    tcc_sp_arr = read_excel(tcc_name + ".xlsx", "F2:AJ25") / 100.0
    white_patch = (
        0 if tcc_name == "handy_tcc" else 12
    )  # ホワイトバランスに用いる色票番号

    # --------------------------------------------------------------
    # 白色点の計算
    white_xyz = np.array([light_sp.dot(cmf_sp) for cmf_sp in cmf_sp_arr])
    k = 100.0 / white_xyz[1]
    white_xyz *= k
    print("白色点 XYZ=", white_xyz)

    # --------------------------------------------------------------
    # TCCのXYZ, L*a*b* 算出（正解値）
    tcc_xyz_arr = np.array(
        [
            [
                (light_sp * tcc_sp_arr[p]).dot(cmf_sp_arr[i]) * k
                for i in range(len(cmf_sp_arr))
            ]
            for p in range(len(tcc_sp_arr))
        ]
    )
    # print("tcc_xyz=")
    # print(tcc_xyz)
    tcc_lab_arr = xyz_to_lab(tcc_xyz_arr, white_xyz)
    print("TCC L*a*b* 正解値=")
    print(tcc_lab_arr)

    # --------------------------------------------------------------
    # カメラシミュレータのホワイトバランス
    white_balance = np.array([0.9, 0.9, 0.9])  # ホワイトバランスのRGB値
    raw_rgb = np.array(
        [
            (light_sp * tcc_sp_arr[white_patch]).dot(filter_sp_arr[i])
            for i in range(len(filter_sp_arr))
        ]
    )
    # print("raw_rgb=", raw_rgb)
    rgb_gain = white_balance / raw_rgb
    # print("rgb_gain=", rgb_gain)

    # --------------------------------------------------------------
    # カメラシミュレータによる色票のRGB値算出
    tcc_rgb_arr = np.array(
        [
            [
                (light_sp * tcc_sp_arr[p]).dot(filter_sp_arr[i]) * rgb_gain[i]
                for i in range(len(filter_sp_arr))
            ]
            for p in range(len(tcc_sp_arr))
        ]
    )
    # print("TCC RGB=")
    # print(tcc_rgb_arr)
    # print(extend_rgb(tcc_rgb))

    # --------------------------------------------------------------
    # カメラシミュレータRGBからXYZ推定行列の算出
    # tcc_rgb_input_arr = tcc_rgb_arr
    tcc_rgb_input_arr = extend_rgb(tcc_rgb_arr)
    conv_rgb_to_xyz = LinearRegression()
    conv_rgb_to_xyz.fit(tcc_rgb_input_arr, tcc_xyz_arr)

    # --------------------------------------------------------------
    # カメラシミュレータRGB値からXYZ, L*a*b*の推定
    tcc_xyz_pred_arr = conv_rgb_to_xyz.predict(tcc_rgb_input_arr)
    # print("TCC XYZ 推定値=")
    # print(tcc_xyz_pred_arr)
    tcc_lab_pred_arr = xyz_to_lab(tcc_xyz_pred_arr, white_xyz)
    print("TCC L*a*b* 推定値=")
    print(tcc_lab_pred_arr)

    # --------------------------------------------------------------
    # TCC 正解値と推定値の色差計算
    delta_e_arr = np.array(
        [
            np.linalg.norm(tcc_lab_pred_arr[i] - tcc_lab_arr[i], ord=2)
            for i in range(len(tcc_lab_arr))
        ]
    )
    delta_e = np.mean(delta_e_arr)
    print("TCC 推定色差 DeltaE=", delta_e)

    # --------------------------------------------------------------
    # 舌のXYZ, L*a*b* 算出（正解値）
    tongue_xyz_arr = np.array(
        [
            [
                (light_sp * tongue_sp_arr[p]).dot(cmf_sp_arr[i]) * k
                for i in range(len(cmf_sp_arr))
            ]
            for p in range(len(tongue_sp_arr))
        ]
    )
    # print("tongue_xyz=")
    # print(tongue_xyz_arr)
    tongue_lab_arr = xyz_to_lab(tongue_xyz_arr, white_xyz)
    print("舌 L*a*b* 正解値=")
    print(tongue_lab_arr)

    # --------------------------------------------------------------
    # カメラシミュレータによる舌のRGB値算出
    tongue_rgb_arr = np.array(
        [
            [
                (light_sp * tongue_sp_arr[p]).dot(filter_sp_arr[i]) * rgb_gain[i]
                for i in range(len(filter_sp_arr))
            ]
            for p in range(len(tongue_sp_arr))
        ]
    )
    print("舌 RGB=")
    print(tongue_rgb_arr)

    # --------------------------------------------------------------
    # 舌のカメラシミュレータRGB値からXYZ, L*a*b*の推定
    tongue_rgb_input_arr = extend_rgb(tongue_rgb_arr)
    tongue_xyz_pred_arr = conv_rgb_to_xyz.predict(tongue_rgb_input_arr)
    print("舌 XYZ 推定値=")
    print(tongue_xyz_pred_arr)
    tongue_lab_pred_arr = xyz_to_lab(tongue_xyz_pred_arr, white_xyz)
    print("舌 L*a*b* 推定値=")
    print(tongue_lab_pred_arr)

    # --------------------------------------------------------------
    # 舌 正解値と推定値の色差計算
    delta_e_arr = np.array(
        [
            np.linalg.norm(tongue_lab_pred_arr[i] - tongue_lab_arr[i], ord=2)
            for i in range(len(tongue_lab_arr))
        ]
    )
    delta_e = np.mean(delta_e_arr)
    print("舌 推定色差 DeltaE=", delta_e)

    fig = plt.figure(figsize=(15, 8))
    ax_ab = fig.add_subplot(1, 2, 1, xlabel="a*", ylabel="b*")
    plt.grid(True)
    ax_ab.set_title(tcc_name + "  a* - b*")
    ax_ab.set_xlim(-40, 60)
    ax_ab.set_ylim(-40, 60)
    ax_ab.text(-35, 55, "red:GT, black:Estim, green:tcc")
    for i in range(len(tongue_lab_arr)):
        ax_ab.plot(
            [tongue_lab_arr[i, 1], tongue_lab_pred_arr[i, 1]],
            [tongue_lab_arr[i, 2], tongue_lab_pred_arr[i, 2]],
            "-",
            c="gray",
        )
    ax_ab.plot(tongue_lab_arr[:, 1], tongue_lab_arr[:, 2], "o", c="r")
    ax_ab.plot(tongue_lab_pred_arr[:, 1], tongue_lab_pred_arr[:, 2], "o", c="k")
    ax_ab.plot(tcc_lab_arr[:, 1], tcc_lab_arr[:, 2], "o", c="g")

    ax_aL = fig.add_subplot(1, 2, 2, xlabel="a*", ylabel="L*")
    plt.grid(True)
    ax_aL.set_title(tcc_name + "  a* - L*")
    ax_aL.set_xlim(-40, 60)
    ax_aL.set_ylim(0, 100)
    ax_aL.text(-35, 95, "red:GT, black:Estim, green:tcc")
    for i in range(len(tongue_lab_arr)):
        ax_aL.plot(
            [tongue_lab_arr[i, 1], tongue_lab_pred_arr[i, 1]],
            [tongue_lab_arr[i, 0], tongue_lab_pred_arr[i, 0]],
            "-",
            c="gray",
        )
    ax_aL.plot(tongue_lab_arr[:, 1], tongue_lab_arr[:, 0], "o", c="r")
    ax_aL.plot(tongue_lab_pred_arr[:, 1], tongue_lab_pred_arr[:, 0], "o", c="k")
    ax_aL.plot(tcc_lab_arr[:, 1], tcc_lab_arr[:, 0], "o", c="g")

    plt.show()