"""Loads twitter user info (json format) and geojson file for 50 us states.
Generates new geojson file with new properties added to each state.
Colums include count, percent, decile, quintile.
Also added female & male counts, and percent-female.

On turtleland2, runs in conda env "minerva-dev", which includes
numpy, pandas, geojson
"""
from __future__ import print_function
from builtins import input

import argparse
import json
import logging
import operator
import os
import sys

import geojson
import numpy as np
import pandas as pd

logging.basicConfig(level=logging.DEBUG)

def make_args():
    '''For debug/test; create argparse.Namespace with predefined inputs
    '''


class Filter:
    def __init__(self, options):
        ''''''
        self.options = options
        self.input_feature_collection = None  # state geometry
        self.input_data = None  # pandas DataFrame
        self.state_data = None  # ditto


        # load the states geojson file here
        if self.options.geojson_states_filename:
            with open(options.geojson_states_filename) as inp:
                geo = geojson.load(inp)
                if geo.get('type') != 'FeatureCollection':
                    raise Exception('geojson_states_filename not a FeatureCollection')
                self.input_feature_collection = geo.get('features')
                print('Input geojson has %d features' % len(self.input_feature_collection))

    @classmethod
    def make_parser(klass):
        description = 'Generate geojson file with twitter user counts'
        parser = argparse.ArgumentParser(description=description)
        parser.add_argument('-a', '--age_group', type=int, \
            help='Filter age_group by INDEX (use -l to list)')
        # parser.add_argument('-d', '--dump_filename', \
        #     help='Output file to dump stats and related info')
        parser.add_argument('-g', '--gender', type=int, \
            help='Filter gender by INDEX (0=female, 1=male)')
        parser.add_argument('-i', '--input_filename', required=True, \
            help='Input file with twitter user data from UAlbany')
        parser.add_argument('-l', '--list-filters', action='store_true', \
            help='List the available filters')
        parser.add_argument('-m', '--max_rows', help='Maximum number of rows to input')

        # User argument group to for output geojson, which requires input states geom
        source_dir = os.path.dirname(os.path.abspath(__file__))
        shared_dir = os.path.join(source_dir, os.pardir, 'shared')
        default_geostates = os.path.join(shared_dir, 'shapefile_50_states.geojson')
        parser.add_argument('-s', '--geojson_states_filename', default=default_geostates, \
            help='Use substitute file for the US states geometry (geojson)')
        parser.add_argument('-o', '--output_filename', help='Output filename (geojson)')

        return parser

    def load_user_data(self):
        '''Loads input data into pandas dataframe
        '''
        if not self.options.input_filename:
            raise Exception('No input_filename specified in options')

        # Read input file into a list of dictionaries
        input_list = list()
        line_count = 0
        with open(self.options.input_filename) as f:
            while True:
                line = f.readline()
                if not line:
                    break

                line_count += 1
                if self.options.max_rows and line_count > self.options.max_rows:
                    break

                data = json.loads(line)
                # Extract geo_loc and move its fields to top level
                geo_loc = data.get('geo_loc')
                if geo_loc:
                    data.update(geo_loc)
                    del data['geo_loc']

                input_list.append(data)

        # Create DataFrame from raw input
        raw_input_data = pd.DataFrame(input_list)
        logging.info('input file size {}'.format(len(raw_input_data)))

        # Prune out duplicates
        dedup = raw_input_data.drop_duplicates('user_id')
        logging.info('after removing duplicates: {}'.format(len(dedup)))

        # Prune out those missing "state" field
        # This will be our starting input data
        self.input_data = dedup.loc[ pd.notna(dedup['state']) ]
        logging.info('after removing users without state field: {}'
            .format(len(self.input_data)))

        # test = self.input_data.sort_values('id', ascending=False)
        # print(test.head())

    def list_filters(self):
        '''Writes out gender and age_group choices for filtering input data

        '''
        format_string = '{:3d}: {}'

        gender_list = self.input_data.gender.unique()
        print()
        print('Gender Choices (option -g):')
        for i,gender in enumerate(sorted(gender_list)):
            print(format_string.format(i, gender))

        age_group_list = self.input_data.age_group.unique()
        print()
        print('Age Group Choices (option -a):')
        for i,age_group in enumerate(sorted(age_group_list)):
            print(format_string.format(i, age_group))

        print()

    def apply_filters(self):
        '''Filters input data based on gender and/or age_group options
        '''
        for name in ['gender', 'age_group']:
            index = getattr(self.options, name)
            if index is None:
                continue
            print('Applying {} filter'.format(name))

            # Get the set of values for the given column
            option_list = self.input_data[name].unique()
            if index < 0 or index >= len(option_list):
                print()
                print('{} option {} not valid; use one of:'
                    .format(name, index))
                format_string = '{:3d}: {}'
                for i,value in enumerate(sorted(option_list)):
                    print(format_string.format(i, value))
                print()
                sys.exit(-101)

            value = option_list[index]  # get the value to match
            df = self.input_data.loc[self.input_data[name] == value]
            self.input_data = df
            logging.info('data size: {}'.format(len(self.input_data)))


    def compile_by_state(self):
        '''Builds DataFrame with results for each state
        '''
        # Start by creating internal dataframe with relevant columns
        df_min = pd.DataFrame()
        df_min['count'] = self.input_data['user_id']
        df_min['state'] = self.input_data['state']
        df_min['age_group'] = self.input_data['age_group']
        df_min['gender'] = self.input_data['gender']

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

        # Copy state and count columns to begin self.state_data frame
        self.state_data = grp_state.filter(['state', 'count'])

        # Add column with the percentage of total users in each state
        total_count = self.state_data['count'].sum()
        #print('total_count {}'.format(total_count))
        percent = 100.0 * self.state_data['count'] / total_count
        self.state_data['percent'] = percent.round(1)
        #print(self.state_data.head())

        # Generate rank column
        self.state_data['rank'] = self.state_data['count'].rank()

        # Generate decile and quintile data
        self.state_data['decile'] = pd.qcut(
            self.state_data['count'], 10, duplicates='drop', labels=False)
        self.state_data['quintile'] = pd.qcut(
            self.state_data['count'], 5, duplicates='drop', labels=False)

        # Compute percent-female column
        # To do this:
        #   1. Group by state & gender
        #   2. Reshape count column into 2 columns (female, male)
        #   3. Create new dataframe with female, male columns
        #   4. Compute percent-female column
        #   5. Append percent-female column to state_data
        grp_state_gender = df_min.groupby(['state', 'gender']).count()
        # Use state for new index, gender for new columns
        state_levels, gender_levels = grp_state_gender.index.levels
        # Shape of new dataframe same as grp_state_gender
        nrows,ncols = grp_state_gender.index.levshape  # ()

        # Get and reshape the count column from df_min
        s_count = grp_state_gender['count']  # (series)
        s_count_reshape = np.reshape(s_count.values,(nrows, ncols))

        # Create dataframe with states as index, and genders as columns
        df_state_gender = pd.DataFrame(
            s_count_reshape, index=state_levels, columns=gender_levels)
        self.state_data['female'] = df_state_gender['female']
        self.state_data['male'] = df_state_gender['male']
        # Compute percent female
        percent_female = 100.0 * df_state_gender['female'] / \
            (df_state_gender['female'] + df_state_gender['male'])

        self.state_data['pct-female'] = percent_female.round(1)

        # Same basic thing to calculate percent by age_group
        # Todo need to reindex in order to fill in missing rows
        # grp_state_age = df_min.groupby(['state', 'age_group']).count()
        # state_levels,age_levels = grp_state_age.index.levels
        # nrows,ncols = grp_state_age.index.levshape
        # s_count = grp_state_age['count']
        # s_count.reshape = np.reshape(s_count.values, (nrows,ncols))
        # df_state_age = pd.DataFrame(
        #     s_count_reshape, index=state_levels, columns=age_levels)
        # print(df_state_age.head())

        print(self.state_data.head(9))
        csv_filename = 'tw2state.csv'
        self.state_data.to_csv(csv_filename)
        print('Wrote {}'.format(csv_filename))


    def write_output(self):
        ''''''
        # Build a dictionary of the input feature collection (state geometry)
        input_dict = dict()
        for feature in self.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_data.iterrows():
            input_feature = input_dict.get(index)
            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 _dump_stats(self, input_count, bystate_dict):
        '''Writes data and some stats out
        '''
        if self.options.dump_filename:
            fp = open(self.options.dump_filename, 'wt')
        else:
            fp = sys.stdout

        fp.write('Total number of inputs: {}\n'.format(input_count))

        for k,v in sorted(bystate_dict.items()):
            fp.write('{:20s}: {:5d}\n'.format(k,v))
        # total_count = sum(bystate_dict.values())
        # print('Number of inputs with state info: {}'.format(total_count))

        if self.options.dump_filename:
            fp.close()
            print('Wrote info file {}'.format(self.options.dump_filename))


if __name__ == '__main__':
    parser = Filter.make_parser()
    options = parser.parse_args()
    #print(options)

    filter = Filter(options)
    filter.load_user_data()

    if options.list_filters:
        filter.list_filters()
        sys.exit(0)

    # Apply filters
    filter.apply_filters()

    filter.compile_by_state()

    if options.output_filename:
        filter.write_output()
