import gradio as gr
from pathlib import Path
import uuid
import random

from utils.data_utils import generate_leaderboard
from utils.plot_utils import plot_ratings
from utils.utils import simulate, submit_rating, generate_matchup
from config import MODE, VIDEOS, MODELS, CRITERIA, default_beta


head = f"""
    <script src="https://ajax.googleapis.com/ajax/libs/jquery/3.2.1/jquery.min.js"></script>
    <script src="https://maxcdn.bootstrapcdn.com/bootstrap/3.3.7/js/bootstrap.min.js"></script>
    <script src="https://cdnjs.cloudflare.com/ajax/libs/plotly.js/1.33.1/plotly.min.js"></script>
    <script>{Path('static/modelViewer.js').read_text()}</script>
    <script>{Path('static/popup.js').read_text()}</script>
    <script>{Path('static/plots.js').read_text()}</script>
"""

with gr.Blocks(title='3D Animation Arena', head=head, css_paths='static/style.css') as arena:

    sessionState = gr.State({
        'video': None,
        'modelLeft': None,
        'modelRight': None,
        'darkMode': False,
        'videos': VIDEOS,
        'currentTab': CRITERIA[0],
        'uuid': None
    })
    
    frontState = gr.JSON(sessionState, visible=False)

    with gr.Row():
        with gr.Column(scale=1):
            gr.HTML('')
        with gr.Column(scale=12):
            gr.HTML("<h1 style='text-align:center; font-size:50px'>3D Animation Arena</h1>")
        with gr.Column(scale=1):
            toggle_dark = gr.Button(value="Dark Mode")

        def update_toggle_dark(state):
            state['darkMode'] = not state['darkMode']
            if state['darkMode']:
                return gr.update(value="Light Mode"), state
            else:
                return gr.update(value="Dark Mode"), state

        toggle_dark.click(
            inputs=[sessionState],
            js="""
            () => {
                document.body.classList.toggle('dark');
            }
            """,
            fn=update_toggle_dark,
            outputs=[toggle_dark, sessionState]
        )

    with gr.Tab(label='Arena'):
        models = gr.HTML('''
            <div class="viewer-container">
                <iframe
                    id="modelViewerLeft" 
                    src="https://d39vhmln1nnc4z.cloudfront.net/index.html" 
                    width="100%" 
                    height="100%" 
                    allow="storage-access"
                ></iframe>

                <iframe 
                    id="modelViewerRight" 
                    src="https://d39vhmln1nnc4z.cloudfront.net/index.html" 
                    width="100%" 
                    height="100%" 
                    allow="storage-access"
                ></iframe>
            </div>''', 
            render=False)

        with gr.Row():
            with gr.Column(scale=1):
                gr.HTML(f"<h1>1. Choose a video below:</h1>")
                video = gr.Video(
                    label='Input Video',
                    interactive=False,
                    autoplay=True,
                    show_download_button=False,
                    loop=True,
                    elem_id='gradioVideo',
                )

                triggerButtons = {}
                for vid in sessionState.value['videos']:
                    triggerButtons[vid] = gr.Button(elem_id=f'triggerBtn_{vid}', visible=False)
                    triggerButtons[vid].click(
                        fn=lambda vid=vid: gr.update(value=f'https://gradio-model-viewer.s3.eu-west-1.amazonaws.com/sample+videos/{vid}.mp4'),
                        outputs=[video]
                    )
                examples = gr.HTML(visible=False)

            with gr.Column(scale=4):
                gr.HTML("""
                    <h1>2. Play around with the models:
                        <span class="glyphicon glyphicon-question-sign popup-btn btn btn-info btn-lg" data-popup-id="instructionsPopup">
                            <span class="popup-text" id="instructionsPopup">You can control the playback in both viewers at the same time by using the video, or control both viewers independently by using mouse and GUI!</span>
                        </span>
                    </h1>
                """)
                with gr.Row():
                    models.render()

                with gr.Row():
                    gr.HTML(f"<h1>3. Choose your favorite model for each criteria:</h1>")
                ratingButtons = {}
                for criteria in CRITERIA:
                    with gr.Row():                
                        with gr.Column():
                            with gr.Row():
                                match criteria:
                                    case 'Global_Appreciation':
                                        instructions = "Your overall appreciation of the models, including general aesthetics and self-contacts if applicable."
                                    case 'Ground_Contacts':
                                        instructions = "The quality of the models' contacts with the ground, including ground penetration and foot sliding."
                                    case 'Fidelity':
                                        instructions = "The fidelity of the models compared to the motion of the original video."
                                    case 'Fluidity':
                                        instructions = "The smoothness and temporal coherence of the models."
                                gr.HTML(f"""
                                    <h2 style='text-align:center;'>{criteria.replace('_', ' ')}
                                    <span class="glyphicon glyphicon-question-sign popup-btn btn btn-info btn-lg" data-popup-id="{criteria}Popup">
                                        <span class="popup-text" id="{criteria}Popup">{instructions}</span>
                                    </span></h2>
                                """)
                            with gr.Row():
                                ratingButtons[criteria] = []
                                with gr.Column(scale=2):
                                    ratingButtons[criteria].append(gr.Button('Left Model', variant='primary', interactive=False))
                                with gr.Column(scale=1, min_width=2):
                                    ratingButtons[criteria].append(gr.Button('Skip', min_width=2, interactive=False))
                                with gr.Column(scale=2):
                                    ratingButtons[criteria].append(gr.Button('Right Model', variant='primary', interactive=False))
    

    # Leaderboard per criteria
    with gr.Tab(label='Leaderboards') as leaderboard_tab:

        if MODE == 'testing':
            # Simulation controls
            with gr.Row():
                simulate_btn = gr.Button('Simulate Matches', variant='primary')
                add_model_btn = gr.Button('Add Model', variant='secondary')
            with gr.Row():
                gr.Markdown('''
                ## Probability of each model to be chosen is updated after each vote following: \
                $$ p_i = \\frac{e^{-\\frac{Matches_i}{\\beta}}}{\\sum_{j=1}^{N} e^{-\\frac{Matches_j}{\\beta}}} $$  
                    ''')
                iterate = gr.Number(label='Number of iterations', value=100, minimum=1, maximum=2000, precision=0, interactive=True)
                beta = gr.Number(label='Beta', value=default_beta, minimum=1, maximum=1000, precision=0, step=10, interactive=True)
        else:
            beta = gr.Number(label='Beta', value=default_beta, render=False)

        leaderboards = {}
        tabs = {}
        for criteria in CRITERIA:
            with gr.Tab(label=criteria.replace('_', ' ')) as tabs[criteria]:
                with gr.Row():
                    gr.HTML(f"<h2 style='text-align:center;'>{criteria.replace('_', ' ')}</h2>")
                with gr.Row():
                    leaderboards[criteria] = gr.Dataframe(value=None, row_count=(len(MODELS), 'fixed'), headers=['Model', 'Elo', 'Wins', 'Matches', 'Win Rate'], interactive=False)

        # Plots
        if MODE == 'testing':
            with gr.Row():
                elo_plot = gr.Plot(value=None, label='Elo Ratings', format='plotly', elem_id='plot')
            with gr.Row():
                wr_plot = gr.Plot(value=None, label='Win Rates', format='plotly', elem_id='plot')
            with gr.Row():
                matches_plot = gr.Plot(value=None, label='Matches played', format='plotly', elem_id='plot')
        elif MODE == 'production':
            elo_plot = gr.Plot(value=None, label='Elo Ratings', format='plotly', elem_id='plot', visible=False)
            wr_plot = gr.Plot(value=None, label='Win Rates', format='plotly', elem_id='plot', visible=False)
            matches_plot = gr.Plot(value=None, label='Matches played', format='plotly', elem_id='plot', visible=False)

    with gr.Tab(label='About'):
        gr.Markdown('''
        ## Thank you for using the 3D Animation Arena!
                        
        This app is designed to compare different models based on human preferences, inspired by dylanebert's [3D Arena](https://huggingface.co/spaces/dylanebert/3d-arena) on Hugging Face. 
        Current rankings often use metrics to assess the quality of a model, but these metrics may not always reflect the complexity behind human preferences.

        The current models competing in the arena are:
        -  4DHumans (https://github.com/shubham-goel/4D-Humans)
        - CLIFF (https://github.com/haofanwang/CLIFF)
        - GVHMR (https://github.com/zju3dv/GVHMR)
        - HybrIK (https://github.com/jeffffffli/HybrIK)
        - WHAM (https://github.com/yohanshin/WHAM)
                
        All inferences are precomputed following the code in the associated GitHub repository. 
        Some post-inference modifications have been made to some models in order to make the comparison possible. 
        These modifications include:
        * Adjusting height to a common ground
        * Fixing the root depth of certain models, when depth was extremely jittery
        * Fixing the root position of certain models, when no root position was available

        All models use the SMPL body model to discard the influence of the body model on the comparison.
        These choices were made without any intention to favor or harm any model. 
        All matchups are generated randomly, don't hesitate to rate the same videos multiple times as the matchups will probably be different!

        ---

        If you have comments, complaints or suggestions, please contact me at 3danimationarena@gmail.com. 
        New models and videos will be added over time, feel free to share your ideas! Keep in mind that I will not add raw inferences from other people to keep it fair.  
        ''')


    # Event handlers
    def randomize_videos(state):
        state['uuid'] = str(uuid.uuid4())
        random.shuffle(state['videos'])
        gallery = "<div class='gallery'>"
        for vid in state['videos']:
            gallery += f"""
                <button class="btn btn-info thumbnail-btn" onclick="(function() {{
                    let gradioVideo = document.getElementById('gradioVideo');
                    let videoComponent = gradioVideo ? gradioVideo.querySelector('video') : null;
                    if (videoComponent && !videoComponent.src.includes('{vid}')) {{
                        Array.from(document.getElementsByClassName('thumbnail-btn')).forEach(btn => btn.disabled = true);
                    }}
                    document.getElementById('triggerBtn_{vid}').click();
                    }})()">
                    <video class="thumbnail" preload="" loop muted onmouseenter="this.play()" onmouseleave="this.pause()">
                        <source src="https://gradio-model-viewer.s3.eu-west-1.amazonaws.com/sample+videos/{vid}.mp4">
                    </video>
                </button>
            """
        gallery += "</div>"
        return state, gallery

    async def display_leaderboards():
        return [await generate_leaderboard(criteria) for criteria in CRITERIA]
    
    arena.load(
        inputs=[sessionState],
        fn=lambda state: randomize_videos(state),
        outputs=[sessionState, examples],
    ).then(
        inputs=[],
        fn=lambda: gr.update(visible=True),
        outputs=[examples]
    ).then(
        inputs=[gr.State(CRITERIA[0])],
        fn=plot_ratings,  
        outputs=[elo_plot, wr_plot, matches_plot]
    ).then(
        inputs=[],
        fn=display_leaderboards,
        outputs=[leaderboards[criteria] for criteria in CRITERIA]
    )

    async def update_models(video, state):
        leaderboard = await generate_leaderboard(CRITERIA[0])
        video_name = video.split('/')[-1].split('.')[0]
        modelLeft, modelRight = generate_matchup(leaderboard=leaderboard, beta=beta.value)

        state['video'] = video_name
        state['modelLeft'] = MODELS[modelLeft]
        state['modelRight'] = MODELS[modelRight]

        return state, state

    video.change(
        inputs=[video, sessionState],
        fn=update_models,
        outputs=[sessionState, frontState]
    )

    # Weird workaround to run JS function on state change, from https://github.com/gradio-app/gradio/issues/3525#issuecomment-2348596861
    frontState.change(
        inputs=[frontState],
        js='(state) => updateViewers(state)',
        fn=lambda state: None,
    ).then(
        inputs=None,
        fn=lambda: tuple(gr.update(interactive=True) for _ in sum(ratingButtons.values(), [])),
        outputs= sum(ratingButtons.values(), [])
    )

    leaderboard_tab.select(
        inputs=None,
        js='() => resetPlots()',
        fn=None,
    ).then(
        fn=lambda: [gr.update(value=None) for _ in range(3)],
        outputs=[elo_plot, wr_plot, matches_plot]
    ).then(
        inputs=[sessionState],
        fn=lambda state: plot_ratings(state['currentTab']),
        outputs=[elo_plot, wr_plot, matches_plot]
    )

    async def process_rating(state, i, criteria):
        return gr.update(value=await submit_rating(
            criteria=criteria, 
            video=state['video'],
            winner=state['modelLeft'] if i == 0 else state['modelRight'] if i == 2 else None, 
            loser=state['modelRight'] if i == 0 else state['modelLeft'] if i == 2 else None,
            uuid=state['uuid']
        ))

    def update_tab(state, criteria):
        state['currentTab'] = criteria
        return state
    
    for criteria in CRITERIA:
        for i, button in enumerate(ratingButtons[criteria]):
            button.click(
            #     fn=lambda i=i, criteria=criteria: gr.Info(f'{"You chose Left Model for " if i == 0 else "You chose Right Model for " if i == 2 else "You skipped "} {criteria.replace("_", " ")}!'),
            # ).then(
                fn=lambda: tuple(gr.update(interactive=False) for _ in range(len(ratingButtons[criteria]))),
                outputs=ratingButtons[criteria]
            ).then(
                inputs=[sessionState, gr.State(i), gr.State(criteria)],
                fn=process_rating,
                outputs=[leaderboards[criteria]],
            )

        tabs[criteria].select(
            fn=lambda: [gr.update(value=None) for _ in range(3)],
            outputs=[elo_plot, wr_plot, matches_plot]
        ).then(
            inputs=[gr.State(criteria)],
            fn=plot_ratings,  
            outputs=[elo_plot, wr_plot, matches_plot]
        ).then(
            inputs=[sessionState, gr.State(criteria)],
            fn=update_tab,
            outputs=[sessionState]
        )


    if MODE == 'testing':
        for criteria in CRITERIA:
            simulate_btn.click(
                inputs=[iterate, beta, gr.State(criteria)],
                fn=simulate,
                outputs=[leaderboards[criteria]],
            ).then(fn=lambda: [gr.update(value=None) for _ in range(3)],
                outputs=[elo_plot, wr_plot, matches_plot]
            ).then(
                inputs=[gr.State(criteria)],
                fn=plot_ratings,  
                outputs=[elo_plot, wr_plot, matches_plot]
            )

        add_model_btn.click(
            fn=lambda: MODELS.append(f'model_{len(MODELS)}'),
        )

if __name__ == '__main__':
    gr.set_static_paths(['static'])
    arena.queue(default_concurrency_limit=50).launch(inbrowser=True, allowed_paths=['static/'])