import gradio as gr
import copy
import datasets

STYLE = """

.main {
  width: 90% !important;
  margin: 0 auto; /* Center the container */
}

.small-font{
  font-size: 12pt !important;
}

.small-font:hover {
  font-size: 20px !important;
  transition: font-size 0.3s ease-out;
  transition-delay: 0.8s;
}

.group {
  padding-top: 10px;
  padding-left: 10px;
  padding-right: 10px;
  padding-bottom: 10px;
  border: 2px dashed gray;
  border-radius: 20px;
  box-shadow: 5px 3px 10px 1px rgba(0, 0, 0, 0.4) !important;
}

.accordion > button > span{
  font-size: 12pt !important;
}

.accordion {
  border-style: dashed !important;
  border-left-width: 2px !important;
  border-bottom-width: 2.5px !important;
  border-top: none !important;
  border-right: none !important;
  box-shadow: none !important;
}

.no-gap {
    gap: 0px;
}

.no-radius {
    border-radius: 0px;    
}

#search_input > label > span {
    display: none;
}

#exp-type > span {
    display: none;
}
"""

dataset_repo_id = "chansung/auto-paper-qa2"
ds = datasets.load_dataset(dataset_repo_id)

title2qna = {}
date2qna = {}
longest_qans = 0

def count_nans(row):
    count = 0

    for _, (k, v) in enumerate(data.items()):
        if v is None:
            count = count + 1

    return count

for data in ds["train"]:
    date = data["target_date"].strftime("%Y-%m-%d")

    if date in date2qna:
        papers = copy.deepcopy(date2qna[date])
        for paper in papers:
            if paper["title"] == data["title"]:
                if count_nans(paper) > count_nans(data):
                    date2qna[date].remove(paper)
        
        date2qna[date].append(data)
        del papers
    else:
        date2qna[date] = [data]

for date in date2qna:
    papers = date2qna[date]
    for paper in papers:
        title2qna[paper["title"]] = paper

titles = title2qna.keys()

sorted_dates = sorted(date2qna.keys())
last_date = sorted_dates[-1]
last_papers = date2qna[last_date]
selected_paper = last_papers[0]

def get_papers(date):
    papers = [paper["title"] for paper in date2qna[date]]
    return gr.Dropdown(
        papers,
        value=papers[0]
    )

def set_paper(date, paper_title):
    selected_paper = None
    for paper in date2qna[date]:
        if paper["title"] == paper_title:
            selected_paper = paper
            break

    return (
        gr.Markdown(f"# {selected_paper['title']}"), gr.Markdown(selected_paper["summary"]),

        gr.Markdown(f"## 🙋 {selected_paper['0_question']}"), 
        gr.Markdown(f"↪ **(ELI5)** {selected_paper['0_answers:eli5']}"), 
        gr.Markdown(f"↪ **(Technical)** {selected_paper['0_answers:expert']}"),
        gr.Markdown(f"## 🙋🙋 {selected_paper['0_additional_depth_q:follow up question']}"),
        gr.Markdown(f"↪ **(ELI5)** {selected_paper['0_additional_depth_q:answers:eli5']}"), 
        gr.Markdown(f"↪ **(Technical)** {selected_paper['0_additional_depth_q:answers:expert']}"),
        gr.Markdown(f"## 🙋🙋 {selected_paper['0_additional_breath_q:follow up question']}"), 
        gr.Markdown(f"↪ **(ELI5)** {selected_paper['0_additional_breath_q:answers:eli5']}"), 
        gr.Markdown(f"↪ **(Technical)** {selected_paper['0_additional_breath_q:answers:expert']}"),

        gr.Markdown(f"## 🙋 {selected_paper['1_question']}"), 
        gr.Markdown(f"↪ **(ELI5)** {selected_paper['1_answers:eli5']}"), 
        gr.Markdown(f"↪ **(Technical)** {selected_paper['1_answers:expert']}"),
        gr.Markdown(f"## 🙋🙋 {selected_paper['1_additional_depth_q:follow up question']}"), 
        gr.Markdown(f"↪ **(ELI5)** {selected_paper['1_additional_depth_q:answers:eli5']}"), 
        gr.Markdown(f"↪ **(Technical)** {selected_paper['1_additional_depth_q:answers:expert']}"),
        gr.Markdown(f"## 🙋🙋 {selected_paper['1_additional_breath_q:follow up question']}"), 
        gr.Markdown(f"↪ **(ELI5)** {selected_paper['1_additional_breath_q:answers:eli5']}"), 
        gr.Markdown(f"↪ **(Technical)** {selected_paper['1_additional_breath_q:answers:expert']}"),

        gr.Markdown(f"## 🙋 {selected_paper['2_question']}"), 
        gr.Markdown(f"↪ **(ELI5)** {selected_paper['2_answers:eli5']}"), 
        gr.Markdown(f"↪ **(Technical)** {selected_paper['2_answers:expert']}"),
        gr.Markdown(f"## 🙋🙋 {selected_paper['2_additional_depth_q:follow up question']}"), 
        gr.Markdown(f"↪ **(ELI5)** {selected_paper['2_additional_depth_q:answers:eli5']}"), 
        gr.Markdown(f"↪ **(Technical)** {selected_paper['2_additional_depth_q:answers:expert']}"),
        gr.Markdown(f"## 🙋🙋 {selected_paper['2_additional_breath_q:follow up question']}"), 
        gr.Markdown(f"↪ **(ELI5)** {selected_paper['2_additional_breath_q:answers:eli5']}"), 
        gr.Markdown(f"↪ **(Technical)** {selected_paper['2_additional_breath_q:answers:expert']}"),
    )

def change_exp_type(exp_type):
    if exp_type == "ELI5":
        return (
            gr.Markdown(visible=True), gr.Markdown(visible=False), gr.Markdown(visible=True), gr.Markdown(visible=False), gr.Markdown(visible=True), gr.Markdown(visible=False),
            gr.Markdown(visible=True), gr.Markdown(visible=False), gr.Markdown(visible=True), gr.Markdown(visible=False), gr.Markdown(visible=True), gr.Markdown(visible=False),
            gr.Markdown(visible=True), gr.Markdown(visible=False), gr.Markdown(visible=True), gr.Markdown(visible=False), gr.Markdown(visible=True), gr.Markdown(visible=False),
        )
    else:
        return (
            gr.Markdown(visible=False), gr.Markdown(visible=True), gr.Markdown(visible=False), gr.Markdown(visible=True), gr.Markdown(visible=False), gr.Markdown(visible=True),
            gr.Markdown(visible=False), gr.Markdown(visible=True), gr.Markdown(visible=False), gr.Markdown(visible=True), gr.Markdown(visible=False), gr.Markdown(visible=True),
            gr.Markdown(visible=False), gr.Markdown(visible=True), gr.Markdown(visible=False), gr.Markdown(visible=True), gr.Markdown(visible=False), gr.Markdown(visible=True),
        )        

def search(search_in, max_results=3):
    results = []

    for title in titles:
        if len(results) > 3:
            break
        else:
            if search_in in title:
                results.append(title)

    return (
        gr.Textbox(
            visible=True if len(results) > 0 else False,
            value=results[0] if len(results) > 0 else ""
        ),
        gr.Textbox(
            visible=True if len(results) > 1 else False,
            value=results[1] if len(results) > 1 else ""
        ),
        gr.Textbox(
            visible=True if len(results) > 2 else False,
            value=results[2] if len(results) > 2 else ""
        )
    )

UPDATE_SEARCH_RESULTS = f"""
function search(searchIn, maxResults = 3) {{
    if (searchIn.trim().length > 0) {{
        const results = [];
        let titles = {list(titles)};

        for (const title of titles) {{ // Assuming 'titles' is an array defined elsewhere
            if (results.length > 3) {{
                break;
            }} else {{
                if (title.toLowerCase().includes(searchIn.toLowerCase())) {{ // JavaScript's equivalent to Python's 'in'
                    results.push(title);
                }}
            }}
        }}

        // Handle UI elements (Explanation below)
        const resultElements = [1, 2, 3].map(index => {{
            return results[index - 1] || '';
        }});

        if (resultElements[0] == '') {{
            document.getElementById('search_r1').style.display = 'none';
        }} else {{
            document.getElementById('search_r1').style.display = 'block';
        }}

        if (resultElements[1] == '') {{
            document.getElementById('search_r2').style.display = 'none';
        }} else {{
            document.getElementById('search_r2').style.display = 'block';
        }}

        if (resultElements[2] == '') {{
            document.getElementById('search_r3').style.display = 'none';
        }} else {{
            document.getElementById('search_r3').style.display = 'block';
        }}

        return resultElements; 
    }} else {{
        document.getElementById('search_r1').style.display = 'none';
        document.getElementById('search_r2').style.display = 'none';
        document.getElementById('search_r3').style.display = 'none';

        return ['', '', '']
    }}
}}
"""

def set_date(title):
    paper = title2qna[title]
    date = paper["target_date"].strftime("%Y-%m-%d")
    return date

def set_papers(date, title):
    papers = [paper["title"] for paper in date2qna[date]]
    return (
        gr.Dropdown(choices=papers, value=title),
        gr.Textbox("")
    )

with gr.Blocks(css=STYLE) as demo:
    gr.Markdown("# Let's explore papers with auto generated Q&As")
    
    with gr.Column(elem_classes=["group"]):
        with gr.Row():
            date_dd = gr.Dropdown(
                sorted_dates, 
                value=last_date, 
                label="Select date", 
                interactive=True,
                scale=3,
            )
            papers_dd = gr.Dropdown(
                [paper["title"] for paper in last_papers],
                value=selected_paper["title"],
                label="Select paper title", 
                interactive=True,
                scale=7,
            )

        with gr.Column(elem_classes=["no-gap"]):
            search_in = gr.Textbox("", placeholder="Enter keywords to search...", elem_id="search_input")
            search_r1 = gr.Button(visible=False, elem_id="search_r1", elem_classes=["no-radius"])
            search_r2 = gr.Button(visible=False, elem_id="search_r2", elem_classes=["no-radius"])
            search_r3 = gr.Button(visible=False, elem_id="search_r3", elem_classes=["no-radius"])

    title = gr.Markdown(f"# {selected_paper['title']}")
    summary = gr.Markdown(f"{selected_paper['summary']}", elem_classes=["small-font"])

    with gr.Row():
        with gr.Column(scale=7):
            gr.Markdown("## Auto generated Questions & Answers")

        exp_type = gr.Radio(choices=["ELI5", "Technical"], value="ELI5", elem_id="exp-type", scale=3)

    # 1
    with gr.Column(elem_classes=["group"], visible=True) as q_0:
        basic_q_0 = gr.Markdown(f"## 🙋 {selected_paper['0_question']}")
        basic_q_eli5_0 = gr.Markdown(f"↪ **(ELI5)** {selected_paper['0_answers:eli5']}", elem_classes=["small-font"]) 
        basic_q_expert_0 = gr.Markdown(f"↪ **(Technical)** {selected_paper['0_answers:expert']}", visible=False, elem_classes=["small-font"]) 

        with gr.Accordion("Additional question #1", open=False, elem_classes=["accordion"]) as aq_0_0:
            depth_q_0 = gr.Markdown(f"## 🙋🙋 {selected_paper['0_additional_depth_q:follow up question']}")
            depth_q_eli5_0 = gr.Markdown(f"↪ **(ELI5)** {selected_paper['0_additional_depth_q:answers:eli5']}", elem_classes=["small-font"])
            depth_q_expert_0 = gr.Markdown(f"↪ **(Technical)** {selected_paper['0_additional_depth_q:answers:expert']}", visible=False, elem_classes=["small-font"])

        with gr.Accordion("Additional question #2", open=False, elem_classes=["accordion"]) as aq_0_1:
            breath_q_0 = gr.Markdown(f"## 🙋🙋 {selected_paper['0_additional_breath_q:follow up question']}")
            breath_q_eli5_0 = gr.Markdown(f"↪ **(ELI5)** {selected_paper['0_additional_breath_q:answers:eli5']}", elem_classes=["small-font"])
            breath_q_expert_0 = gr.Markdown(f"↪ **(Technical)** {selected_paper['0_additional_breath_q:answers:expert']}", visible=False, elem_classes=["small-font"])

    # 2
    with gr.Column(elem_classes=["group"], visible=True) as q_1:
        basic_q_1 = gr.Markdown(f"## 🙋 {selected_paper['1_question']}")
        basic_q_eli5_1 = gr.Markdown(f"↪ **(ELI5)** {selected_paper['1_answers:eli5']}", elem_classes=["small-font"]) 
        basic_q_expert_1 = gr.Markdown(f"↪ **(Technical)** {selected_paper['1_answers:expert']}", visible=False, elem_classes=["small-font"]) 

        with gr.Accordion("Additional question #1", open=False, elem_classes=["accordion"]) as aq_1_0:
            depth_q_1 = gr.Markdown(f"## 🙋🙋 {selected_paper['1_additional_depth_q:follow up question']}")
            depth_q_eli5_1 = gr.Markdown(f"↪ **(ELI5)** {selected_paper['1_additional_depth_q:answers:eli5']}", elem_classes=["small-font"])
            depth_q_expert_1 = gr.Markdown(f"↪ **(Technical)** {selected_paper['1_additional_depth_q:answers:expert']}", visible=False, elem_classes=["small-font"])

        with gr.Accordion("Additional question #2", open=False, elem_classes=["accordion"]) as aq_1_1:
            breath_q_1 = gr.Markdown(f"## 🙋🙋 {selected_paper['1_additional_breath_q:follow up question']}")
            breath_q_eli5_1 = gr.Markdown(f"↪ **(ELI5)** {selected_paper['1_additional_breath_q:answers:eli5']}", elem_classes=["small-font"])
            breath_q_expert_1 = gr.Markdown(f"↪ **(Technical)** {selected_paper['1_additional_breath_q:answers:expert']}", visible=False, elem_classes=["small-font"])

    # 3
    with gr.Column(elem_classes=["group"], visible=True) as q_2:
        basic_q_2 = gr.Markdown(f"## 🙋 {selected_paper['2_question']}")
        basic_q_eli5_2 = gr.Markdown(f"↪ **(ELI5)** {selected_paper['2_answers:eli5']}", elem_classes=["small-font"]) 
        basic_q_expert_2 = gr.Markdown(f"↪ **(Technical)** {selected_paper['2_answers:expert']}", visible=False, elem_classes=["small-font"]) 

        with gr.Accordion("Additional question #1", open=False, elem_classes=["accordion"]) as aq_2_0:
            depth_q_2 = gr.Markdown(f"## 🙋🙋 {selected_paper['2_additional_depth_q:follow up question']}")
            depth_q_eli5_2 = gr.Markdown(f"↪ **(ELI5)** {selected_paper['2_additional_depth_q:answers:eli5']}", elem_classes=["small-font"])
            depth_q_expert_2 = gr.Markdown(f"↪ **(Technical)** {selected_paper['2_additional_depth_q:answers:expert']}", visible=False, elem_classes=["small-font"])

        with gr.Accordion("Additional question #2", open=False, elem_classes=["accordion"]) as aq_2_1:
            breath_q_2 = gr.Markdown(f"## 🙋🙋 {selected_paper['2_additional_breath_q:follow up question']}")
            breath_q_eli5_2 = gr.Markdown(f"↪ **(ELI5)** {selected_paper['2_additional_breath_q:answers:eli5']}", elem_classes=["small-font"])
            breath_q_expert_2 = gr.Markdown(f"↪ **(Technical)** {selected_paper['2_additional_breath_q:answers:expert']}", visible=False, elem_classes=["small-font"])

    gr.Markdown("The target papers are collected from [Hugging Face 🤗 Daily Papers](https://huggingface.co/papers) on a daily basis. "
                "The entire data is generated by [Google's Gemini 1.0](https://deepmind.google/technologies/gemini/) Pro. "
                "If you are curious how it is done, visit the [Auto Paper Q&A Generation project repository](https://github.com/deep-diver/auto-paper-analysis) "
                "Also, the generated dataset is hosted on Hugging Face 🤗 Dataset repository as well([Link](https://huggingface.co/datasets/chansung/auto-paper-qa2)). ")
    
    search_r1.click(
        set_date,
        search_r1,
        date_dd
    ).then(
        set_papers,
        inputs=[date_dd, search_r1],
        outputs=[papers_dd, search_in]
    )

    search_r2.click(
        set_date,
        search_r2,
        date_dd
    ).then(
        set_papers,
        inputs=[date_dd, search_r2],
        outputs=[papers_dd, search_in]
    )

    search_r3.click(
        set_date,
        search_r3,
        date_dd
    ).then(
        set_papers,
        inputs=[date_dd, search_r3],
        outputs=[papers_dd, search_in]
    )

    date_dd.input(
        get_papers,
        date_dd,
        papers_dd
    ).then(
        set_paper,
        [date_dd, papers_dd],
        [
            title, summary,
            basic_q_0, basic_q_eli5_0, basic_q_expert_0,
            depth_q_0, depth_q_eli5_0, depth_q_expert_0,
            breath_q_0, breath_q_eli5_0, breath_q_expert_0,

            basic_q_1, basic_q_eli5_1, basic_q_expert_1,
            depth_q_1, depth_q_eli5_1, depth_q_expert_1,
            breath_q_1, breath_q_eli5_1, breath_q_expert_1,

            basic_q_2, basic_q_eli5_2, basic_q_expert_2,
            depth_q_2, depth_q_eli5_2, depth_q_expert_2,
            breath_q_2, breath_q_eli5_2, breath_q_expert_2
        ]        
    )

    papers_dd.change(
        set_paper,
        [date_dd, papers_dd],
        [
            title, summary,
            basic_q_0, basic_q_eli5_0, basic_q_expert_0,
            depth_q_0, depth_q_eli5_0, depth_q_expert_0,
            breath_q_0, breath_q_eli5_0, breath_q_expert_0,

            basic_q_1, basic_q_eli5_1, basic_q_expert_1,
            depth_q_1, depth_q_eli5_1, depth_q_expert_1,
            breath_q_1, breath_q_eli5_1, breath_q_expert_1,

            basic_q_2, basic_q_eli5_2, basic_q_expert_2,
            depth_q_2, depth_q_eli5_2, depth_q_expert_2,
            breath_q_2, breath_q_eli5_2, breath_q_expert_2
        ]
    )

    search_in.change(
        inputs=[search_in],
        outputs=[search_r1, search_r2, search_r3],
        js=UPDATE_SEARCH_RESULTS,
        fn=None
    )

    exp_type.select(
        change_exp_type,
        exp_type,
        [
            basic_q_eli5_0, basic_q_expert_0, depth_q_eli5_0, depth_q_expert_0, breath_q_eli5_0, breath_q_expert_0,
            basic_q_eli5_1, basic_q_expert_1, depth_q_eli5_1, depth_q_expert_1, breath_q_eli5_1, breath_q_expert_1,
            basic_q_eli5_2, basic_q_expert_2, depth_q_eli5_2, depth_q_expert_2, breath_q_eli5_2, breath_q_expert_2
        ]
    )

demo.launch(share=True)