import gradio as gr
import pandas as pd
import logging
import re
from task_visualizations import TaskVisualizations
import plotly.graph_objects as go

logging.basicConfig(level=logging.INFO)


class AppConfig:
    repo_representations_path = "data/repo_representations.jsonl"
    task_counts_path = "data/repos_task_counts.csv"
    selected_task_counts_path = "data/selected_repos_task_counts.csv"
    tasks_path = "data/paperswithcode_tasks.csv"


def load_repo_df(repo_representations_path):
    data = pd.read_json(repo_representations_path, lines=True, orient="records")
    return data.assign(
        text=data["text"]
        .str.replace(r"<img.*\/>", "", regex=True)
        .str.replace("│", "\n")
        .str.replace("⋮", "\n")
    )


def display_representations(repo, representation1, representation2):
    repo_data = repos_df[repos_df["repo_name"] == repo]
    logging.info(f"repo_data: {repo_data}")
    text1 = (
        repo_data[repo_data["representation"] == representation1]["text"].iloc[0]
        if not repo_data[repo_data["representation"] == representation1].empty
        else "No data available"
    )
    text2 = (
        repo_data[repo_data["representation"] == representation2]["text"].iloc[0]
        if not repo_data[repo_data["representation"] == representation2].empty
        else "No data available"
    )

    return text1, text2


def setup_repository_representations_tab(repos, representation_types):
    gr.Markdown("Select a repository and two representation types to compare them.")

    with gr.Row():
        repo = gr.Dropdown(choices=repos, label="Repository", value=repos[0])
        representation1 = gr.Dropdown(
            choices=representation_types, label="Representation 1", value="readme"
        )
        representation2 = gr.Dropdown(
            choices=representation_types,
            label="Representation 2",
            value="generated_readme",
        )

    with gr.Row():
        with gr.Column(
            elem_id="column1",
            variant="panel",
            scale=1,
            min_width=300,
        ):
            text1 = gr.Markdown()
        with gr.Column(
            elem_id="column2",
            variant="panel",
            scale=1,
            min_width=300,
        ):
            text2 = gr.Markdown()

    def update_representations(repo, representation1, representation2):
        text1_content, text2_content = display_representations(
            repo, representation1, representation2
        )
        return (
            f"### Representation 1: {representation1}\n\n{text1_content}",
            f"### Representation 2: {representation2}\n\n{text2_content}",
        )

    # Initial call to populate textboxes with default values
    text1.value, text2.value = update_representations(
        repos[0], "readme", "generated_readme"
    )

    for component in [repo, representation1, representation2]:
        component.change(
            fn=update_representations,
            inputs=[repo, representation1, representation2],
            outputs=[text1, text2],
        )


## main
repos_df = load_repo_df()
repos = list(repos_df["repo_name"].unique())
representation_types = list(repos_df["representation"].unique())
logging.info(f"found {len(repos)} repositories")
logging.info(f"representation types: {representation_types}")
task_visualizations = TaskVisualizations(
    AppConfig.task_counts_path,
    AppConfig.selected_task_counts_path,
    AppConfig.tasks_path,
)

with gr.Blocks() as demo:
    with gr.Tab("Explore Repository Representations"):
        setup_repository_representations_tab(repos, representation_types)
    with gr.Tab("Explore PapersWithCode Tasks"):
        gr.Markdown("## PapersWithCode Tasks Visualization")
        with gr.Row():
            min_task_count = gr.Slider(minimum=1, maximum=100, value=10, step=1, label="Minimum Task Count")
            update_button = gr.Button("Update Plots")
        
        with gr.Row():
            with gr.Column():
                gr.Markdown("### All Repositories")
                all_repos_plot = gr.Plot()
            with gr.Column():
                gr.Markdown("### Selected Repositories")
                selected_repos_plot = gr.Plot()
        
        def update_plots(min_count):
            all_plot, selected_plot = task_visualizations.display_tasks_sunburst_charts(min_count)
            return all_plot, selected_plot
        
        update_button.click(
            fn=update_plots,
            inputs=[min_task_count],
            outputs=[all_repos_plot, selected_repos_plot]
        )

demo.launch()