import gradio as gr
import copy
import datasets

STYLE = """
.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-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;
}
"""

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

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]

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']}"),
    )

with gr.Blocks(css=STYLE) as demo:
    date_dd = gr.Dropdown(
        sorted_dates, 
        value=last_date, 
        label="Select date", 
        interactive=True
    )
    papers_dd = gr.Dropdown(
        [paper["title"] for paper in last_papers],
        value=selected_paper["title"],
        label="Select paper title", 
        interactive=True
    )

    date_dd.input(
        get_papers,
        date_dd,
        papers_dd
    )

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

    gr.Markdown("## Auto generated Questions & Answers")    

    # 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']}", 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']}", 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']}", 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']}", 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']}", 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']}", 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']}", 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']}", 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']}", elem_classes=["small-font"])

    papers_dd.input(
        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
        ]
    )

demo.launch(share=True)