# Copyright (c) OpenMMLab. All rights reserved.
import os

os.system('pip install -U openmim')
os.system('mim install mmengine')
os.system('mim install "mmcv>=2.0.0"')

from argparse import ArgumentParser

import cv2
from mmengine.model.utils import revert_sync_batchnorm

from mmseg.apis import inference_model, init_model
from mmseg.apis.inference import show_result_pyplot
import torch
import time
import gradio as gr
import plotly.express as px
import json

def main():
    parser = ArgumentParser()
    parser.add_argument('--config', default='configs/snnet/setr_naive_512x512_160k_b16_ade20k_deit_3_s_l_224_snnetv2.py', help='Config file')
    parser.add_argument('--checkpoint', help='Checkpoint file', default='setr_naive_512x512_160k_b16_ade20k_snnetv2_deit3_s_l_lora_16_iter_160000.pth')
    # parser.add_argument('--video', help='Video file or webcam id')

    parser.add_argument(
        '--device', default='cuda:0', help='Device used for inference')
    parser.add_argument(
        '--palette',
        default='cityscapes',
        help='Color palette used for segmentation map')
    parser.add_argument(
        '--show', action='store_true', help='Whether to show draw result')
    parser.add_argument(
        '--show-wait-time', default=1, type=int, help='Wait time after imshow')
    parser.add_argument(
        '--output-file', default=None, type=str, help='Output video file path')
    parser.add_argument(
        '--output-fourcc',
        default='MJPG',
        type=str,
        help='Fourcc of the output video')
    parser.add_argument(
        '--output-fps', default=30, type=int, help='FPS of the output video')
    parser.add_argument(
        '--output-height',
        default=-1,
        type=int,
        help='Frame height of the output video')
    parser.add_argument(
        '--output-width',
        default=-1,
        type=int,
        help='Frame width of the output video')
    parser.add_argument(
        '--opacity',
        type=float,
        default=0.5,
        help='Opacity of painted segmentation map. In (0, 1] range.')
    args = parser.parse_args()

    # build the model from a config file and a checkpoint file
    model = init_model(args.config, args.checkpoint, device=args.device)
    if args.device == 'cpu':
        model = revert_sync_batchnorm(model)

    from mmseg.models.backbones.snnet import get_stitch_configs_bidirection
    stitch_configs_info, _, _, anchor_ids, sl_ids, ls_ids, lsl_ids, sls_ids = get_stitch_configs_bidirection([12, 24])

    stitch_configs_info = {i: cfg for i, cfg in enumerate(stitch_configs_info)}


    with open('./model_flops/snnet_flops_setr_naive_512x512_160k_b16_ade20k_deit_3_s_l_224_snnetv2.json', 'r') as f:
        flops_params = json.load(f)

    with open('./results/eval_single_scale_20230507_235400.json', 'r') as f:
        results = json.load(f)

    config_ids = list(results.keys())
    flops_res = {}
    eval_res = {}
    total_data = {}
    for i, cfg_id in enumerate(config_ids):
        flops = flops_params[cfg_id]
        miou_res = results[cfg_id]['metric']['mIoU'] * 100
        eval_res[int(cfg_id)] = miou_res
        flops_res[int(cfg_id)] = flops / 1e9
        total_data[int(cfg_id)] = [flops // 1e9, miou_res]


    def visualize_stitch_pos(stitch_id):
        if stitch_id == 13:
            # 13 is equivalent to 0
            stitch_id = 0

        names = [f'ID {key}' for key in flops_res.keys()]

        fig = px.scatter(x=flops_res.values(), y=eval_res.values(), hover_name=names)
        fig.update_layout(
            title=f"SN-Netv2 - Stitch ID - {stitch_id}",
            title_x=0.5,
            xaxis_title="GFLOPs",
            yaxis_title="mIoU",
            font=dict(
                family="Courier New, monospace",
                size=18,
                color="RebeccaPurple"
            ),
            legend=dict(
                yanchor="bottom",
                y=0.99,
                xanchor="left",
                x=0.01),
        )
        # continent, DarkSlateGrey
        fig.update_traces(marker=dict(size=10,
                                      line=dict(width=2)),
                          selector=dict(mode='markers'))

        fig.add_scatter(x=[flops_res[stitch_id]], y=[eval_res[stitch_id]], mode='markers', marker=dict(size=15), name='Current Stitch')
        return fig


    def segment_video(video, stitch_id):

        if stitch_id == 13:
            # 13 is equivalent to 0
            stitch_id = 0

        model.backbone.reset_stitch_id(stitch_id)
        output_video_path = './temp_video.avi'
        cap = cv2.VideoCapture(video)
        assert (cap.isOpened())
        input_height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
        input_width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
        input_fps = cap.get(cv2.CAP_PROP_FPS)


        fourcc = cv2.VideoWriter_fourcc(*args.output_fourcc)
        output_fps = args.output_fps if args.output_fps > 0 else input_fps
        output_height = args.output_height if args.output_height > 0 else int(
            input_height)
        output_width = args.output_width if args.output_width > 0 else int(
            input_width)
        writer = cv2.VideoWriter(output_video_path, fourcc, output_fps,
                                 (output_width, output_height), True)

        try:
            while True:
                start_time = time.time()
                flag, frame = cap.read()
                if not flag:
                    break

                # test a single image
                result = inference_model(model, frame)

                # blend raw image and prediction
                draw_img = show_result_pyplot(model, frame, result,
                                              show=False,
                                              with_labels=False,
                                              )

                if draw_img.shape[0] != output_height or draw_img.shape[
                    1] != output_width:
                    draw_img = cv2.resize(draw_img,
                                          (output_width, output_height))
                writer.write(draw_img)
        finally:
            if writer:
                writer.release()
            cap.release()

        fig = visualize_stitch_pos(stitch_id)
        
        return output_video_path, fig

    def segment_image(image, stitch_id):
        if stitch_id == 13:
            # 13 is equivalent to 0
            stitch_id = 0

        model.backbone.reset_stitch_id(stitch_id)
        result = inference_model(model, image)
        draw_img = show_result_pyplot(model, image, result,
                                      show=False,
                                      with_labels=True,
                                      )
        fig = visualize_stitch_pos(stitch_id)
        return draw_img, fig



    with gr.Blocks() as image_demo:
        with gr.Row():
            with gr.Column():
                image_input = gr.Image(label='Input Image')
                stitch_slider = gr.Slider(minimum=0, maximum=134, step=1, label="Stitch ID")
                with gr.Row():
                    clear_button = gr.ClearButton()
                    submit_button = gr.Button()

            with gr.Column():
                image_output = gr.Image(label='Segmentation Results')
                stitch_plot = gr.Plot(label='Stitch Position')

            submit_button.click(
                fn=segment_image,
                inputs=[image_input, stitch_slider],
                outputs=[image_output, stitch_plot],
            )

            stitch_slider.change(
                fn=visualize_stitch_pos,
                inputs=[stitch_slider],
                outputs=[stitch_plot],
                show_progress=False
            )

            clear_button.click(
                lambda: [None, 0, None, None],
                outputs=[image_input, stitch_slider, image_output, stitch_plot],
            )

        gr.Examples(
            [
                ['./demo_1.jpg', 0],
                ['./demo_2.jpg', 1],
                ['./demo_3.jpg', 93],
                ['./demo_4.jpg', 3],
            ],
            inputs=[
                image_input,
                stitch_slider
            ],
            outputs=[
                image_input,
                stitch_plot
            ],
        )

    with gr.Blocks() as demo:
        with gr.Column():
            gr.HTML("""
                <h1 align="center" style=" display: flex; flex-direction: row; justify-content: center; font-size: 25pt; ">Stitched ViTs are Flexible Vision Backbones</h1>
                <div align="center"> <img align="center" src='file/gradio_banner.png' width="70%"> </div>
                <h3 align="center" >This is the classification demo page of SN-Netv2, a flexible vision backbone that allows for 100+ runtime speed and performance trade-offs. You can also run this gradio demo on your local GPUs at <a href="https://github.com/ziplab/SN-Netv2">https://github.com/ziplab/SN-Netv2</a>, Paper link: <a href="https://arxiv.org/abs/2307.00154">https://arxiv.org/abs/2307.00154</a>.</h3>
                """)
        tabbed_page = gr.TabbedInterface([image_demo,], ['Image'])


    demo.launch(allowed_paths=['./'])


if __name__ == '__main__':
    main()