import gradio as gr
from core import Ladeco
from matplotlib.figure import Figure
import matplotlib.pyplot as plt
import matplotlib as mpl
import spaces
from PIL import Image


mpl.style.use("dark_background")
ladeco = Ladeco()


@spaces.GPU
def infer(img: str) -> tuple[Figure, Figure]:
    out = ladeco.predict(img)

    seg = out.visualize(level=2)[0].image
    colormap = out.color_map(level=2)

    area = out.area()[0]

    # match the color of segmentation image and pie chart
    colors = []
    l2_area = {}
    for labelname, area_ratio in area.items():
        if labelname.startswith("l2") and area_ratio > 0:
            colors.append(colormap[labelname])
            labelname = labelname.replace("l2_", "").capitalize()
            l2_area[labelname] = area_ratio

    pie = plot_pie(l2_area, colors=colors)

    return seg, pie


def plot_pie(data: dict[str, float], colors=None) -> Figure:
    fig, ax = plt.subplots()

    labels = list(data.keys())
    sizes = list(data.values())

    *_, autotexts = ax.pie(sizes, labels=labels, autopct="%1.1f%%", colors=colors)

    for percent_text in autotexts:
        percent_text.set_color("k")

    ax.axis("equal")

    return fig


def choose_example(imgpath: str) -> gr.Image:
    img = Image.open(imgpath)
    width, height = img.size
    ratio = 512 / max(width, height)
    img = img.resize((int(width * ratio), int(height * ratio)))
    return gr.Image(value=img, label="輸入影像（不支援 SVG 格式）", type="filepath")


css = """
.reference {
    text-align: center;
    font-size: 1.2em;
    color: #d1d5db;
    margin-bottom: 20px;
}
.reference a {
    color: #FB923C;
    text-decoration: none;
}
.reference a:hover {
    text-decoration: underline;
    color: #FB923C;
}
.description {
    text-align: center;
    font-size: 1.1em;
    color: #d1d5db;
    margin-bottom: 25px;
}
.footer {
    text-align: center;
    margin-top: 30px;
    padding-top: 20px;
    border-top: 1px solid #ddd;
    color: #d1d5db;
    font-size: 14px;
}
.main-title {
    font-size: 24px;
    font-weight: bold;
    text-align: center;
    margin-bottom: 20px;
}
.selected-image {
    height: 756px;
}
.example-image {
    height: 220px;
}
""".strip()
theme = gr.themes.Base(
    primary_hue="orange",
    secondary_hue="orange",
    neutral_hue="gray",
    font=gr.themes.GoogleFont("Source Sans Pro"),
).set(
    background_fill_primary="*neutral_950",  # 主背景色（深黑）
    button_primary_background_fill="*primary_500",  # 按鈕顏色（橘色）
    body_text_color="*neutral_200",  # 文字顏色（淺色）
)
with gr.Blocks(css=css, theme=theme) as demo:
    gr.HTML(
        """
        <div class="main-title">LaDeco 景觀環境影像語意分析模型</div>
        <div class="reference">
          引用資料：
          <a href="https://www.sciencedirect.com/science/article/pii/S1574954123003187" target="_blank">
            Li-Chih Ho (2023), LaDeco: A Tool to Analyze Visual Landscape Elements, Ecological Informatics, vol. 78.
          </a>
        </div>
    """.strip()
    )
    with gr.Row(equal_height=True):
        with gr.Group():
            img = gr.Image(
                label="輸入影像（不支援 SVG 格式）",
                type="filepath",
                elem_classes="selected-image",
            )
            gr.Label("範例影像", show_label=False)
            with gr.Row():
                ex1 = gr.Image(
                    value="examples/beach.jpg",
                    show_label=False,
                    type="filepath",
                    elem_classes="example-image",
                    interactive=False,
                    show_download_button=False,
                    show_fullscreen_button=False,
                )
                ex2 = gr.Image(
                    value="examples/field.jpg",
                    show_label=False,
                    type="filepath",
                    elem_classes="example-image",
                    interactive=False,
                    show_download_button=False,
                    show_fullscreen_button=False,
                )
                ex3 = gr.Image(
                    value="examples/sky.jpg",
                    show_label=False,
                    type="filepath",
                    elem_classes="example-image",
                    interactive=False,
                    show_download_button=False,
                    show_fullscreen_button=False,
                )

        with gr.Column():
            seg = gr.Plot(label="語意分割")
            pie = gr.Plot(label="元素面積比例")

    start = gr.Button("開始分析", variant="primary")

    gr.HTML(
        """
        <div class="footer">
            © 2024 LaDeco 版權所有<br>
            開發者：何立智、楊哲睿
        </div>
    """.strip()
    )

    start.click(fn=infer, inputs=img, outputs=[seg, pie])

    ex1.select(fn=choose_example, inputs=ex1, outputs=img)
    ex2.select(fn=choose_example, inputs=ex2, outputs=img)
    ex3.select(fn=choose_example, inputs=ex3, outputs=img)


if __name__ == "__main__":
    demo.launch()