"""Class definition and Function to create citation file"""

import ads
class papers:
    
    def __init__(self, username, sort_by='year', authorship='first_author', filename='citation_list.txt'):
        self.username = username
        self.sort_by = sort_by
        self.authorship = authorship
        self.filename = filename
        
    def query_ADS_first_author(self):

        # NOTE: Wrong person papers are included because of name ambiguity

        self.paper_data = list(ads.SearchQuery(first_author=self.username, sort=self.sort_by, 
                               fl=['author', 'title', 'pub', 'year', 'volume', 'page', 'citation_count']))
        
    def query_ADS_nth_author(self):

        # NOTE: Wrong person papers are included because of name ambiguity

        self.paper_data = list(ads.SearchQuery(author=self.username, sort=self.sort_by, 
                               fl=['author', 'title', 'pub', 'year', 'volume', 'page', 'citation_count']))
    
    def generate_citation(self):

        # NOTE: Provide option for number of authors to be displayed -- print et al. after that number
        # NOTE: Formatting for nth author

        if(self.authorship == 'first_author'):
            self.query_ADS_first_author()
        else:
            # NOTE: First author papers need to be excluded
            self.query_ADS_nth_author()
            
        paper_citations = []
        paper_titles = []

        for index in range(len(self.paper_data)):
            
            flag = 1
            citation_string = ''
            paper_title = self.paper_data[index].title[0]
            
            if(self.paper_data[index].year!=None and self.paper_data[index].pub!=None and 
            self.paper_data[index].volume!=None and self.paper_data[index].page[0]!=None):
        
                for i, author_name in enumerate(self.paper_data[index].author):
                    split_name = author_name.split(',')

                    if(len(split_name)==1):
                        split_name.append(split_name[0][0]+' '+split_name[0][1])

                    split_last_name = split_name[1].split(' ')

                    if(len(split_last_name)==3):
                        citation_string += split_name[0]+', '+split_last_name[1][0]+'. '+split_last_name[2][0]+'., '
                    else:
                        citation_string += split_name[0]+', '+split_last_name[1][0]+'., '
                    
                    # For nth author
                    if(i==0):
                        if(self.authorship != 'first_author'):
                            # -2 to exclude space and comma
                            if self.username == citation_string[0:len(citation_string)-2]:
                                flag = 0
                                break
                if(flag):
                    citation_string += self.paper_data[index].year+', '+self.paper_data[index].pub+', '+self.paper_data[index].volume+', '+self.paper_data[index].page[0]
                    paper_citations.append(citation_string)
                    paper_titles.append(paper_title)

        self.paper_citations = paper_citations
        self.paper_titles = paper_titles
                
    def write_citation(self):
        cite_file = open(self.filename,"w")
        for index in range(len(self.paper_titles)):
            cite_file.write(self.paper_titles[index])
            cite_file.write('\n')
            cite_file.write(self.paper_citations[index])
            cite_file.write('\n')
            cite_file.write('\n')

        cite_file.close()

def create_citation_file(username, sort_by='year', authorship='first_author', filename='citation_list.txt'):
    """Citation File Creation

    Generate .txt file with paper titles and citations for given author name.

    Args:
        username (str): Name of author in 'Last Name, First Name/First Initial' format. E.g. 'Sanghi, Aniket' or 'Shah, N.'
        sort_by (str): Sort the papers queried by either 'year' or 'citation_count'
        authorship (str): Select either 'first_author' papers or 'nth_author' papers. 
        filename (str): Filename for the .txt file with paper titles and citations. E.g. 'citation_list.txt'
        
    Returns:
        citation file: .txt file with paper titles and citations.
    """
    author_obj = papers(username=username, sort_by=sort_by, authorship=authorship, filename=filename)
    author_obj.generate_citation()
    author_obj.write_citation()

