import altair as alt
import gradio as gr
import numpy as np
import pandas as pd
from vega_datasets import data


def make_plot(plot_type):
    if plot_type == "interactive_barplot":
        source = data.movies.url

        pts = alt.selection(type="single", encodings=['x'])

        rect = alt.Chart(data.movies.url).mark_rect().encode(
            alt.X('IMDB_Rating:Q', bin=True),
            alt.Y('Rotten_Tomatoes_Rating:Q', bin=True),
            alt.Color('count()',
                scale=alt.Scale(scheme='greenblue'),
                legend=alt.Legend(title='Total Records')
            )
        )

        circ = rect.mark_point().encode(
            alt.ColorValue('grey'),
            alt.Size('count()',
                legend=alt.Legend(title='Records in Selection')
            )
        ).transform_filter(
            pts
        )

        bar = alt.Chart(source).mark_bar().encode(
            x='Major_Genre:N',
            y='count()',
            color=alt.condition(pts, alt.ColorValue("steelblue"), alt.ColorValue("grey"))
        ).properties(
            width=550,
            height=200
        ).add_selection(pts)

        plot = alt.vconcat(
            rect + circ,
            bar
        ).resolve_legend(
            color="independent",
            size="independent"
        )
        return plot
    elif plot_type == "multiline":
        source = data.stocks()

        highlight = alt.selection(type='single', on='mouseover',
                                fields=['symbol'], nearest=True)

        base = alt.Chart(source).encode(
            x='date:T',
            y='price:Q',
            color='symbol:N'
        )

        points = base.mark_circle().encode(
            opacity=alt.value(0)
        ).add_selection(
            highlight
        ).properties(
            width=600
        )

        lines = base.mark_line().encode(
            size=alt.condition(~highlight, alt.value(1), alt.value(3))
        )

        return points + lines


with gr.Blocks() as demo:
    button = gr.Radio(label="Plot type",
                      choices=['interactive_barplot', "multiline"], value='interactive_barplot')
    plot = gr.Plot(label="Plot")
    button.change(make_plot, inputs=button, outputs=[plot])
    demo.load(make_plot, inputs=[button], outputs=[plot])


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