import gradio as gr
import numpy as np
from PIL import Image

theme_id = "HokageM/EVA_01"
theme = gr.themes.ThemeClass.from_hub(theme_id)

HEADER = """
This is a color vision correction tool.

The `Deuteranomaly` and `Protanomaly` sliders control the strength of the correction. In this context, correction means that the image colors are adjusted to make color differences more distinguishable for people with color vision deficiencies.
"""

DESCRIPTION_FORMULA = r"""
I have a mix of deuteranomaly and protanomaly, which motivated me to create this tool. With the correction enabled, colorblind tests become much more solvable for me.

The correction formula comes from this [paper](https://arxiv.org/abs/1711.10662) by Jinmi Lee:

$$
\begin{bmatrix}
R' \\
G' \\
B'
\end{bmatrix}
=
\begin{bmatrix}
1 - \frac{\alpha_d}{2} & \frac{\alpha_d}{2} & 0 \\
\frac{\alpha_p}{2} & 1 - \frac{\alpha_p}{2} & 0 \\
\frac{\alpha_p}{4} & \frac{\alpha_d}{4} & 1 - \frac{\alpha_p + \alpha_d}{4}
\end{bmatrix}
\begin{bmatrix}
R \\
G \\
B
\end{bmatrix}
$$

You can set $\alpha_d$ by adjusting the `Deuteranomaly` slider and $\alpha_p$ by adjusting the `Protanomaly` slider.
"""


def correction_for_colorblindness(image_array: np.array, degree_protanomaly: float, degree_deuteranomaly: float):
    """Apply a simple protanomaly/deuteranomaly correction matrix.

    The input image is converted to float during calculation to avoid uint8
    overflow/underflow. The result is clipped back into the valid RGB range.
    """

    if degree_protanomaly < 0 or degree_deuteranomaly < 0:
        raise ValueError("Colorblindness correction degrees must be >= 0.")

    r = image_array[..., 0].astype(float)
    g = image_array[..., 1].astype(float)
    b = image_array[..., 2].astype(float)

    corrected = image_array.astype(float).copy()

    corrected[..., 0] = (1 - degree_deuteranomaly / 2) * r + (
            degree_deuteranomaly / 2
    ) * g
    corrected[..., 1] = (degree_protanomaly / 2) * r + (1 - degree_protanomaly / 2) * g
    corrected[..., 2] = (
            (degree_protanomaly / 4) * r
            + (degree_deuteranomaly / 4) * g
            + (1 - (degree_deuteranomaly + degree_protanomaly) / 4) * b
    )

    return np.clip(corrected, 0, 255).astype(np.uint8)


def load_demo():
    return Image.open("./color_blind_test.jpg")


def unload():
    return None, None


with gr.Blocks() as demo:
    gr.Markdown("# Byakugan Visualizer")
    gr.Markdown(HEADER)
    with gr.Row():
        with gr.Column():
            deuter_degree = gr.Slider(step=0.5, minimum=0, maximum=5, label=r"Deuteranomaly")
            prot_degree = gr.Slider(step=0.5, minimum=0, maximum=5, label="Protanomaly")
            correction_btn = gr.Button("Create Correction", variant="primary")
        gr.Markdown(DESCRIPTION_FORMULA, latex_delimiters=[{"left": "$$", "right": "$$", "display": True},
                                                           {"left": "$", "right": "$", "display": False}])
    with gr.Row():
        img = gr.Image()
        processed_img = gr.Image()

    with gr.Row():
        load_demo_btn = gr.Button("Load Colorblind Test", variant="primary")
        unload_btn = gr.Button("Unload Images", variant="secondary")

    deuter_degree.change(
        fn=correction_for_colorblindness,
        inputs=[img, deuter_degree, prot_degree],
        outputs=[processed_img],
    )
    prot_degree.change(
        fn=correction_for_colorblindness,
        inputs=[img, deuter_degree, prot_degree],
        outputs=[processed_img],
    )
    correction_btn.click(
        fn=correction_for_colorblindness,
        inputs=[img, deuter_degree, prot_degree],
        outputs=[processed_img],
    )
    load_demo_btn.click(
        fn=load_demo,
        outputs=[img],
    )
    unload_btn.click(
        fn=unload,
        outputs=[img, processed_img],
    )

demo.launch(theme=theme)