# -*- coding: utf-8 -*-
"""
Created on Thu Jan  9 12:15:14 2025

@author: acer
"""

import arxiv
import pandas as pd
from nltk.tokenize import sent_tokenize
from nltk.corpus import stopwords
from sklearn.feature_extraction.text import TfidfVectorizer
import nltk

from sklearn.feature_extraction.text import TfidfVectorizer
from nltk.corpus import stopwords
from nltk.tokenize import sent_tokenize

# Ensure required NLTK data is downloaded
nltk.download('punkt')
nltk.download('stopwords')

# Function to summarize text using NLTK
def summarize_text(text, sentence_count=3):
    sentences = sent_tokenize(text)
    
    # If the abstract is short, return the full abstract
    if len(sentences) <= sentence_count:
        return text  # No need to summarize if it's already short

    stop_words = stopwords.words('english')  # Keep this as a list
    vectorizer = TfidfVectorizer(stop_words=stop_words)  # Pass the list directly
    tfidf_matrix = vectorizer.fit_transform(sentences)
    sentence_scores = tfidf_matrix.sum(axis=1).A1  # Sum TF-IDF scores per sentence
    top_sentence_indices = sorted(
        range(len(sentence_scores)), key=lambda i: sentence_scores[i], reverse=True
    )[:sentence_count]

    summary = " ".join([sentences[i] for i in sorted(top_sentence_indices)])
    return summary

# Function to fetch papers from arXiv
def fetch_papers(query, max_results=10):
    search = arxiv.Search(
        query=query,
        max_results=max_results,
        sort_by=arxiv.SortCriterion.Relevance
    )
    papers = []
    for result in search.results():
        papers.append({
            "title": result.title,
            "authors": [author.name for author in result.authors],
            "abstract": result.summary,  # Capture the abstract
            "published_date": result.published,
            "url": result.entry_id
        })
    return papers

# Function to save results to CSV
def save_to_csv(papers, filename="papers3.csv"):
    df = pd.DataFrame(papers)
    df.to_csv(filename, index=False)
    print(f"Saved {len(papers)} papers to {filename}")

# Main Function
def main():
    query = "error correcting codes"
    print(f"Fetching papers on: {query}")
    
    # Step 1: Fetch Papers
    papers = fetch_papers(query, max_results=5)
    print(f"Fetched {len(papers)} papers.")

    # Step 2: Summarize Abstracts
    for paper in papers:
        print(f"Summarizing: {paper['title']}")
        paper['summary_short'] = summarize_text(paper['abstract'])  # Summarize the abstract

    # Step 3: Save to CSV
    save_to_csv(papers)

if __name__ == "__main__":
    main()