import requests
import json
import yaml
import scipdf
import os
import time
import aiohttp
import asyncio
import numpy as np
import random

def get_content_between_a_b(start_tag, end_tag, text):
    extracted_text = ""
    start_index = text.find(start_tag)
    while start_index != -1:
        end_index = text.find(end_tag, start_index + len(start_tag))
        if end_index != -1:
            extracted_text += text[start_index + len(start_tag) : end_index] + " "
            start_index = text.find(start_tag, end_index + len(end_tag))
        else:
            break
    return extracted_text.strip()


def extract(text, type):
    if text:
        target_str = get_content_between_a_b(f"<{type}>", f"</{type}>", text)
        if target_str:
            return target_str
        else:
            return text
    else:
        return ""
    
def download(url):
    try:
        response = requests.get(url)
        if response.status_code == 200:
            return response.content
        else:
            print(f"Failed to download the file from the URL: {url}")
            return None
    except requests.RequestException as e:
        print(f"An error occurred while downloading the file from the URL: {url}")
        print(e)
        return None
    except Exception as e:
        print(f"An unexpected error occurred while downloading the file from the URL: {url}")
        print(e)
        return None
    
class Result:
    def __init__(self,title="",abstract="",article = "",citations_conut = 0,year = None) -> None:
        self.title = title
        self.abstract = abstract
        self.article = article
        self.citations_conut = citations_conut
        self.year = year

# Define the API endpoint URL

semantic_fields = ["title", "abstract", "year", "authors.name", "authors.paperCount", "authors.citationCount","authors.hIndex","url","referenceCount","citationCount","influentialCitationCount","isOpenAccess","openAccessPdf","fieldsOfStudy","s2FieldsOfStudy","embedding.specter_v1","embedding.specter_v2","publicationDate","citations"]


fieldsOfStudy = ["Computer Science","Medicine","Chemistry","Biology","Materials Science","Physics","Geology","Art","History","Geography","Sociology","Business","Political Science","Philosophy","Art","Literature","Music","Economics","Philosophy","Mathematics","Engineering","Environmental Science","Agricultural and Food Sciences","Education","Law","Linguistics"]

# citations.paperId, citations.title, citations.year, citations.authors.name, citations.authors.paperCount, citations.authors.citationCount, citations.authors.hIndex, citations.url, citations.referenceCount, citations.citationCount, citations.influentialCitationCount, citations.isOpenAccess, citations.openAccessPdf, citations.fieldsOfStudy, citations.s2FieldsOfStudy, citations.publicationDate

# publicationDateOrYear： 2019-03-05 ； 2019-03 ； 2019 ； 2016-03-05:2020-06-06 ； 1981-08-25: ； :2020-06-06 ； 1981:2020

# publicationTypes: Review ； JournalArticle CaseReport ； ClinicalTrial ； Dataset ； Editorial ； LettersAndComments ； MetaAnalysis ； News ； Study ； Book ； BookSection



def process_fields(fields):
   return ",".join(fields)


class SementicSearcher:
    def __init__(self, ban_paper = []) -> None:
        self.ban_paper = ban_paper
    
    def search_papers(self, query, limit=5, offset=0, fields=["title", "paperId", "abstract", "isOpenAccess", 'openAccessPdf', "year","publicationDate","citations.title","citations.abstract","citations.isOpenAccess","citations.openAccessPdf","citations.citationCount","citationCount","citations.year"],
                            publicationDate=None, minCitationCount=0, year=None, 
                            publicationTypes=None, fieldsOfStudy=None):
        url = 'https://api.semanticscholar.org/graph/v1/paper/search'
        fields = process_fields(fields) if isinstance(fields, list) else fields
        
        # More specific query parameter
        query_params = {
            'query': query,
            "limit": limit,
            "offset": offset,
            'fields': fields,
            'publicationDateOrYear': publicationDate,
            'minCitationCount': minCitationCount,
            'year': year,
            'publicationTypes': publicationTypes,
            'fieldsOfStudy': fieldsOfStudy
        }
        # Load the API key from the configuration file 
        api_key = os.environ.get('SEMENTIC_SEARCH_API_KEY',None)
        headers = {'x-api-key': api_key} if api_key else None
        try:
            filtered_query_params = {key: value for key, value in query_params.items() if value is not None}
            response = requests.get(url, params=filtered_query_params, headers=headers)

            if response.status_code == 200:
                response_data = response.json()
                return response_data
            elif response.status_code == 429:
                time.sleep(1)  
                print(f"Request failed with status code {response.status_code}: begin to retry")
                return self.search_papers(query, limit, offset, fields, publicationDate, minCitationCount, year, publicationTypes, fieldsOfStudy)
            else:
                print(f"Request failed with status code {response.status_code}: {response.text}")
                return None
        except requests.RequestException as e:
            print(f"An error occurred: {e}")
            return None
                
    def cal_cosine_similarity(self, vec1, vec2):
        return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))
    
    def cal_cosine_similarity_matric(self,matric1, matric2):
        if isinstance(matric1, list):
            matric1 = np.array(matric1)
        if isinstance(matric2, list):
            matric2 = np.array(matric2)
        if len(matric1.shape) == 1:
            matric1 = matric1.reshape(1, -1)
        if len(matric2.shape) == 1:
            matric2 = matric2.reshape(1, -1)
        dot_product = np.dot(matric1, matric2.T)
        norm1 = np.linalg.norm(matric1, axis=1)
        norm2 = np.linalg.norm(matric2, axis=1)

        cos_sim = dot_product / np.outer(norm1, norm2)
        scores = cos_sim.flatten()
        return scores.tolist()

    def read_arxiv_from_path(self, pdf_path):
        def is_pdf(binary_data):
            pdf_header = b'%PDF-'
            return binary_data.startswith(pdf_header)
        try:
            flag = is_pdf(pdf_path)
            if not flag:
                return None
        except Exception as e:
            pass
        try:
            article_dict = scipdf.parse_pdf_to_dict(pdf_path)
        except Exception as e:
            print(f"Failed to parse the PDF file: {pdf_path}")
            return None
        return article_dict
        
    def get_paper_embbeding_and_score(self,query_embedding, paper,llm):
        paper_content = f"""
Title: {paper['title']}
Abstract: {paper['abstract']}
"""
        paper_embbeding = llm.get_embbeding(paper_content)
        paper_embbeding = np.array(paper_embbeding)
        score = self.cal_cosine_similarity(query_embedding,paper_embbeding)
        return [paper,score]

    
    def rerank_papers(self, query_embedding, paper_list,llm):
        if len(paper_list) == 0:
            return []
        paper_list = [paper for paper in paper_list if paper]
        if len(paper_list) >= 50:
            paper_list = random.sample(paper_list,50)
        paper_contents = []
        for paper in paper_list:
            paper_content = f"""
Title: {paper['title']}
Abstract: {paper['abstract']}
"""
            paper_contents.append(paper_content)
        paper_contents_embbeding = llm.get_embbeding(paper_contents)
        paper_contents_embbeding = np.array(paper_contents_embbeding)
        scores = self.cal_cosine_similarity_matric(query_embedding,paper_contents_embbeding)
            
        # 根据score对paper_list进行排序 
        paper_list = sorted(zip(paper_list,scores),key = lambda x: x[1],reverse = True)
        paper_list = [paper[0] for paper in paper_list]
        return paper_list
        
    
    def search(self,query,max_results = 5 ,paper_list = None ,rerank_query = None,llm = None,year = None,publicationDate = None,need_download = True,fields = ["title", "paperId", "abstract", "isOpenAccess", 'openAccessPdf', "year","publicationDate","citationCount"]):
        if rerank_query:
            rerank_query_embbeding = llm.get_embbeding(rerank_query)
            rerank_query_embbeding = np.array(rerank_query_embbeding)

        readed_papers = []
        if paper_list:
            if isinstance(paper_list,set):
                paper_list = list(paper_list)
            if len(paper_list) == 0 :
                pass
            elif isinstance(paper_list[0], str):
                readed_papers = paper_list
            elif isinstance(paper_list[0], Result):
                readed_papers = [paper.title for paper in paper_list]

        print(f"Searching for papers related to the query: <{query}>")
        results = self.search_papers(query,limit = 10 * max_results,year=year,publicationDate = publicationDate,fields = fields)
        if not results or "data" not in results:
            return []
        
        new_results = []
        for result in results['data']:
            if result['title'] in self.ban_paper:
                continue
            new_results.append(result)
        results = new_results

        final_results = []
        if need_download:
            paper_candidates = []
            for result in results:
                if not result['isOpenAccess'] or  not result['openAccessPdf'] or result['title'] in readed_papers:
                    continue
                else:
                    paper_candidates.append(result)
        else:
            paper_candidates = results
        
        if llm and rerank_query:
            paper_candidates = self.rerank_papers(rerank_query_embbeding, paper_candidates,llm)
        
        if need_download:
            for result in paper_candidates:
                pdf_link = result['openAccessPdf']["url"]
                try:
                    content = self.download_pdf(pdf_link)
                    if not content:
                        continue
                except Exception as e:
                    continue
                title = result['title']
                abstract = result['abstract']
                citationCount = result['citationCount']
                year = result['year']
                article = self.read_arxiv_from_path(content)
                if not article:
                    continue
                final_results.append(Result(title,abstract,article,citationCount,year))
                if len(final_results) >= max_results:
                    break
        else:
            for result in paper_candidates:
                title = result['title']
                abstract = result['abstract']
                citationCount = result['citationCount']
                year = result['year']
                final_results.append(Result(title,abstract,None,citationCount,year))
                if len(final_results) >= max_results:
                    break
        return final_results
    
    def search_related_paper(self,title,need_citation = True,need_reference = True,rerank_query = None,llm = None,paper_list = []):
        print(f"Searching for the related papers of <{title}>, need_citation: {need_citation}, need_reference: {need_reference}")
        fileds = ["title","abstract","citations.title","citations.abstract","citations.citationCount","references.title","references.abstract","references.citationCount","citations.isOpenAccess","citations.openAccessPdf","references.isOpenAccess","references.openAccessPdf","citations.year","references.year"]
        results = self.search_papers(title,limit = 3,fields=fileds)
        related_papers = []
        related_papers_title = []
        if not results or "data" not in results:
            return None
        for result in results["data"]:
            if not result:
                continue
            if need_citation:
                for citation in result["citations"]:
                    if "openAccessPdf" not in citation or not citation["openAccessPdf"]:
                        continue
                    elif citation["title"] in related_papers_title or citation["title"] in self.ban_paper or citation["title"] in paper_list:
                        continue
                    elif citation["isOpenAccess"] == False or citation["openAccessPdf"] == None:
                        continue
                    else:
                        related_papers.append(citation)
                        related_papers_title.append(citation["title"])
            if need_reference:
                for reference in result["references"]:
                    if "openAccessPdf" not in reference or not reference["openAccessPdf"]:
                        continue
                    elif reference["title"] in related_papers_title or reference["title"] in self.ban_paper or reference["title"] in paper_list:
                        continue
                    elif reference["isOpenAccess"] == False or reference["openAccessPdf"] == None:
                        continue
                    else:
                        related_papers.append(reference)
                        related_papers_title.append(reference["title"])
            if result:
                break
        if len(related_papers) >= 200:
            related_papers = related_papers[:200]

        if rerank_query and llm:
            rerank_query_embbeding = llm.get_embbeding(rerank_query)
            rerank_query_embbeding = np.array(rerank_query_embbeding)
            related_papers = self.rerank_papers(rerank_query_embbeding, related_papers,llm)
            related_papers = [[paper["title"],paper["abstract"],paper["openAccessPdf"]["url"],paper["citationCount"],paper['year']] for paper in related_papers]
        else:
            related_papers = [[paper["title"],paper["abstract"],paper["openAccessPdf"]["url"],paper["citationCount"],paper['year']] for paper in related_papers]
            related_papers = sorted(related_papers,key = lambda x: x[3],reverse = True)
        print(f"Found {len(related_papers)} related papers")
        for paper in related_papers:
            url = paper[2]
            content = self.download_pdf(url)
            if content:
                article = self.read_arxiv_from_path(content)
                if not article:
                    continue
                result = Result(paper[0],paper[1],article,paper[3],paper[4])
                return result
        return None
    
    
    def download_pdf(self, pdf_link):
        content = download(pdf_link)
        return content

    def read_paper_title_abstract(self,article):
        title = article["title"]
        abstract = article["abstract"]
        paper_content = f"""
Title: {title}
Abstract: {abstract}
        """
        return paper_content

    def read_paper_content(self,article):
        paper_content = self.read_paper_title_abstract(article)
        for section in article["sections"]:
            paper_content += f"section: {section['heading']}\n content: {section['text']}\n ref_ids: {section['publication_ref']}\n"
        return paper_content
    
    def read_paper_content_with_ref(self,article):
        paper_content = self.read_paper_content(article)
        paper_content += "<References>\n"
        i = 1
        for refer in article["references"]:
            ref_id = refer["ref_id"]
            title = refer["title"]
            year = refer["year"]
            paper_content += f"Ref_id:{ref_id} Title: {title} Year: ({year})\n"
            i += 1
        paper_content += "</References>\n"
        return paper_content