"""Combine state json files into single dataframe
"""

from __future__ import print_function
import argparse
import json
import logging
import os
import string
import sys

import geojson
import numpy as np
import pandas as pd


# List of corresponding input & output files
# Hard-coded because Birdsong naming is inconsistent
# (filename, postal code, name, total number of posts)
states_table = [
    ('arizona.json',     'AZ', 'Arizona',     244818),
    ('california.json',  'CA', 'California',  928607),
    ('illinois.json',    'IL', 'Illinois',    319234),
    ('nevada.json',      'NV', 'Nevada',      220359),
    ('newyork.json',     'NY', 'New York',    391298),
    ('northdakota.json', 'ND', 'North Dakota', 11880),
    ('oklahoma.json',    'OK', 'Oklahoma',     63661),
    ('oregon.json',      'OR', 'Oregon',      208552),
    ('southdakota.json', 'SD', 'South Dakota', 14051),
    ('texas.json',       'TX', 'Texas',       654348),
    ('utah.json',        'UT', 'Utah',        110497),
    ('vermont.json',     'VT', 'Vermont',      27368)
]

# states_table = [
#     ('northdakota.json', 'ND', 'North Dakota', 11880),
#     ('vermont.json',     'VT', 'Vermont', 27368)
# ]

class Converter:
    def __init__(self, options):
        self.options = options
        self.input_frame = None
        self.state_frame = None

    def load_input_files(self):
        os.chdir(options.directory)
        self.input_frame = pd.DataFrame()

        # Read each file as dataframe; concatenate into list
        df_list = list()
        for i,(infile,postal,state, n) in enumerate(states_table):
            if options.index and i != options.index:
                continue
            print('Processing {}'.format(infile))
            df = self.load_json(infile)
            if not 'state' in df:
                df['state'] = state
            df_list.append(df)
        self.input_frame = pd.concat(df_list)
        # print(self.input_frame.head())
        # print(self.input_frame.info())
        # print(self.input_frame.tail())

    def load_json(self, filename):
        '''Loads json input, traverses selected columns, and returns DataFrame
        '''
        read_file = False
        with open(filename) as f:
            text = f.read()
            self.input_json = json.loads(text)
            read_file = True

        if not read_file:
            print('Did not read file -- skipping {}'.format(filename))
            return None
        #print('Checkpoint 1000 {}'.format(len(self.input_json)))

        # Traverse input and build list of tuples, 1 for each input object
        #for post in self.input_json:
        # Base columns on those previously extracted from MongoDB
        # Main exception - drop tags and convert caption text to ascii letters, numbers & whitespace
        input_fields = 'user.id,user.username,createdTime,caption.text,images.standard_resolution.url' + \
            ',location.latitude,location.longitude'
        field_names = input_fields.split(',')

        # Build records
        record_list = list()
        for i,post in enumerate(self.input_json):
            item_list = list()
            for field in field_names:
                keys = field.split('.')
                value = post
                for key in keys:
                    if key in value:
                        value = value[key]
                    else:
                        value = None
                        break

                # Convert caption to lower case ascii letters, numbers, & whitespace
                if field == 'caption.text':
                    value = self._to_ascii_string(value)
                item_list.append(value)
            if 'state' in post:
                item_list.append(post['state'])
            elif self.options.state:
                item_list.append(self.options.state)
            else:
                item_list.append(None)
            # if i < 5:
            #     print('{}'.format(item_list))
            record_list.append(tuple(item_list))

        columns = 'id,username,createdTime,caption,image_url,latitude,longitude,state'
        df = pd.DataFrame(data=record_list,columns=columns.split(','))
        return df

    def compile_by_state(self):
        ''''''
        # Start by creating internal dataframe with relevant columns
        df = pd.DataFrame()
        df['match'] = self.input_frame['id']
        df['state'] = self.input_frame['state']
        # df['age_group'] = self.input_data['age_group']
        # df['gender'] = self.input_data['gender']

        # Group by state in order to get number of users info
        grp_state = df.groupby('state', as_index=False).count()
        #print(grp_state.head())

        # Add column with total number of posts
        if self.options.index:
            grp_state['posts'] = states_table[self.options.index][3]
        else:
            input_counts = [t[3] for t in states_table]
            #print('input_count {}'.format(input_counts))
            grp_state['posts'] = input_counts

        # Copy columns to begin self.state_frame
        self.state_frame = grp_state.filter(['state', 'posts', 'match'])
        # Compute number of matching posts per 1000 total posts
        self.state_frame['perK'] = (1000 * self.state_frame['match'] / self.state_frame['posts']).round(1)

        # # Generate rank column
        self.state_frame['rank'] = self.state_frame['perK'].rank()
        print(self.state_frame)

    def write_output(self):
        ''''''
        # Get states geometry file
        source_dir = os.path.dirname(os.path.abspath(__file__))
        shared_dir = os.path.join(source_dir, os.pardir, 'shared')
        geostates_filename = os.path.join(shared_dir, 'shapefile_50_states.geojson')

        with open(geostates_filename) as inp:
            geo = geojson.load(inp)
            if geo.get('type') != 'FeatureCollection':
                raise Exception('geojson_states_filename not a FeatureCollection')
            input_feature_collection = geo.get('features')
            print('Input geojson has %d features' % len(input_feature_collection))

        # Build a dictionary of the input feature collection (state geometry)
        input_dict = dict()
        for feature in input_feature_collection:
            properties = feature.get('properties', {})
            name = properties.get('name')
            if not name:
                logging.error('input feature missing name')
                continue
            input_dict[name] = feature

        # Build output features, using geometry in input_dict
        output_feature_list = list()
        for index,row in self.state_frame.iterrows():
            input_feature = input_dict.get(row['state'])
            if input_feature is None:
                logging.info('input geometry has no feature named {} (skipping)'
                    .format(index))
                continue

            # Instantiate geojson feature
            geometry = input_feature.get('geometry')
            geo_properties = row.to_dict()
            output_feature = geojson.Feature(
                geometry=geometry, properties=geo_properties)
            output_feature_list.append(output_feature)

        output_feature_collection = geojson.FeatureCollection(output_feature_list)

        # Write feature collection
        with open(self.options.output_filename, 'w') as f:
            geojson.dump(output_feature_collection, f)
            f.write('\n')
            print('Wrote {}'.format(self.options.output_filename))





    def _to_ascii_string(self, text):
        '''Also converts to lowercase and removes punctuation

        '''
        if not text:
            return None
        # Remove chars that are not ascii letters or whitespace
        filter_chars = set(string.ascii_letters + string.whitespace)
        filter_text = filter(lambda x: x in filter_chars, text)
        ascii_text = filter_text.encode('ascii', 'ignore')
        return ascii_text.lower()



if __name__ == '__main__':
    description = 'Combine state json files and generate summary info'
    parser = argparse.ArgumentParser()
    parser.add_argument('-d', '--directory', required=True, \
        help='Directory for input & output files')
    parser.add_argument('-i', '--index', type=int, \
        help='only process one state, by index')
    parser.add_argument('-l', '--limit', type=int, \
        help='Limit number of items processed (for debug)')
    parser.add_argument('-v', '--verbose_level', default='0', \
        help='Verbose output level (default 0)')
    parser.add_argument('-o', '--output_filename', help='Output filename (geojson)')
    options = parser.parse_args()

    converter = Converter(options)
    converter.load_input_files()
    converter.compile_by_state()
    if options.output_filename:
        converter.write_output()
