import streamlit as st
import pandas as pd
import ast
import openai
from openai import OpenAI
import PyPDF2
from PyPDF2 import PdfReader
import openreview
import json
import urllib.request
import requests
import os
import pickle
my_api_key = os.environ["OPENAI_API_KEY"]
client_openai = OpenAI(api_key=my_api_key)


# Check if the password matches the env variable
def check_password():
    def password_entered():
        if st.session_state["password"] == os.environ["APP_PASSCODE"]:
            st.session_state["password_correct"] = True
            del st.session_state["password"]  # Clear password from memory
        else:
            st.session_state["password_correct"] = False

    if "password_correct" not in st.session_state:
        st.text_input("Enter Password", type="password", on_change=password_entered, key="password")
        return False
    elif not st.session_state["password_correct"]:
        st.error("Password incorrect")
        return False
    else:
        return True

# Ask for password at the beginning
if not check_password():
    st.stop()

# converts paper pdf into text
def pdf_to_text(pdf):
    paper_text = ''

    reader = PdfReader(pdf)
    number_of_pages = len(reader.pages)
    for p in range(number_of_pages):
        page = reader.pages[p].get_object()
        text = page.extract_text()
        paper_text += text

    return paper_text

# loads all 2024 ICLR submissions from OpenReview
submissions_path = "./submissions_iclr2024.pkl"
# API V2
client = openreview.api.OpenReviewClient(
    baseurl='https://api2.openreview.net',
)

# Get ICLR venues
venues = client.get_group(id='venues').members
iclr_venues = [v for v in venues if "iclr.cc" in v.lower() and "conference" in v.lower()]

@st.cache_data
def load_ICLR_submissions():
    if os.path.exists(submissions_path):
        with open(submissions_path, "rb") as f:
            return pickle.load(f)
    else:
        # Select the second latest venue
        venue_id = iclr_venues[-2]
        venue_group = client.get_group(venue_id)
        submission_name = venue_group.content['submission_name']['value']
        submissions = client.get_all_notes(invitation=f'{venue_id}/-/{submission_name}', details='replies')
        with open(submissions_path, "wb") as f:
            pickle.dump(submissions, f)
        return submissions

# returns review and pdf paper text given paper id and reviewer id
def parse_openreview_id(submissions, paper_id, reviewer_id):
    paper_url = f"https://openreview.net/pdf?id={paper_id}"
    fields = ['summary', 'soundness', 'presentation', 'contribution', 'strengths', 'weaknesses', 'questions']

    ind = next((index for index, entry in enumerate(submissions) if entry.id == paper_id), None)
    curr_review = submissions[ind]
    bad_review = next((i for i, entry in enumerate(curr_review.details['replies'])
              if entry['signatures'][0].split('_')[-1] == reviewer_id), None)

    try:
        review_content = curr_review.details["replies"][bad_review]["content"]
    except:
        st.error("Paper ID or reviewer ID are incorrect")

    review_content = {key: review_content[key] if key in fields else None for key in review_content.keys() if key in fields}

    urllib.request.urlretrieve(paper_url, f"{paper_id}.pdf")
    pdf = f"{paper_id}.pdf"
    pdf_text = pdf_to_text(pdf)

    return review_content, pdf, pdf_text


def create_feedback(review, pdf_text, agent_prompt):

    messages = [{
        "role": "user",
        "content": [
      {
        "type": "text",
        "text": agent_prompt
      },
      {
        "type": "text",
        "text": "Here is the ML conference review"
      },
      {
        "type": "text",
        "text": json.dumps(review) #json.dumps(review)
      },
      {
        "type": "text",
        "text": "Finally, read the paper this review was written about"
      },
      {
        "type": "text",
        "text": pdf_text
      }
      ]}]

    headers = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {my_api_key}"
    }
    payload = {
    "model": "gpt-4o",
    "messages": messages,
    "max_tokens": 1000,
    "temperature": 0.2
    }

    try:
        response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
        feedback = response.json()["choices"][0]["message"]["content"]
    except Exception as e:
        print(f"An unexpected error occurred: {e}")
        feedback = "an error occured"

    return feedback


def critic(review, feedback, pdf_text, critic_prompt):
    
    messages = [{
        "role": "user",
        "content": [
      {
        "type": "text",
        "text": critic_prompt
      },
      {
        "type": "text",
        "text": "Here is the ML conference review"
      },
      {
        "type": "text",
        "text": json.dumps(review) #json.dumps(review)
      },
      {
        "type": "text",
        "text": "Here is the feedback about the review"
      },
      {
        "type": "text",
        "text": feedback
      },
      {
        "type": "text",
        "text": "Finally, read the paper this review was written about"
      },
      {
        "type": "text",
        "text": pdf_text
      }
      ]}]

    headers = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {my_api_key}"
    }
    payload = {
    "model": "gpt-4o",
    "messages": messages,
    "max_tokens": 1000,
    "temperature": 0.2
    }

    try:
        response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
        revised_feedback = response.json()["choices"][0]["message"]["content"]
    except Exception as e:
        print(f"An unexpected error occurred: {e}")
        revised_feedback = "an error occured"

    return revised_feedback

agent_prompt = """
You are given a peer review of a machine learning paper submitted to a top-tier ML conference on OpenReview.
First, you will read the text of the review and then the paper about which it was written.
Give feedback to the reviewer so that it becomes a high-quality review.
Here is a checklist. For each item, determine if the review passes or fails. If it fails, provide feedback accordingly. Follow the below guidelines for your feedback:

1. The reviewer does not ask for anything that is OBVIOUSLY already present in the paper. You should be ABSOLUTELY certain the reviewer made an error, otherwise do NOT comment on their point.
If you find such a case, quote the specific statement from the review and provide the exact quote from the paper that directly addresses or contradicts it. 
Ensure that the quote from the paper specifically answers the claim or request made by the reviewer, and explain why the selected quote is relevant.
If the reviewer has a question, you do not need to respond. If you respond, point to a specific section or quote from the paper and ask “does this provide what you were looking for?”
2. The review does not make any vague, unjustifiable, or unsupported claims. 
If there are any vague comments in the review, ask the reviewer to be more SPECIFIC and DETAILED.
3. If the reviewer claims the paper isn't novel, they should explain why and specify the paper(s) they believe this work is similar to. 
If they have not already done this, ask the reviewer to do so.
4. The reviewer does not make any personal attacks against the paper and/or author(s). 
For example, they call the work "incompetent" without justifying why. 
If they do, suggest the reviewer revise their feedback to remove those attacks.

Be concise, limiting your response to 2-3 sentences per checklist item. Do not repeat anything that the reviewer already included in their review, and do not summarize your feedback at the end.

Here are some examples of good responses to reviews that may fall under each of these categories.

1. A question or claim that is either already answered in the paper or contradicts what is provided in the paper.

For such cases, we would like to point the reviewer specifically to the relevant part of the paper through direct quoting. Only send a quote if it verbatim exists in the paper or review.

Example 1:

**Reviewer Comment:** *..In Figure 4, the efficiency experiments have no results for Transformer models, which is a key limitation of the paper.I*

**Feedback to reviewer:** You may want to check Section 3, Figure 5 of the paper which has the Transformer results. See: “In Transformers, the proposed technique provides 25% relative improvement in wall-clock time (Figure 5)”.

Example 2: 
**Reviewer Comment:** *Figure 2. Are people forced to select a choice or could they select 'I don't know'? Did you monitor response times to see if the manipulated images required longer times for individuals to pass decisions? In Appendix A, you mention “the response times will also be released upon publication”, however, I do not see further information about this in the paper.*

As the reviewer already refers to the key part of the paper with quotes and asks a question that is not answered in the paper, we do not need to give feedback to this comment.

2. Lack of clarity or justification for criticism: The reviewer criticizes the paper without providing a concrete justification. This results in points that are not actionable or harder to respond to. 

For such cases, we would like to nudge the reviewer to justify their claim.

Example 1:
**Reviewer Comment:** *A key limitation is that this paper motivates from the TTT perspective, but no TTT experiments are performed and compared, e.g., comparing with (Sun et al., 2020).*

**Feedback to reviewer:** In general, a “TTT experiment” is not a broadly known terminology. It would be very helpful if you could describe what specific experiment you are referring to, and why it is important for this paper.

Example 2: 
**Reviewer Comment:** *W1: Many of the key pieces of this work have been considered before separately. I think some additional discussion of related work is needed.

W2: It appears that the linear mode connectivity results may be somewhat brittle.*

**Feedback to reviewer:** For W1, It would be very helpful to specify or cite what key pieces are considered in earlier works, how they relate to the current work, and specify what is missing. For W2, could you elaborate on why you see the results as brittle and what would be convincing evidence so that the feedback is more actionable?

Example 3: 
**Reviewer Comment:** *The data and model size are not large-scale, thus the paper will not be impactful.*

**Feedback to reviewer:** Please consider putting this claim in context. If possible, please discuss why you consider the experiments to be not large-scale using precedents, why it is important to have larger-scale experiments in this context, and how it would inform or improve the manuscript.

3. Novelty claims: The reviewer claims a lack of novelty without providing concrete justification.

For this case, we would like to nudge the reviewer to justify the claim, by prompting them to provide the most relevant references, the relationships, and specifying similarities or differences.

Example 1:
**Reviewer Comment:** *.. The paper’s novelty is limited considering the ICLR standards, as there are very close works [1,2,3]. *

**Feedback to reviewer:** Please consider the reasons for why the novelty is limited, and specify what ICLR standards are in this context. In particular, it would be very helpful to all parties involved if you give examples of the closest papers, their similarities, and differences with the methods or results in the current paper.

Example 2: 
**Reviewer Comment:** *.. DASHA is a mash-up between MARINA and existing distributed nonconvex optimization methods. Other than the fact that three variants of DASHA get rid of the uncompressed synchronization in MARINA, this reviewer could not pinpoint a difference between MARINA and DASHA. As such, the main novelty of this work seems to be in terms of theoretical analysis of MARINA when the uncompressed synchronization step is removed. The authors could have done a better job of clarifying where does this novelty lie in the analysis (e.g., pinpointing the key analytical approaches in the lemma that helped improve the analysis)*

As the novelty claim is already well-justified and an actionable question is asked, we do not need to give feedback to this review.

4. Personal attack, ad hominem: The review contains an attack on the authors. This can be about the personality, the knowledge, or the experience of the authors. 

For this case, we would like to kindly warn the reviewer about their comment. 

Example 1: 
**Reviewer Comment:** *.. The authors clearly do not live in the real world and do not care about people or downstream effects of their research. *

**Feedback to the reviewer:** We kindly suggest you revise this comment, as it includes remarks about the personalities or intents of the authors.

Your instructions: Go through all comments in the review. Check each comment using the above guidelines. Only address issues with the review that fail according to the above guidelines. If the review is in a format where it has a "summary" and "strengths" section, DO NOT address those sections. Only address weaknesses or questions.  Make sure you only address issues with the review.
Your feedback will be sent directly to reviewers. Follow the below format when sending your review:

- Comment: {{the comment of interest}}
- Feedback: {{write your short feedback}}

and send a list of comments and feedbacks. If you do not identify issues for a comment, do not add it to the list or send feedback. DO NOT mention that there is a checklist or guidelines.

If you cannot identify any issues, respond with "Thanks for your hard work!"
"""

critic_prompt = """
You are given feedback about a peer review of a machine learning paper submitted to a top-tier ML conference on OpenReview. The aim of the feedback is to make the review high-quality.
First, you will read the original review, then you will read the feedback, and then you will read the paper the review was written about.
You need to edit the feedback for correctness, removing anything that is incorrect and revising anything that needs to be fixed.
You should also edit the feedback for clarity, removing anything that may frustrate or confuse a reviewer.

The original feedback was written according to the following checklist. For each item, the agent determined if the review passed or failed. If it failed, the agent provided feedback accordingly. Here is the checklist:

1. The reviewer does not ask for anything that is OBVIOUSLY already present in the paper. You should be ABSOLUTELY certain the reviewer made an error, otherwise do NOT comment on their point.
If you find such a case, quote the specific statement from the review and provide the exact quote from the paper that directly addresses or contradicts it. 
Ensure that the quote from the paper specifically answers the claim or request made by the reviewer, and explain why the selected quote is relevant.
If the reviewer has a question, you do not need to respond. If you respond, point to a specific section or quote from the paper and ask “does this provide what you were looking for?”
2. The review does not make any vague, unjustifiable, or unsupported claims. 
If there are any vague comments in the review, ask the reviewer to be more specific.
3. If the reviewer claims the paper isn't novel, they should explain why and specify the paper(s) they believe this work is similar to. 
If they have not already done this, ask the reviewer to do so.
4. The reviewer does not make any personal attacks against the paper and/or author(s). 
For example, they call the work "incompetent" without justifying why. 
If they do, suggest the reviewer revise their feedback to remove those attacks.

The agent was told to “be concise, limiting your response to 2-3 sentences per checklist item. Do not repeat anything that the reviewer already included in their review, and do not summarize your feedback at the end.”

Here are some examples of good responses to reviews that the agent was given that may fall under each of these categories.

1. A question or claim that is either already answered in the paper or contradicts what is provided in the paper.
For such cases, we would like to point the reviewer specifically to the relevant part of the paper through direct quoting. Only send a quote if it verbatim exists in the paper or review.
Example 1:
**Reviewer Comment:** *..In Figure 4, the efficiency experiments have no results for Transformer models, which is a key limitation of the paper.I*
**Feedback to reviewer:** You may want to check Section 3, Figure 5 of the paper which has the Transformer results. See: “In Transformers, the proposed technique provides 25% relative improvement in wall-clock time (Figure 5)”.
Example 2: 
**Reviewer Comment:** *Figure 2. Are people forced to select a choice or could they select 'I don't know'? Did you monitor response times to see if the manipulated images required longer times for individuals to pass decisions? In Appendix A, you mention “the response times will also be released upon publication”, however, I do not see further information about this in the paper.*
As the reviewer already refers to the key part of the paper with quotes and asks a question that is not answered in the paper, we do not need to give feedback to this comment.

2. Lack of clarity or justification for criticism: The reviewer criticizes the paper without providing a concrete justification. This results in points that are not actionable or harder to respond to. 
For such cases, we would like to nudge the reviewer to justify their claim.
Example 1:
**Reviewer Comment:** *A key limitation is that this paper motivates from the TTT perspective, but no TTT experiments are performed and compared, e.g., comparing with (Sun et al., 2020).*
**Feedback to reviewer:** In general, a “TTT experiment” is not a broadly known terminology. It would be very helpful if you could describe what specific experiment you are referring to, and why it is important for this paper.
Example 2: 
**Reviewer Comment:** *W1: Many of the key pieces of this work have been considered before separately. I think some additional discussion of related work is needed.
W2: It appears that the linear mode connectivity results may be somewhat brittle.*
**Feedback to reviewer:** For W1, It would be very helpful to specify or cite what key pieces are considered in earlier works, how they relate to the current work, and specify what is missing. For W2, could you elaborate on why you see the results as brittle and what would be convincing evidence so that the feedback is more actionable?
Example 3: 
**Reviewer Comment:** *The data and model size are not large-scale, thus the paper will not be impactful.*
**Feedback to reviewer:** Please consider putting this claim in context. If possible, please discuss why you consider the experiments to be not large-scale using precedents, why it is important to have larger-scale experiments in this context, and how it would inform or improve the manuscript.

3. Novelty claims: The reviewer claims a lack of novelty without providing concrete justification.
For this case, we would like to nudge the reviewer to justify the claim, by prompting them to provide the most relevant references, the relationships, and specifying similarities or differences.
Example 1:
**Reviewer Comment:** *.. The paper’s novelty is limited considering the ICLR standards, as there are very close works [1,2,3]. *
**Feedback to reviewer:** Please consider the reasons for why the novelty is limited, and specify what ICLR standards are in this context. In particular, it would be very helpful to all parties involved if you give examples of the closest papers, their similarities, and differences with the methods or results in the current paper.
Example 2: 
**Reviewer Comment:** *.. DASHA is a mash-up between MARINA and existing distributed nonconvex optimization methods. Other than the fact that three variants of DASHA get rid of the uncompressed synchronization in MARINA, this reviewer could not pinpoint a difference between MARINA and DASHA. As such, the main novelty of this work seems to be in terms of theoretical analysis of MARINA when the uncompressed synchronization step is removed. The authors could have done a better job of clarifying where does this novelty lie in the analysis (e.g., pinpointing the key analytical approaches in the lemma that helped improve the analysis)*
As the novelty claim is already well-justified and an actionable question is asked, we do not need to give feedback to this review.

4. Personal attack, ad hominem: The review contains an attack on the authors. This can be about the personality, the knowledge, or the experience of the authors. 
For this case, we would like to kindly warn the reviewer about their comment. 
Example 1: 
**Reviewer Comment:** *.. The authors clearly do not live in the real world and do not care about people or downstream effects of their research. *
**Feedback to the reviewer:** We kindly suggest you revise this comment, as it includes remarks about the personalities or intents of the authors.


Your instructions: Do not change the format of the feedback, only edit for content clarity and correctness. DO NOT ADD any new points unless the previous feedback OBVIOUSLY missed something important. You need to check every quote and factual claim in the feedback and edit for correctness, as it is imperative all the feedback is correct. If the review is in a format where it has a "summary" and "strengths" section, DO NOT address those sections. Make sure you only address issues with the review.
Your feedback will be sent directly to reviewers. Follow the below format when sending your review:
- Comment: {{the comment of interest}}
- Feedback: {{write your short feedback}}
and send a list of comments and feedbacks. If you do not identify issues for a comment, do not add it to the list or send feedback. DO NOT mention that there is a checklist or guidelines, and do NOT mention that this is the edited feedback. Do not add a preamble, go right to the comments and feedback.
If you have no feedback, then respond with "Thanks for your hard work!"
"""

# Set page to wide mode
st.set_page_config(layout="wide")

# Sidebar for navigation
st.sidebar.title('Input review and PDF or OpenReview ID')

option = st.sidebar.radio(
    "Choose one input option:",
    ("Upload PDF with Review", "OpenReview paper ID"))

# # Add a text input box to the sidebar for user input
# user_input = st.sidebar.text_input("Enter OpenReview ID (e.g. xC8xh2RSs2)")

user_input = False

if option == "OpenReview paper ID":
    paper_id = st.text_input("Enter OpenReview ID (e.g. xC8xh2RSs2):")
    paper_id = paper_id.strip()
    reviewer_id = st.text_input("Enter reviewer ID (e.g. gNxe):")
    reviewer_id = reviewer_id.strip()
    if paper_id and reviewer_id:
        upload_file = False
        user_input = True
        submissions = load_ICLR_submissions()
        review, pdf, pdf_text = parse_openreview_id(submissions, paper_id, reviewer_id)
        # st.write(f"Review text: {pdf_text}")

elif option == "Upload PDF with Review":
    user_text = st.text_area("Enter review:")
    uploaded_file = st.file_uploader("Upload PDF", type="pdf")
    if user_text and uploaded_file:
        upload_file = True
        user_input = True
        review, pdf = user_text, uploaded_file
        pdf_text = pdf_to_text(pdf)

        # st.write(f"Review text: {pdf_text}")
        # st.write(f"You entered review: {user_text}")
        # st.write(f"Uploaded file: {uploaded_file.name}")

if user_input:
    # Run your pipeline to generate the dataframe based on user input
    feedback = create_feedback(review, pdf_text, agent_prompt)

    revised_feedback = critic(review, feedback, pdf_text, critic_prompt)

    st.title(f'Review feedback')

    # Create three columns
    col1, col2, col3 = st.columns(3)

    # Column 1: Display review fields
    with col1:
        st.subheader("Review")
        st.write(json.dumps(review))

        if not upload_file:
            pdf_url = f"https://openreview.net/pdf?id={paper_id}"
            st.markdown(f"[Click here to view PDF]({pdf_url})")

    # Column 2: Display feedback
    with col2:
        st.subheader('Feedback')
        st.write(feedback)

    # Column 3: Display revised feedback (from critic)
    with col3:
        st.subheader('Revised Feedback')
        st.write(revised_feedback)

else:
    st.title('Please enter OpenReview ID or upload PDF with review to generate feedback')