from typing import List

import matplotlib.colors as mcolors
import matplotlib.pyplot as plt
import numpy as np

# set font by default to courier new
plt.rcParams["font.family"] = "PT Mono"
# # sort by the proximity to the colors in viridis
# # Importing necessary functions

# # Calculate the proximity of each color in XKCD_COLORS to the viridis colormap
# def calculate_proximity_to_viridis(color):
#     rgb_triple = mcolors.to_rgb(mcolors.XKCD_COLORS[color])
#     distances = [(sum((a - b)**2 for a, b in zip(rgb_triple, viridis(i))), i) for i in range(256)]
#     _, closest_viridis_value = min(distances, key=lambda x: x[0])  # Find the viridis color with the minimum distance
#     return closest_viridis_value / 255  # Normalize to range (0, 1)

# # Calculate the proximity values for each color
# proximity_values = {color: calculate_proximity_to_viridis(color) for color in colors}

# # Sort the colors based on their proximity values
# sorted_colors = sorted(proximity_values.keys(), key=lambda x: proximity_values[x])


def create_color_calibration_image(
    colors, ppi: List[int] = [100], index=0, solid_capstyle="butt", antialiased=True
):
    first_color = colors[0]
    last_color = colors[-1]
    print(f"Processing color range: {first_color} to {last_color}")
    # Conversion factor: 1 inch = dpi pixels
    vert_space = 1 / 2  # inches
    fontsize = 12  # points

    # Figure settings
    fig_width = 4.0  # inches
    fig_height = len(colors) * vert_space
    fig, ax = plt.subplots(figsize=(fig_width, fig_height))

    # plot a vertical black rectangle between x=3 and x=4
    ax.axvspan(0.25, 0.75, facecolor="black")
    # ax.axvspan(-0.325-0.175, -0.125, facecolor="black")
    # Loop through each color
    if index % 2 == 0:
        skip = -1
    else:
        skip = 1
    for idx, color in enumerate(colors[::skip]):
        # print(color)
        y_position = 0.5 - 0.125 + idx * vert_space  # Offset each color by 0.3 inches

        # Draw color name
        rgb_triple = mcolors.to_rgb(mcolors.XKCD_COLORS[color])
        # round to 4 decimal places
        rgb_triple = tuple(round(x, 4) for x in rgb_triple)
        # format as string with fixed decimal places
        rgb_triple = ", ".join([f"{x:1.4f}" for x in rgb_triple])
        hex_code = mcolors.to_hex(mcolors.XKCD_COLORS[color])
        ax.text(
            1.0,
            y_position,
            color.replace("xkcd:", ""),
            va="center",
            fontsize=fontsize,
            # bbox=dict(facecolor='gray', alpha=0.21),
        )
        ax.text(
            1.25,
            y_position - 1.5 * fontsize / 72,
            f"{hex_code}\n({rgb_triple})",
            va="center",
            fontsize=6,
            # bbox=dict(facecolor='gray', alpha=0.33),
        )
        # ax.text(
        #     1.125,
        #     y_position - 2 * fontsize / 72,
        #     f"{hex_code}",
        #     va="center",
        #     fontsize=fontsize * 0.6,
        # )

        # Draw color square
        rect_height = 0.25
        rect_width = 0.25
        square_x_start = 0.75 - rect_width / 2  # Offset from the left
        square_y_start = y_position - rect_height  # + 0.25
        ax.add_patch(
            plt.Rectangle(
                (square_x_start, square_y_start), rect_width, rect_height, fc=color
            )
        )

        # Draw lines with varying stroke sizes
        line_x_start = 0  # Offset from the left
        line_length = 0.5
        line_y_start = y_position - rect_height - 0.075
        line_widths = [0.25, 0.5, 0.75, 1, 1.25, 1.5, 1.75, 2.0, 3.0][::-1]
        for width in line_widths:
            ax.plot(
                [line_x_start, line_x_start + line_length],
                [line_y_start, line_y_start],
                linewidth=width,
                color=color,
                antialiased=antialiased,
                solid_capstyle=solid_capstyle,
            )
            line_y_start += 0.05

        # # now repeat but vertical of height 1
        # line_x_start = 3.125 - 0.05  # Offset from the left
        # line_y_start = y_position + dpi / 960
        # for width in line_widths:
        #     ax.plot(
        #         [line_x_start, line_x_start],
        #         [line_y_start, line_y_start - 0.5],
        #         linewidth=width,
        #         color=color,
        #         antialiased=True,
        #     )
        #     ax.plot(
        #         [0.5 + line_x_start, 0.5 + line_x_start],
        #         [line_y_start, line_y_start - 0.5],
        #         linewidth=width,
        #         color=color,
        #         antialiased=True,
        #     )
        #     line_x_start += 0.05

    # Save the image
    # Remove axes
    ax.axis("off")

    # ax.set_aspect("equal")
    # plt.tight_layout(pad=0)
    ax.set_ylim([0, fig_height])
    # ax.set_xlim([0, fig_width])
    ax.set_xlim([0, fig_width])
    # pad = 0.108
    pad = 0
    fig.subplots_adjust(left=0, right=1, top=1, bottom=0)
    output_paths = []
    for _dpi in ppi:
        _out_path = f"/tmp/color_calibration-{_dpi}_{index:02d}.png"
        plt.savefig(_out_path, pad_inches=pad, dpi=_dpi)
        output_paths.append(_out_path)
    # plt.show()
    plt.close()

    return output_paths


if __name__ == "__main__":
    import argparse
    import os

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--rows", type=int, default=73, help="Number of entries per column"
    )
    parser.add_argument(
        "--dir", type=str, default="/Volumes/TMP/tests", help="Directory to save images"
    )
    parser.add_argument(
        "-k",
        "--kind",
        type=str,
        nargs="+",
        default=["hsv", "lex", "lab", "umap"],
        help="Kinds of sorting",
    )
    parser.add_argument(
        "--ppi", action="append", type=int, default=[300], help="Pixels per inch"
    )
    parser.add_argument("--aliased", action="store_true", help="Disable antialiasing")
    parser.add_argument(
        "--capstyle", type=str, default="butt", help="Capstyle of lines"
    )
    args = parser.parse_args()
    COLUMN_LENGTH = args.rows
    KINDS = args.kind
    PPIS = args.ppi
    DIR = args.dir
    ANTIALIASED = not args.aliased
    CAPSTYLE = args.capstyle

    # COLUMN_LENGTH = 73  # results in 13 unfiltered columns (perfect)
    # COLUMN_LENGTH = (
    #     106  # results in 9 unfiltered columns (last one short), square-ish image
    # )
    OMITTED_COLORS = [
        # "black",
        # "white",
        # "poop",
        # "poo brown",
        # "shit",
        # "shit brown",
    ]
    # OMITTED_COLORS = list(map(lambda s: f"xkcd:{s}", OMITTED_COLORS))
    # KIND = "hsv"  # choose from umap, hsv
    for KIND in KINDS:
        colors = list(mcolors.XKCD_COLORS.keys())
        sorted_indices = np.load(f"scripts/{KIND}_sorted_indices.npy")
        sorted_colors = [colors[idx] for idx in sorted_indices]

        colors = sorted_colors
        colors = [c for c in colors if c not in OMITTED_COLORS]
        print(f"Total number of colors: {len(colors)}")

        chunks = [
            colors[i : i + COLUMN_LENGTH] for i in range(0, len(colors), COLUMN_LENGTH)
        ]

        for idx, color_part in enumerate(chunks):
            image_path = create_color_calibration_image(
                colors=color_part,
                ppi=PPIS,
                index=idx,
                antialiased=ANTIALIASED,
                solid_capstyle=CAPSTYLE,
            )
            os.system(f"identify {image_path[0]}")

        for PPI in PPIS:
            # use imagemagick to stitch together the images horizontally
            os.system(
                f"convert +append /tmp/color_calibration-{PPI}_*.png /tmp/color_calibration-{PPI}.png"
            )
            os.system(f"rm /tmp/color_calibration-{PPI}_*")
            print(f"Final image saved to /tmp/color_calibration-{PPI}.png")
            os.system(
                f"mkdir -p {DIR} && cp /tmp/color_calibration-{PPI}.png {DIR}/xkcd_{COLUMN_LENGTH}_{KIND}_{PPI}.png"
            )
            print(f"Copied to {DIR}/xkcd_{COLUMN_LENGTH}_{KIND}_{PPI}.png")