"""
Queries arxiv API and downloads papers (the query is a parameter).
The script is intended to enrich an existing database pickle (by default db.p),
so this file will be loaded first, and then new results will be added to it.
"""
# import app
from config import Config as config
import os
import time
import pickle
import random
import argparse
import urllib.request
import feedparser
import requests
import sqlite3
from utils import safe_pickle_dump


# import app

def chunks(l, n):
    # For item i in a range that is a length of l,
    for i in range(0, len(l), n):
        # Create an index range for l of n items:
        yield l[i:i+n]

def get_api_data(batch):
    '''Retrieve API data for a particular preprint id'''
    l = ','.join(batch)
    url = r'http://export.arxiv.org/api/query?id_list={}&max_results={}'.format(l,len(batch))
    data = requests.get(url).text
    js = feedparser.parse(data)
    return js


def arx_data_from_pids(pids, batch_size=100):
    '''
    Retrieves data in batches of 100 from ArXiv.
    '''
    arx_data = []
    k=0
    for batch_pids in chunks(pids,batch_size):
        # sleep 3s between api calls
        time.sleep(3)
        api_output = get_api_data(batch_pids)
        api_output = api_output['entries']

        # i=0
        for item in api_output:
            arx_data.append(item)
            # i+=1
            k+=1
        print(k, ' items processed')
    return arx_data

def encode_feedparser_dict(d):
    """
    helper function to get rid of feedparser bs with a deep copy.
    I hate when libs wrap simple things in their own classes.
    """
    if isinstance(d, feedparser.FeedParserDict) or isinstance(d, dict):
        j = {}
        for k in d.keys():
            j[k] = encode_feedparser_dict(d[k])
        return j
    elif isinstance(d, list):
        l = []
        for k in d:
            l.append(encode_feedparser_dict(k))
        return l
    else:
        return d

def parse_arxiv_url(url):
    """
    examples is http://arxiv.org/abs/1512.08756v2
    we want to extract the raw id and the version
    """
    ix = url.rfind('/')
    idversion = url[ix+1:] # extract just the id (and the version)
    parts = idversion.split('v')
    assert len(parts) == 2, 'error parsing url ' + url
    return parts[0], int(parts[1])

if __name__ == "__main__":

    # parse input arguments
    parser = argparse.ArgumentParser()
    parser.add_argument('--search-query', type=str,
                      default='cat:gr-qc',
                      help='query used for arxiv API. See http://arxiv.org/help/api/user-manual#detailed_examples')
    parser.add_argument('--start-index', type=int, default=0, help='0 = most recent API result')
    parser.add_argument('--max-index', type=int, default=500, help='upper bound on paper index we will fetch')
    parser.add_argument('--results-per-iteration', type=int, default=100, help='passed to arxiv API')
    parser.add_argument('--wait-time', type=float, default=5.0, help='lets be gentle to arxiv API (in number of seconds)')
    parser.add_argument('--break-on-no-added', type=int, default=200, help='break out early if all returned query papers are already in db? 1=yes, 0=no')
    args = parser.parse_args()

    # misc hardcoded variables
    base_url = 'http://export.arxiv.org/api/query?' # base api query url
    print('Searching arXiv for %s' % (args.search_query, ))

    # lets load the existing database to memory
    try:
        db = pickle.load(open(config.DB_PATH, 'rb'))
    except Exception as e:
        print('error loading existing database:')
        print(e)
        print('starting from an empty database')
        db = {}
        # time.sleep(5)  # useful in debugging

    # -----------------------------------------------------------------------------
    # main loop where we fetch the new results
    print('database has %d entries at start' % (len(db), ))
    num_added_total = 0


    # -----------------------------------------------------------------------------
    # -----------------------------------------------------------------------------

    # AD Edit
    # add preprints from users' publication histories
    if not os.path.isfile(config.DATABASE_PATH):
        print("the database file as.db should exist. You can create an empty database with sqlite3 as.db < schema.sql")
        sys.exit()

    sqldb = sqlite3.connect(config.DATABASE_PATH)
    sqldb.row_factory = sqlite3.Row # to return dicts rather than tuples

    def query_db(query, args=(), one=False):
        """Queries the database and returns a list of dictionaries."""
        cur = sqldb.execute(query, args)
        rv = cur.fetchall()
        return (rv[0] if rv else None) if one else rv

    publications = query_db('''select * from publication''')
    pids = []
    for publication in publications:
        pid = publication['paper_id']
        pids.append(pid)
    pids = list(set(pids))
    # print('pids', pids)

    parse = arx_data_from_pids(pids=pids,batch_size=100)
    # print('PARSE', parse)
    num_added = 0
    num_skipped = 0
    for e in parse:

        j = encode_feedparser_dict(e)
        # print('j', j)
        # extract just the raw arxiv id and version for this paper
        rawid, version = parse_arxiv_url(j['id'])
        j['_rawid'] = rawid
        j['_version'] = version

        # add to our database if we didn't have it before, or if this is a new version
        if not rawid in db or j['_version'] > db[rawid]['_version']:
            db[rawid] = j
            print('Updated %s added %s' % (j['updated'].encode('utf-8'), j['title'].encode('utf-8')))
            num_added += 1
            num_added_total += 1
        else:
            num_skipped += 1


    # -----------------------------------------------------------------------------
    # -----------------------------------------------------------------------------

    for i in range(args.start_index, args.max_index, args.results_per_iteration):

        print("Results %i - %i" % (i,i+args.results_per_iteration))
        query = 'search_query=%s&sortBy=lastUpdatedDate&start=%i&max_results=%i' % (args.search_query,
                                                               i, args.results_per_iteration)
        with urllib.request.urlopen(base_url+query) as url:
            response = url.read()
        parse = feedparser.parse(response)
        num_added = 0
        num_skipped = 0
        for e in parse.entries:

            j = encode_feedparser_dict(e)

            # extract just the raw arxiv id and version for this paper
            rawid, version = parse_arxiv_url(j['id'])
            j['_rawid'] = rawid
            j['_version'] = version

            # add to our database if we didn't have it before, or if this is a new version
            if not rawid in db or j['_version'] > db[rawid]['_version']:
                db[rawid] = j
                print('Updated %s added %s' % (j['updated'].encode('utf-8'), j['title'].encode('utf-8')))
                num_added += 1
                num_added_total += 1
            else:
                num_skipped += 1

        # print some information
        print('Added %d papers, already had %d.' % (num_added, num_skipped))

        if len(parse.entries) == 0:
            print('Received no results from arxiv. Rate limiting? Exiting. Restart later maybe.')
            print(response)
            break

        if num_added == 0 and args.break_on_no_added == 1:
            print('No new papers were added. Assuming no new papers exist. Exiting.')
            break

        print('Sleeping for %i seconds' % (args.wait_time , ))
        time.sleep(args.wait_time + random.uniform(0, 3))

    # save the database before we quit, if we found anything new
    if num_added_total > 0:
        print('Saving database with %d papers to %s' % (len(db), config.DB_PATH))
        safe_pickle_dump(db, config.DB_PATH)
