from __future__ import print_function
import os
import glob
import argparse
import json
import string

from alcohol_surveillance_utils import alcohol_keywords
from keyword_matching_utils import make_aho_keyword_automaton
from keyword_matching_utils import has_any_keyword

from tqdm import tqdm


class InstagramAlcoholPostFilter:

    def __init__(self, match_full_words=False):

        if match_full_words:
            keyword_list = [' ' + w.strip() + ' '
                            for w in alcohol_keywords]
        else:
            keyword_list = alcohol_keywords
        self.alcohol_kwa = make_aho_keyword_automaton(keyword_list)

        self.printable = set(string.printable)
        self.punctuation = set(string.punctuation)

    def _clean_text(self, input_text):

        input_text = filter(lambda x: x in self.printable,
                            input_text).lower()

        input_text = input_text.encode('ascii', 'ignore')

        input_text = ''.join(
            [x if x not in self.punctuation else ' ' for x in input_text])

        return input_text

    def get_post_text(self, post):

        # Get caption and tags
        caption_text = post.get('caption', {}).get('text')
        tag_list = post.get('tags')

        # create string with caption and tags
        # NOTE: this may not be necessary as caption seems to contain all tags
        input_text = ''

        if caption_text:
            input_text += caption_text

        if len(tag_list) > 0:
            input_text += ' '.join(tag_list)

        return input_text

    def filter_post(self, post):
        """
        Checks if the given post json has one of the alcohol keywords
        """

        # get post text
        post_text = self.get_post_text(post)

        if not post_text:
            return False

        # clean up text
        post_text = self._clean_text(post_text)

        if not post_text:
            return False

        # look for keywords
        match = has_any_keyword(post_text, self.alcohol_kwa)

        return match

    def filter_posts(self, input_json_file, output_json_file=None,
                     limit=None, state=None):
        """
        Filters all the posts in the input json file
        """

        print('\n---------------------------------------')

        # read input json file
        print('\n>>> Loading input file - {}'.format(input_json_file))

        with open(input_json_file) as f:
            text = f.read()
            input_json = json.loads(text)

        if not text:
            raise Exception('Input file missing, empty, or unreadable ({})'
                            .format(input_json_file))

        print('Number of input posts: {}'.format(len(input_json)))

        # filter posts
        print('\n>>> Running filter')

        output_json = []

        if limit is not None:
            input_json = input_json[:limit]

        for i in tqdm(range(len(input_json))):

            post = input_json[i]

            match = self.filter_post(post)

            if match:
                if state:
                    post['state'] = state
                output_json.append(post)

        print('Number of matching posts: {}'.format(len(output_json)))

        # write to output file
        if output_json_file is not None:

            print('\n>>> Writing output file - {}'.format(output_json_file))

            with open(output_json_file, 'w') as f:
                f.write(json.dumps(output_json))
                f.write('\n')

        else:

            return output_json


if __name__ == '__main__':

    parser = argparse.ArgumentParser()

    parser.add_argument('-if', '--input_filename',
                        help='Input json file containing instagram posts')

    parser.add_argument('-id', '--input_dir',
                        help='Input directory containing json files with '
                             'instagram posts')

    parser.add_argument('-of', '--output_filename',
                        help='Output json file containing the filtered posts')

    parser.add_argument('-od', '--output_dir',
                        help='Output directory for writing json files with '
                             'filtered posts')

    parser.add_argument('-e', '--match_full_words', action='store_true',
                        help='match full words instead of partial')

    parser.add_argument('-l', '--limit', type=int,
                        help='Limit number of items processed')

    parser.add_argument('-s', '--state', help='Add state field to data')

    options = parser.parse_args()

    print(options)

    # check args
    assert options.input_filename or options.input_dir, \
        "Either input_filename of input_dir must be provided"
    assert not (options.input_filename and options.input_dir), \
        "Only one of input_filename of input_dir must be provided"

    # filter posts
    if options.input_dir:

        assert options.output_dir, 'Output directory was not provided'

        # Create output dir if it doesnt exist
        if not os.path.isdir(options.output_dir):
            os.makedirs(options.output_dir)
        else:
            Warning('Output directory already exists!')

        # Look for json files in input directory
        print('>>> Found the following json files in the input directory:')

        json_file_list = sorted(glob.glob(
            os.path.join(options.input_dir, '*.json')))

        for f in json_file_list:
            print(f)

        # create classification filter
        iap_filter = InstagramAlcoholPostFilter(
            match_full_words=options.match_full_words)

        # Run the filter
        for input_json_file in json_file_list:

            output_json_file = os.path.join(
                options.output_dir,
                os.path.basename(input_json_file)
            )

            iap_filter.filter_posts(
                input_json_file,
                output_json_file=output_json_file,
                limit=options.limit
            )

    if options.input_filename:

        assert options.output_filename, 'Output filename was not provided'

        # create classification filter
        iap_filter = InstagramAlcoholPostFilter(
            match_full_words=options.match_full_words)

        # Run the filter
        iap_filter.filter_posts(
            options.input_filename,
            output_json_file=options.output_filename,
            limit=options.limit,
            state=options.state
        )
