import json
import os

import faiss
import gradio as gr
import pandas as pd
import spaces
import torch
from datasets import load_dataset
from huggingface_hub import InferenceClient, hf_hub_download
from huggingface_hub import login as hf_hub_login
from huggingface_hub import upload_file
from sentence_transformers import SentenceTransformer

from arxiv_stuff import ARXIV_CATEGORIES_FLAT

# Get HF_TOKEN from environment variables
HF_TOKEN = os.getenv("HF_TOKEN")

# Login to Hugging Face Hub
hf_hub_login(token=HF_TOKEN, add_to_git_credential=True)

# Check if using persistent storage
persistent_storage = os.path.exists("/data")
if persistent_storage:
    # Use persistent storage
    print("Using persistent storage")

# Dataset details
dataset_name = "nomadicsynth/arxiv-dataset-abstract-embeddings"
dataset_revision = "v1.0.0"
local_index_path = "arxiv_faiss_index.faiss"

# Embedding model details
embedding_model_name = "nomadicsynth/research-compass-arxiv-abstracts-embedding-model"
embedding_model_revision = "2025-01-28_23-06-17-1epochs-12batch-32eval-512embed-final"

# Amalysis model details

# Settings for Llama-3.3-70B-Instruct
reasoning_model_id = "meta-llama/Llama-3.3-70B-Instruct"
max_length = 1024 * 4
temperature = None
top_p = None
presence_penalty = None

# Settings for QwQ-32B
# reasoning_model_id = "Qwen/QwQ-32B"
# reasoning_start_tag = "<think>"
# reasoning_end_tag = "</think>"
# max_length = 1024 * 4
# temperature = 0.6
# top_p = 0.95
# presence_penalty = 0.1

# Global variables
dataset = None
embedding_model = None
reasoning_model = None


def save_faiss_index_to_hub():
    """Save the FAISS index to the Hub for easy access"""
    global dataset, local_index_path
    # 1. Save the index to a local file
    dataset["train"].save_faiss_index("embedding", local_index_path)
    print(f"FAISS index saved locally to {local_index_path}")

    # 2. Upload the index file to the Hub
    remote_path = upload_file(
        path_or_fileobj=local_index_path,
        path_in_repo=local_index_path,  # Same name on the Hub
        repo_id=dataset_name,  # Use your dataset repo
        token=HF_TOKEN,
        repo_type="dataset",  # This is a dataset file
        revision=dataset_revision,  # Use the same revision as the dataset
        commit_message="Add FAISS index",  # Commit message
    )

    print(f"FAISS index uploaded to Hub at {remote_path}")

    # Remove the local file. It's now stored on the Hub.
    os.remove(local_index_path)


def setup_dataset():
    """Load dataset with FAISS index"""
    global dataset
    print("Loading dataset from Hugging Face...")

    # Load dataset
    dataset = load_dataset(
        dataset_name,
        revision=dataset_revision,
    )

    # Try to load the index from the Hub
    try:
        print("Downloading pre-built FAISS index...")
        index_path = hf_hub_download(
            repo_id=dataset_name,
            filename="arxiv_faiss_index.faiss",
            revision=dataset_revision,
            token=HF_TOKEN,
            repo_type="dataset",
        )

        print("Loading pre-built FAISS index...")
        dataset["train"].load_faiss_index("embedding", index_path)
        print("Pre-built FAISS index loaded successfully")

    except Exception as e:
        print(f"Could not load pre-built index: {e}")
        print("Building new FAISS index...")

        # Add FAISS index if it doesn't exist
        if not dataset["train"].features.get("embedding"):
            print("Dataset doesn't have 'embedding' column, cannot create FAISS index")
            raise ValueError("Dataset doesn't have 'embedding' column")

        dataset["train"].add_faiss_index(
            column="embedding",
            metric_type=faiss.METRIC_INNER_PRODUCT,
            string_factory="HNSW,RFlat",  # Using reranking
        )

        # Save the FAISS index to the Hub
        save_faiss_index_to_hub()

    print(f"Dataset loaded with {len(dataset['train'])} items and FAISS index ready")


def init_embedding_model(model_name_or_path: str, model_revision: str = None) -> SentenceTransformer:
    global embedding_model

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    embedding_model = SentenceTransformer(
        model_name_or_path,
        revision=model_revision,
        token=HF_TOKEN,
        device=device,
    )


def init_reasoning_model(model_name: str) -> InferenceClient:
    global reasoning_model
    reasoning_model = InferenceClient(
        model=model_name,
        provider="hf-inference",
        api_key=HF_TOKEN,
    )
    return reasoning_model


def generate(messages: list[dict[str, str]]) -> str:
    """
    Generate a response to a list of messages.

    Args:
        messages: A list of message dictionaries with a "role" and "content" key.

    Returns:
        The generated response as a string.
    """
    global reasoning_model

    system_message = {
        "role": "system",
        "content": "You are an expert in evaluating connections between research papers.",
    }

    messages.insert(0, system_message)

    response_schema = r"""{
  "$schema": "http://json-schema.org/draft-07/schema#",
  "title": "Generated schema for Root",
  "type": "object",
  "properties": {
    "reasoning": {
      "type": "string"
    },
    "key_connections": {
      "type": "array",
      "items": {
        "type": "object",
        "properties": {
          "connection": {
            "type": "string"
          },
          "description": {
            "type": "string"
          }
        },
        "required": [
          "connection",
          "description"
        ]
      }
    },
    "synergies_and_complementarities": {
      "type": "array",
      "items": {
        "type": "object",
        "properties": {
          "type": {
            "type": "array",
            "items": {
              "type": "string"
            }
          },
          "description": {
            "type": "string"
          }
        },
        "required": [
          "type",
          "description"
        ]
      }
    },
    "research_potential": {
      "type": "array",
      "items": {
        "type": "object",
        "properties": {
          "potential": {
            "type": "string"
          },
          "description": {
            "type": "string"
          }
        },
        "required": [
          "potential",
          "description"
        ]
      }
    },
    "rating": {
      "type": "number"
    },
    "confidence": {
      "type": "number"
    }
  },
  "required": [
    "reasoning",
    "key_connections",
    "synergies_and_complementarities",
    "research_potential",
    "rating",
    "confidence"
  ]
}"""

    response_format = {
        "type": "json",
        "value": response_schema,
    }

    result = reasoning_model.chat.completions.create(
        messages=messages,
        max_tokens=max_length,
        temperature=temperature,
        presence_penalty=presence_penalty,
        response_format=response_format,
        top_p=top_p,
    )

    output = result.choices[0].message.content.strip()
    return output


@spaces.GPU
def embed_text(text: str | list[str]) -> torch.Tensor:
    global embedding_model

    # Strip any leading/trailing whitespace
    text = text.strip() if isinstance(text, str) else [t.strip() for t in text]
    embed_text = embedding_model.encode(text, normalize_embeddings=True)  # Ensure vectors are normalized
    return embed_text


def analyse_abstracts(query_abstract: str, compare_abstract: dict) -> str:
    """Analyze the relationship between two abstracts and return formatted analysis"""
    global reasoning_model
    # Check if the compare_abstract is valid
    if not isinstance(compare_abstract, dict) or "abstract" not in compare_abstract:
        return "Invalid compare_abstract format. Expected a dictionary with 'abstract' key."
    if not query_abstract or not compare_abstract["abstract"]:
        return "Invalid input. Please provide both query_abstract and compare_abstract."
    # Check if the query_abstract is a string
    if not isinstance(query_abstract, str):
        return "Invalid query_abstract format. Expected a string."
    # Check if the compare_abstract is a string
    if not isinstance(compare_abstract["abstract"], str):
        return "Invalid compare_abstract format. Expected a string."
    # Check if the query_abstract is empty
    if not query_abstract.strip():
        return "Invalid query_abstract format. Expected a non-empty string."
    # Check if the compare_abstract is empty
    if not compare_abstract["abstract"].strip():
        return "Invalid compare_abstract format. Expected a non-empty string."

    messages = [
    {
        "role": "user",
        "content": f"""You are trained in evaluating conceptual and methodological connections between research papers. Please **identify and analyze the reasoning-based links** between the following two papers:

Paper 1 Abstract:
{query_abstract}

Paper 2 Abstract:
{compare_abstract["abstract"]}

In your evaluation, consider the following dimensions:

* **Methodological Cross-Pollination**: Do the methods or approaches from one paper **directly inform, enhance, or contrast with** the other?
* **Principle or Mechanism Extension**: Do the papers **share core principles, mechanisms, or assumptions** that could be **combined or extended** to generate new understanding or tools?
* **Interdisciplinary Bridges**: Are there clear opportunities for **knowledge transfer or collaboration** across fields or problem domains?
* **Solution or Application Overlap**: Can the solutions, frameworks, or applications in one paper be **adapted or repurposed** to benefit the work in the other, leading to **tangible, novel outcomes**?

Assess these connections in both directions (Paper 1 → Paper 2 and Paper 2 → Paper 1). Focus on **relevant and practically meaningful links** — especially those that might be **missed in practice** due to the sheer volume of publications or the separation between research communities. These are often connections that would be **immediately apparent to an expert** familiar with both papers, but easily overlooked otherwise.

Return a valid JSON object in the following structure:
{{
    "reasoning": "Step-by-step conceptual analysis of how the papers relate, highlighting **key connections**, complementary methods, or shared ideas. Emphasize the most **relevant, practically useful takeaways**, and use markdown bold to highlight major points.",

    "key_connections": [
        {{
            "connection": "connection 1",
            "description": "1–2 sentence explanation of the **main conceptual or methodological link**, emphasizing its practical or theoretical relevance."
        }},
        ...
    ],

    "complementarities": [
        {{
            "type": ["Methodological Cross-Pollination", "Principle or Mechanism Extension", "Interdisciplinary Bridges", "Solution or Application Overlap"],  # Use only the most relevant label per entry
            "description": "A concise explanation (1–2 sentences) of the **identified complementarity** or **productive relationship**, including a specific example or outcome it could enable."
        }},
        ...
    ],

    "research_potential": [
        {{
            "potential": "Potential application or outcome 1",
            "description": "1–2 sentence explanation of the **concrete potential impact**, framed in terms of a **realistic scenario or use case**."
        }},
        ...
    ],

    "rating": 1-5,  # Overall strength of the connection:
                    # 1 = No meaningful connection
                    # 2 = Weak or speculative connection
                    # 3 = Plausible but unproven connection
                    # 4 = Solid connection with future potential
                    # 5 = Strong, well-aligned connection with immediate, valuable implications

    "confidence": 0.0-1.0  # Confidence score in your assessment (e.g., 0.85 for high confidence, 1.0 for absolute certainty)
    # Note: The confidence score should reflect your level of certainty in the analysis, not the strength of the connection itself.
    # A score of 0.0 indicates no confidence in the analysis, while 1.0 indicates absolute certainty.
}}

Return only the JSON object. All key names and string values must be in double quotes.
""",
    },
]

    # Generate analysis
    try:
        output = generate(messages)
    except Exception as e:
        return f"Error: {e}"

    # Parse the JSON output
    try:
        output = json.loads(output)
    except Exception as e:
        return f"Error: {e}"

    # Format the output as markdown for better display
    key_connections = ""
    synergies_and_complementarities = ""
    research_potential = ""
    if "key_connections" in output:
        for connection in output["key_connections"]:
            key_connections += f"- {connection['connection']}: {connection['description']}\n"

    if "synergies_and_complementarities" in output:
        for synergy in output["synergies_and_complementarities"]:
            synergies_and_complementarities += f"- {', '.join(synergy['type'])}: {synergy['description']}\n"

    if "research_potential" in output:
        for potential in output["research_potential"]:
            research_potential += f"- {potential['potential']}: {potential['description']}\n"

    formatted_output = f"""## Synergy Analysis

**Rating**: {'★' * output['rating']}{'☆' * (5-output['rating'])}        **Confidence**: {'★' * round(output['confidence'] * 5)}{'☆' * round((1-output['confidence']) * 5)}

### Key Connections
{key_connections}

### Synergies and Complementarities
{synergies_and_complementarities}

### Research Potential
{research_potential}

### Reasoning
{output['reasoning']}
"""
    return formatted_output
    # return '```"""\n' + output + '\n"""```'


# arXiv Embedding Dataset Details
# DatasetDict({
#     train: Dataset({
#         features: ['id', 'submitter', 'authors', 'title', 'comments', 'journal-ref', 'doi', 'report-no', 'categories', 'license', 'abstract', 'update_date', 'embedding', 'timestamp', 'embedding_model'],
#         num_rows: 2689088
#     })
# })


def find_synergistic_papers(abstract: str, limit=25) -> list[dict]:
    """Find papers synergistic with the given abstract using FAISS with cosine similarity"""
    global dataset

    # Generate embedding for the query abstract (normalized for cosine similarity)
    abstract_embedding = embed_text(abstract)

    # Search for similar papers using FAISS with inner product (cosine similarity for normalized vectors)
    scores, examples = dataset["train"].get_nearest_examples("embedding", abstract_embedding, k=limit)

    papers = []
    for i in range(len(scores)):
        # With cosine similarity, higher scores are better (closer to 1)
        paper_dict = {
            "id": examples["id"][i],
            "title": examples["title"][i],
            "authors": examples["authors"][i],
            "categories": examples["categories"][i],
            "abstract": examples["abstract"][i],
            "update_date": examples["update_date"][i],
            "synergy_score": float(scores[i]),  # Convert to float for serialization
        }
        papers.append(paper_dict)

    return papers


def format_search_results(abstract: str) -> tuple[pd.DataFrame, list[dict]]:
    """Format search results as a DataFrame for display"""
    # Find papers synergistic with the given abstract
    papers = find_synergistic_papers(abstract)

    # Convert to DataFrame for display
    df = pd.DataFrame(
        [
            {
                "Title": p["title"],
                "Authors": p["authors"][:50] + "..." if len(p["authors"]) > 50 else p["authors"],
                "Categories": p["categories"],
                "Date": p["update_date"],
                "Match Score": f"{int(p['synergy_score'] * 100)}%",
                "ID": p["id"],  # Hidden column for reference
            }
            for p in papers
        ]
    )

    return df, papers  # Return both DataFrame and original data


def format_paper_as_markdown(paper: dict) -> str:
    # Convert category codes to full names, handling unknown categories
    subjects = []
    for subject in paper["categories"].split():
        if subject in ARXIV_CATEGORIES_FLAT:
            subjects.append(ARXIV_CATEGORIES_FLAT[subject])
        else:
            subjects.append(f"Unknown Category ({subject})")

    paper["title"] = paper["title"].replace("\n", " ").strip()
    paper["authors"] = paper["authors"].replace("\n", " ").strip()

    return f"""# {paper["title"]}
### {paper["authors"]}
#### {', '.join(subjects)} | {paper["update_date"]} | **Score**: {int(paper['synergy_score'] * 100)}%
**[arxiv:{paper["id"]}](https://arxiv.org/abs/{paper["id"]})** - [PDF](https://arxiv.org/pdf/{paper["id"]})<br>

{paper["abstract"]}
"""


latex_delimiters = [
    {"left": "$$", "right": "$$", "display": True},
    # {"left": "$", "right": "$", "display": False},
    # {"left": "\\(", "right": "\\)", "display": False},
    # {"left": "\\begin{equation}", "right": "\\end{equation}", "display": True},
    # {"left": "\\begin{align}", "right": "\\end{align}", "display": True},
    # {"left": "\\begin{alignat}", "right": "\\end{alignat}", "display": True},
    # {"left": "\\begin{gather}", "right": "\\end{gather}", "display": True},
    # {"left": "\\begin{CD}", "right": "\\end{CD}", "display": True},
    # {"left": "\\[", "right": "\\]", "display": True},
    # {"left": "\\underline{", "right": "}", "display": False},
    # {"left": "\\textit{", "right": "}", "display": False},
    # {"left": "\\textit{", "right": "}", "display": False},
    # {"left": "{", "right": "}", "display": False},
]


def create_interface():
    # Create CSV loggers
    analysis_logger = gr.CSVLogger()
    paper_match_logger = gr.CSVLogger()

    with gr.Blocks(
        css="""
    .cell-menu-button {
        display: none;
    }"""
    ) as demo:
        gr.HTML(
            """
            <div style="text-align: center; margin-bottom: 1rem">
                <h1>Inkling</h1>
                <p>Discover papers with deep conceptual connections to your research</p>
                <p>An experiment in AI-assisted research discovery and insight generation</p>
            </div>
        """
        )

        with gr.Accordion(label="Instructions and Privacy Policy", open=False):
            gr.Markdown(
                """
                This tool helps you uncover research papers with **deep, meaningful connections** to your ideas.
                It uses AI to go beyond keyword or semantic similarity — analyzing how papers relate **conceptually** and **contextually**, 
                even when the surface topics differ.

                The focus is on surfacing *novel insights* — connections that may not be obvious at a glance, 
                but could **spark new perspectives**, **deepen understanding**, or **highlight relationships that might otherwise be overlooked**.

                It’s designed to act more like a research collaborator than a search engine — helping you explore conceptual bridges and 
                unexpected pathways in the literature.

                Please ask any questions or provide feedback on the tool to help us improve it by starting a discussion on
                the [Community Tab](https://huggingface.co/spaces/nomadicsynth/inkling/discussions).

                **Privacy Policy**: The abstract or research description you provide will be included in any feedback
                you submit and may be used to improve the model, and published in a public dataset.
                Please ensure that you have the right to share this information.
                By submitting feedback, you agree to the use of this information for research purposes.
                Do not include personally identifiable, proprietary, or sensitive information.

                **Disclaimer**: This tool is in alpha testing and is not intended for production use.
                The results are not guaranteed to be accurate or reliable. Use at your own risk.
                The tool is provided "as is" without any warranties or guarantees.
                The developers are not responsible for any consequences of using this tool.
                By using this tool, you agree to the terms and conditions outlined in this disclaimer.
                """
            )
            gr.Markdown(
                """
                1. **Enter Abstract**: Paste an abstract or describe your research question or idea in the text box.
                2. **Find Related Papers**: Click the button to explore conceptually related research.
                3. **Select a Paper**: Click on a row in the results table to view more details.
                4. **Analyze Connection**: Click the analysis button to explore the potential connection between the papers.
                5. **Insight Analysis**: Review the model’s reasoning about how and why these papers may relate meaningfully.
                """
            )

        abstract_input = gr.Textbox(
            label="Paper Abstract or Description",
            placeholder="Paste an abstract or describe research details...",
            lines=8,
            key="abstract",
        )
        search_btn = gr.Button("Find Related Papers", variant="primary")

        # Store full paper data
        paper_data_state = gr.State([])

        # Store query abstract
        query_abstract_state = gr.State("")

        # Store selected paper
        selected_paper_state = gr.State(None)

        # Use Dataframe for results
        results_df = gr.Dataframe(
            headers=["Title", "Authors", "Categories", "Date", "Match Score"],
            datatype=["markdown", "markdown", "str", "date", "str"],
            latex_delimiters=latex_delimiters,
            label="Related Papers",
            interactive=False,
            wrap=False,
            line_breaks=False,
            column_widths=["40%", "20%", "20%", "10%", "10%", "0%"],  # Hide ID column
            key="results",
        )

        with gr.Row():
            with gr.Column(scale=1):
                paper_details_output = gr.Markdown(
                    value="# Paper Details",
                    label="Paper Details",
                    latex_delimiters=latex_delimiters,
                    show_copy_button=True,
                    key="paper_details",
                )
                analyze_btn = gr.Button("Analyze Connection", variant="primary", visible=False)
                with gr.Accordion(label="Feedback and Flagging", open=True, visible=False) as paper_feedback_accordion:
                    gr.Markdown(
                        """
                        Please provide feedback on the relevance of this paper to your input.
                        This helps us improve how well the system identifies meaningful research connections.
                        """
                    )
                    paper_feedback = gr.Radio(
                        ["👍 Good Match", "👎 Poor Match"],
                        label="Is this paper meaningfully related to your query?",
                    )
                    paper_expert = gr.Checkbox(label="I am an expert in this field", value=False)
                    paper_comment = gr.Textbox(label="Additional feedback on this match (optional)")
                    flag_paper_btn = gr.Button("Submit Paper Feedback")

            with gr.Column(scale=1):
                analysis_output = gr.Markdown(
                    value="# Connection Analysis",
                    label="Connection Analysis",
                    latex_delimiters=latex_delimiters,
                    show_copy_button=True,
                    key="analysis_output",
                )
                with gr.Accordion(label="Feedback and Flagging", open=True, visible=False) as analysis_feedback_accordion:
                    gr.Markdown(
                        """
                        This connection analysis was generated by an AI model trained to reason about conceptual links between research papers.
                        If you find the explanation helpful, unclear, or off-base, your feedback will help refine the model’s reasoning process.
                        """
                    )
                    analysis_feedback = gr.Radio(
                        ["👍 Helpful", "👎 Not Helpful"],
                        label="Was this explanation useful in understanding the connection?",
                    )
                    analysis_expert = gr.Checkbox(label="I am an expert in this field", value=False)
                    analysis_comment = gr.Textbox(label="Additional feedback on the analysis (optional)")
                    flag_analysis_btn = gr.Button("Submit Analysis Feedback")

        # Set up logging directories
        data_path = "/data" if persistent_storage else "./data"
        os.makedirs(data_path + "/flagged_paper_matches", exist_ok=True)
        os.makedirs(data_path + "/flagged_analyses", exist_ok=True)

        # Set up loggers
        paper_match_logger.setup(
            [abstract_input, paper_details_output, paper_feedback, paper_expert, paper_comment],
            data_path + "/flagged_paper_matches",
        )
        analysis_logger.setup(
            [abstract_input, paper_details_output, analysis_output, analysis_feedback, analysis_expert, analysis_comment],
            data_path + "/flagged_analyses",
        )

        # Display paper details when row is selected
        def on_select(evt: gr.SelectData, papers, query):
            selected_index = evt.index[0]  # Get the row index
            selected = papers[selected_index]

            # Format paper details
            details_md = format_paper_as_markdown(selected)

            return details_md, selected

        # Connect search button to the search function
        search_btn.click(
            format_search_results,
            inputs=[abstract_input],
            outputs=[results_df, paper_data_state],
            api_name="search",
        ).then(
            lambda x: x,  # Identity function to pass through the abstract
            inputs=[abstract_input],
            outputs=[query_abstract_state],
            api_name=False,
        ).then(
            lambda: None,  # Reset selected paper
            outputs=[selected_paper_state],
            api_name=False,
        ).then(
            lambda: (gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)),  # Hide analyze button and feedback accordions
            outputs=[analyze_btn, paper_feedback_accordion, analysis_feedback_accordion],
            api_name=False,
        ).then(
            lambda: ("# Paper Details", "# Synergy Analysis"),  # Clear previous outputs
            outputs=[paper_details_output, analysis_output],
            api_name=False,
        )

        # Use built-in select event from Dataframe
        results_df.select(
            on_select,
            inputs=[paper_data_state, query_abstract_state],
            outputs=[paper_details_output, selected_paper_state],
            api_name=False,
        ).then(
            lambda: (gr.update(visible=True), gr.update(visible=True)),  # Show analyze button and feedback accordion
            outputs=[analyze_btn, paper_feedback_accordion],
            api_name=False,
        )

        # Connect analyze button to run analysis
        analyze_btn.click(
            analyse_abstracts,
            inputs=[query_abstract_state, selected_paper_state],
            outputs=[analysis_output],
            show_progress_on=[paper_details_output, analysis_output],
            api_name=False,
        ).then(
            lambda: gr.update(visible=True),  # Show feedback accordion
            outputs=[analysis_feedback_accordion],
            api_name=False,
        )

        # Add flagging handlers
        flag_paper_btn.click(
            lambda *args: paper_match_logger.flag(list(args)),
            inputs=[abstract_input, paper_details_output, paper_feedback, paper_expert, paper_comment],
            preprocess=False,
            api_name=False,
        )

        flag_analysis_btn.click(
            lambda *args: analysis_logger.flag(list(args)),
            inputs=[abstract_input, paper_details_output, analysis_output, analysis_feedback, analysis_expert, analysis_comment],
            preprocess=False,
            api_name=False,
        )

    return demo


if __name__ == "__main__":
    # Load dataset with FAISS index
    setup_dataset()

    # Initialize the embedding model
    init_embedding_model(embedding_model_name, embedding_model_revision)

    # Initialize the reasoning model
    reasoning_model = init_reasoning_model(reasoning_model_id)

    demo = create_interface()
    demo.queue(api_open=False).launch(ssr_mode=False, show_api=True)