import os

import gradio as gr
from PIL import Image

elements = {
    "Becane": ("data/data_becane_1_0_2_014_before_img6", "data/data_becane_1_0_2_014_after_img7", "models/RoadBecane_f.stl"),
    "BTWIN": ("data/data_btwin_1_26_3_114_before_img0", "data/data_btwin_1_26_3_114_after_img13", "models/BTWIN.stl"),
    "Classic Road": ("data/data_croad_4_26_4_112_before_img3", "data/data_croad_4_26_4_112_after_img7", "models/ClassicRoad.stl"),
    "Domane": ("data/data_domane_2_0_39_012_before_img6", "data/data_domane_2_0_39_012_after_img13", "models/domane.stl"),
    "Enduro": ("data/data_enduro_2_18_18_105_before_img6", "data/data_enduro_2_18_18_105_after_img10", "models/enduro.stl"),
    "G1": ("data/data_g1_2_36_14_010_before_img6", "data/data_g1_2_36_14_010_after_img8", "models/g1.stl"),
    "GBike": ("data/data_gbike_3_31_8_018_before_img2", "data/data_gbike_3_31_8_018_after_img8", "models/gbike.stl"),
    "Holland": (
        "data/data_holland_4_24_18_00_before_img4", "data/data_holland_4_24_18_00_after_img7", "models/holland.stl"),
    "Huffy": (
        "data/data_huffy_6_4_35_012_before_img4", "data/data_huffy_6_4_35_012_after_img12", "models/huffy.stl"),
    "Kuota": (
        "data/data_kuota_6_5_30_112_before_img4", "data/data_kuota_6_5_30_112_after_img12", "models/kuota.stl"),
    "MFactory": (
        "data/data_mfactory_1_17_8_00_before_img4", "data/data_mfactory_1_17_8_00_after_img12", "models/mfactory.stl"),
    "Mirage": (
        "data/data_mirage_6_11_36_117_before_img4", "data/data_mirage_6_11_36_117_after_img12", "models/mirage.stl"),
    "Old Bike": (
        "data/data_oldbike_3_30_15_100_before_img4", "data/data_oldbike_3_30_15_100_after_img12", "models/oldbike.stl"),
    "Freeride": (
        "data/data_freeride_6_24_0_00_before_img4", "data/data_freeride_6_24_0_00_after_img12", "models/freeride_1Step_f.stl"),
    "RondoRuut": (
        "data/data_rondo_2_11_39_106_before_img4", "data/data_rondo_2_11_39_106_after_img12",
        "models/RondoRuutClean_1Step_f.stl"),
    "MTB Ghost": (
        "data/data_ghost_1_0_4_116_before_img4", "data/data_ghost_1_0_4_116_after_img12", "models/MTB_Ghost.stl"),
    "Scalpel": (
        "data/data_scalpel_6_2_12_101_before_img4", "data/data_scalpel_6_2_12_101_after_img12", "models/scalpel.stl"),
    "Verdona": (
        "data/data_verdona_1_14_10_118_before_img4", "data/data_verdona_1_14_10_118_after_img12", "models/verdona.stl"),
    "Vintage": (
        "data/data_vintage_6_6_33_011_before_img4", "data/data_vintage_6_6_33_011_after_img12", "models/Vintage.stl"),
    "WBike": (
        "data/data_wbike_3_19_23_00_before_img4", "data/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(os.path.join(os.path.dirname(__file__),f"{img_b}.png"))
            with gr.Column(scale=1):
                gr.Markdown("#### Background")
                img_bb = gr.Image(os.path.join(os.path.dirname(__file__),f"{img_b}_background.png"))
            with gr.Column(scale=1):
                gr.Markdown("#### Foreground")
                img_bf = gr.Image(os.path.join(os.path.dirname(__file__),f"{img_b}_foreground.png"))
            with gr.Column(scale=1):
                gr.Markdown("#### Segmentation")
                img_bs = gr.Image(os.path.join(os.path.dirname(__file__),f"{img_b}_segmentation.png"))

    gr.Markdown("### After damage")
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("#### Render")
            img_an = gr.Image(os.path.join(os.path.dirname(__file__),f"{img_a}.png"))
        with gr.Column(scale=1):
            gr.Markdown("#### Background")
            img_ab = gr.Image(os.path.join(os.path.dirname(__file__),f"{img_a}_background.png"))
        with gr.Column(scale=1):
            gr.Markdown("#### Foreground")
            img_af = gr.Image(os.path.join(os.path.dirname(__file__),f"{img_a}_foreground.png"))
        with gr.Column(scale=1):
            gr.Markdown("#### Segmentation")
            img_as = gr.Image(os.path.join(os.path.dirname(__file__),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])
    gr.Markdown("### Citation Information \n```\n@inproceedings{bbb_2022,\n  title={Bent & Broken Bicycles: Leveraging"
                " synthetic data for damaged object re-identification},\n  author={Luca Piano, Filippo Gabriele Pratticò,"
                " Alessandro Sebastian Russo, Lorenzo Lanari, Lia Morra, Fabrizio Lamberti},\n  booktitle={2022 IEEE "
                "Winter Conference on Applications of Computer Vision (WACV)},\n  year={2022},\n  "
                "organization={IEEE}\n}\n```")

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