#!/usr/bin/env python
# coding: utf-8

# # Harvest agencies associated with *all* functions
# 
# This notebook loops through the list of functions that were [extracted from the RecordSearch interface](harvesting_functions_from_recordsearch.ipynb) and saves basic details of the agencies responsible for each function. To keep down the file size and avoid too much duplication it doesn't include the full range of relationships that an agency might have. If you want the full agency data, use [this notebook](get_agencies_associated_with_function.ipynb) to harvest agencies associated with an indivividual function or hierarchy.
# 
# The JSON data file created has the following structure:
# 
# ``` json
# [
#     {
#         'term': FUNCTION NAME
#         'agencies': [
#             'agency_id': AGENCY IDENTIFIER,
#             'title': AGENCY NAME,
#             'dates': {
#                 'date_str': AGENCY LIFE DATES AS A STRING,
#                 'start_date': AGENCY START DATE (YYYY-MM-DD),
#                 'end_date': AGENCY END DATE (YYYY-MM-DD),
#             },
#             'agency_status': TYPE/LEVEL OF AGENCY,
#             'location': AGENCY LOCATION,
#             'function_start_date': DATE AGENCY STARTED BEING RESPONSIBLE FOR THIS FUNCTION (YYYY-MM-DD),
#             'function_end_date': DATE AGENCY STOPPED BEING RESPONSIBLE FOR THIS FUNCTION (YYYY-MM-DD),
#         ]
#     }
# ]
# ```

# ## Set up the harvesting code

# In[1]:


import json
import math
import time
import ipywidgets as widgets
from IPython.display import display, HTML, FileLink, clear_output
import pandas as pd
from tqdm import tqdm_notebook
from slugify import slugify
from tinydb import TinyDB, Query
from recordsearch_tools.client import RSAgencySearchClient
from recordsearch_tools.utilities import retry
import arrow


# In[2]:


class AgencyHarvester(object):
    '''
    Searches for agencies associated with a particular function.
    Loops through pages in the results set saving agency details.
    '''
    
    def __init__(self, function):
        self.function = function
        self.total_pages = 0
        self.total_results = 0
        self.agencies = []
        self.client = RSAgencySearchClient()
        self.prepare_harvest()

    @retry(ConnectionError, tries=20, delay=10, backoff=1)
    def prepare_harvest(self):
        '''
        Finds the number of results and calculates how many pages need to be harvested.
        '''
        # Setting results_per_page to zero makes things much faster
        response = self.client.search_agencies(function=self.function, sort='1', results_per_page=0)
        total_results = self.client.total_results
        print('{} agencies'.format(total_results if total_results else 'No'))
        if total_results:
            self.total_pages = math.floor(int(total_results) / self.client.results_per_page) + 1

    @retry(ConnectionError, tries=20, delay=10, backoff=1)
    def start_harvest(self, start=1):
        '''
        Loop through each page of results saving the results.
        '''
        if self.total_pages:
            for page in tqdm_notebook(range(start, self.total_pages + 1), unit='page', desc='Pages:'):
                response = self.client.search_agencies(function=self.function, page=page, sort='1', date_format='iso')
                self.agencies += response['results']
                time.sleep(1)
                
def get_children(function):
    '''
    Gets child terms of a given function.
    '''
    f_list = []
    if 'narrower' in function:
        for subf in function['narrower']:
            f_list.append(subf['term'])
            f_list += get_children(subf)
    return f_list

def load_functions():
    '''
    Loads a pre-harvested JSON file containing functions data.
    Returns a flat list of functions.
    '''
    functions_list = []
    with open('data/functions.json', 'r') as json_file:
        functions = json.load(json_file)
    for function in functions:
        functions_list.append(function['term'])
        functions_list += get_children(function)
    # Get rid of duplicates
    functions_list = set(functions_list)
    # Sort terms
    functions_list = sorted(functions_list)
    return functions_list

def get_function_dates(function, agency):
    '''
    Get the dates an agency was responsible for a given function.
    '''
    dates = {}
    # Loop through the functions associated with an agency
    for f in agency['functions']:
        # Find the current function
        if f['identifier'].lower() == function:
            # Get the dates this agency was responsible for the current function
            dates['function_start_date'] = f['start_date']
            dates['function_end_date'] = f['end_date']
            break
    return dates

def get_all_agencies():
    '''
    Sends function terms off to the harvester to get related agencies.
    '''
    clear_output()
    Record = Query()
    # Get a list of functions
    functions = load_functions()
    db = TinyDB('data/db_agencies_by_function')
    # Loop through the list of functions
    for function in functions:
        clear_output()
        print('\nHarvesting "{}"'.format(function))
        # Fire up the harvester for this function
        client = AgencyHarvester(function=function)
        client.start_harvest()
        agencies = []
        # Create a subset of the agency data to limit the filesize
        for a in client.agencies:
            # Keep the fields we want
            agency = {k: a[k] for k in ['agency_id', 'title', 'dates', 'agency_status', 'location']}
            # Add extra fields to show when the agency was responsible for this function
            agency.update(get_function_dates(function, a))
            agencies.append(agency)
        db.upsert({'term': function, 'agencies': agencies}, Record.term == function)


# ## Start the harvest

# In[4]:


get_all_agencies()


# ## Save the results for download

# In[14]:


def save_json():
    db = TinyDB('data/db_agencies_by_function')
    functions = db.all()
    filename = 'data/agencies_by_function.json'
    with open(filename, 'w') as json_file:
        json.dump(functions, json_file, indent=4)
        display(FileLink(filename))


# In[15]:


save_json()