import gradio as gr
from PIL import Image

elements = {
    "Becane": ("data/becane_1_0_2_014/before/img6", "data/becane_1_0_2_014/after/img7", "models/RoadBecane_f.stl"),
    "BTWIN": ("data/btwin_1_26_3_114/before/img0", "data/btwin_1_26_3_114/after/img13", "models/BTWIN.stl"),
    "Classic Road": ("data/croad_4_26_4_112/before/img3", "data/croad_4_26_4_112/after/img7", "models/ClassicRoad.stl"),
    "Domane": ("data/domane_2_0_39_012/before/img6", "data/domane_2_0_39_012/after/img13", "models/domane.stl"),
    "Enduro": ("data/enduro_2_18_18_105/before/img6", "data/enduro_2_18_18_105/after/img10", "models/enduro.stl"),
    "G1": ("data/g1_2_36_14_010/before/img6", "data/g1_2_36_14_010/after/img8", "models/g1.stl"),
    "GBike": ("data/gbike_3_31_8_018/before/img2", "data/gbike_3_31_8_018/after/img8", "models/gbike.stl"),
    "Holland": (
        "data/holland_4_24_18_00/before/img4", "data/holland_4_24_18_00/after/img7", "models/holland.stl"),
    "Huffy": (
        "data/huffy_6_4_35_012/before/img4", "data/huffy_6_4_35_012/after/img12", "models/huffy.stl"),
    "Kuota": (
        "data/kuota_6_5_30_112/before/img4", "data/kuota_6_5_30_112/after/img12", "models/kuota.stl"),
    "MFactory": (
        "data/mfactory_1_17_8_00/before/img4", "data/mfactory_1_17_8_00/after/img12", "models/mfactory.stl"),
    "Mirage": (
        "data/mirage_6_11_36_117/before/img4", "data/mirage_6_11_36_117/after/img12", "models/mirage.stl"),
    "Old Bike": (
        "data/oldbike_3_30_15_100/before/img4", "data/oldbike_3_30_15_100/after/img12", "models/oldbike.stl"),
    "Freeride": (
        "data/freeride_6_24_0_00/before/img4", "data/freeride_6_24_0_00/after/img12", "models/freeride_1Step_f.stl"),
    "RondoRuut": (
        "data/rondo_2_11_39_106/before/img4", "data/rondo_2_11_39_106/after/img12",
        "models/RondoRuutClean_1Step_f.stl"),
    "MTB Ghost": (
        "data/ghost_1_0_4_116/before/img4", "data/ghost_1_0_4_116/after/img12", "models/MTB_Ghost.stl"),
    "Scalpel": (
        "data/scalpel_6_2_12_101/before/img4", "data/scalpel_6_2_12_101/after/img12", "models/scalpel.stl"),
    "Verdona": (
        "data/verdona_1_14_10_118/before/img4", "data/verdona_1_14_10_118/after/img12", "models/verdona.stl"),
    "Vintage": (
        "data/vintage_6_6_33_011/before/img4", "data/vintage_6_6_33_011/after/img12", "models/Vintage.stl"),
    "WBike": (
        "data/wbike_3_19_23_00/before/img4", "data/wbike_3_19_23_00/after/img12", "models/wbike.stl"),
}


def generate_ui(key):
    img_b, img_a, model = elements[key]
    gr.Markdown("### Before damage")
    with gr.Row():
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("#### Render")
                img_bn = gr.Image(f"{img_b}.png")
            with gr.Column(scale=1):
                gr.Markdown("#### Background")
                img_bb = gr.Image(f"{img_b}_background.png")
            with gr.Column(scale=1):
                gr.Markdown("#### Foreground")
                img_bf = gr.Image(f"{img_b}_foreground.png")
            with gr.Column(scale=1):
                gr.Markdown("#### Segmentation")
                img_bs = gr.Image(f"{img_b}_segmentation.png")

    gr.Markdown("### After damage")
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("#### Render")
            img_an = gr.Image(f"{img_a}.png")
        with gr.Column(scale=1):
            gr.Markdown("#### Background")
            img_ab = gr.Image(f"{img_a}_background.png")
        with gr.Column(scale=1):
            gr.Markdown("#### Foreground")
            img_af = gr.Image(f"{img_a}_foreground.png")
        with gr.Column(scale=1):
            gr.Markdown("#### Segmentation")
            img_as = gr.Image(f"{img_a}_segmentation.png")

    model = gr.Model3D(model, label="3D model preview")
    return img_bn, img_bb, img_bf, img_bs, img_an, img_ab, img_af, img_as, model


def get_values(key):
    img_b, img_a, model = elements[key]
    gr.Markdown("### Before damage")
    img_bn = Image.open(f"{img_b}.png")
    img_bb = Image.open(f"{img_b}_background.png")
    img_bf = Image.open(f"{img_b}_foreground.png")
    img_bs = Image.open(f"{img_b}_segmentation.png")
    img_an = Image.open(f"{img_a}.png")
    img_ab = Image.open(f"{img_a}_background.png")
    img_af = Image.open(f"{img_a}_foreground.png")
    img_as = Image.open(f"{img_a}_segmentation.png")

    return img_bn, img_bb, img_bf, img_bs, img_an, img_ab, img_af, img_as, model


block = gr.Blocks()

with block:
    with open("page.md", "r") as f:
        gr.Markdown(f.read())
        first_key = list(elements.keys())[0]
        dropdown = gr.Dropdown(choices=list(elements.keys()), value=first_key, label="Model", interactive=True)
        img_bn, img_bb, img_bf, img_bs, img_an, img_ab, img_af, img_as, model = generate_ui(first_key)

    dropdown.change(get_values, inputs=[dropdown],
                    outputs=[img_bn, img_bb, img_bf, img_bs, img_an, img_ab, img_af,
                             img_as, model])

block.queue(concurrency_count=40, max_size=20).launch(max_threads=150)