import pandas as pd
import numpy as np
import re
import json
import os
import warnings
import shutil
from pathlib import Path
import codecs

"""
Helper functions to load the article data. The main method to use
is load_data().
"""

# Caution! Modifying this in code will have no effect since the
# default arguments are populated with this reference at creation
# time, so post-hoc modifications will do nothing.
__data_folder = os.path.join(os.path.split(__file__)[0], '..', 'data')


def clean_string(s):
    """
    Clean all the HTML/Unicode nastiness out of a string.
    Replaces newlines with spaces.
    """

    return s.replace('\r', '').replace('\n', ' ').replace('\xa0', ' ').strip()


def load_articles(data_folder=__data_folder, nrows=None):
    """
    Loads the articles CSV. Can optionally only load the first
    `nrows` number of rows.
    """
    column_names = ['id',
                    'feedname',
                    'url',
                    'orig_html',
                    'title',
                    'bodytext',
                    'relevant',
                    'created',
                    'last_modified',
                    'news_source_id',
                    'author']

    return pd.read_csv(os.path.join(data_folder,
                                    'newsarticles_article.csv'),
                       header=None,
                       names=column_names,
                       nrows=nrows,
                       dtype={'orig_html': str, 'author': str})


def load_taggings(data_folder=__data_folder):
    """Loads the type-of-crime human tagging of the articles."""
    uc_column_names = ['id', 'date', 'relevant',
                       'article_id', 'user_id', 'locations']

    uc = pd.read_csv(os.path.join(data_folder,
                                  'newsarticles_usercoding.csv'),
                     header=None,
                     names=uc_column_names)

    uc.set_index('id', drop=True, inplace=True)

    uc_tags_column_names = ['id', 'usercoding_id', 'category_id']

    uc_tags = pd.read_csv(
        os.path.join(data_folder, 'newsarticles_usercoding_categories.csv'),
        header=None,
        names=uc_tags_column_names
    )
    uc_tags.set_index('usercoding_id', drop=True, inplace=True)

    uc_tags['article_id'] = uc.loc[uc_tags.index, 'article_id']
    return uc_tags


def load_model_categories(data_folder=__data_folder):
    tcr_names = ['id', 'relevance', 'category_id', 'coding_id']
    tc_names = ['id', 'date', 'model_info', 'relevance', 'article_id']

    tcr = pd.read_csv(
        os.path.join(data_folder, 'newsarticles_trainedcategoryrelevance.csv'),
        names=tcr_names
    )
    tc = pd.read_csv(
        os.path.join(data_folder, 'newsarticles_trainedcoding.csv'),
        names=tc_names
    ).set_index('id', drop=True)
    tcr['article_id'] = tc.loc[tcr['coding_id']]['article_id'].values
    return tcr


def load_model_locations(data_folder=__data_folder):
    tl_names = ['id', 'text', 'latitude', 'longitude', 'coding_id']
    tc_names = ['id', 'date', 'model_info', 'relevance', 'article_id']

    tl = pd.read_csv(
        os.path.join(data_folder, 'newsarticles_trainedlocation.csv'),
        names=tl_names
    )
    tc = pd.read_csv(
        os.path.join(data_folder, 'newsarticles_trainedcoding.csv'),
        names=tc_names
    ).set_index('id', drop=True)
    tl['article_id'] = tc.loc[tl['coding_id']]['article_id'].values
    return tl


def load_locations(data_folder=__data_folder):
    """Load the human-extracted locations from the articles."""
    uc_column_names = ['id', 'date', 'relevant',
                       'article_id', 'user_id', 'locations']

    uc = pd.read_csv(os.path.join(data_folder,
                                  'newsarticles_usercoding.csv'),
                     header=None,
                     names=uc_column_names)

    uc['locations'] = uc['locations'].apply(lambda x: json.loads(x))

    return uc


def load_categories(data_folder=__data_folder):
    """Loads the mapping of id to names/abbrevations of categories"""
    column_names = ['id', 'category_name', 'abbreviation', 'created',
                    'active', 'kind']

    return pd.read_csv(os.path.join(data_folder, 'newsarticles_category.csv'),
                       header=None,
                       names=column_names)


def load_data(data_folder=__data_folder, nrows=None):
    """
    Creates a dataframe of the article information and k-hot encodes the tags
    into columns called cat_NUMBER. The k-hot encoding is done assuming that
    the categories are 1-indexed and there are as many categories as the
    maximum value of the numerical cateogry_id column.

    Inputs:
        data_folder:
            A folder containing the data files in CSV format.
        nrows:
            Number of articles to load. Defaults to all, which uses about 4
            GB of memory.
    """

    df = load_articles(data_folder=data_folder, nrows=nrows)
    df['relevant'] = df['relevant'] == 't'
    df.rename(columns={'id': 'article_id'}, inplace=True)
    df.set_index('article_id', drop=True, inplace=True)
    # hopefully this will save some memory/space, can add back if needed
    del(df['orig_html'])

    tags_df = load_taggings(data_folder)
    # will help cacheing
    tags_df.sort_values(by='article_id', inplace=True)
    tags_df = tags_df.loc[tags_df['article_id'].isin(
        df.index.intersection(tags_df['article_id']))]

    locs_df = load_locations(data_folder)
    locs_df.sort_values(by='article_id', inplace=True)
    locs_df = locs_df.loc[locs_df['article_id'].isin(
        df.index.intersection(locs_df['article_id']))]

    model_tags_df = load_model_categories(data_folder)
    # will help cacheing
    model_tags_df.sort_values(by='article_id', inplace=True)
    model_tags_df = model_tags_df.loc[model_tags_df['article_id'].isin(
        df.index.intersection(model_tags_df['article_id']))]

    # init with empty lists
    df['locations'] = np.empty([df.shape[0], 0]).tolist()
    loc_article_ids = locs_df['article_id'].values
    df.loc[loc_article_ids, 'locations'] = locs_df['locations'].values

    def find_loc_in_string(locs, string):
        """
        The locations are generated from JavaScript, which means there's
        going to be some problems getting things to line up exactly and
        neatly. This function will hopefully performa all necessary
        transformations to find the given location text within the
        larger string.

        Inputs:
            locs: list of locations as loaded by load_locations
            string: bodytext of article in which to find locs
        Returns:
            updated_locs: list of locations as loaded by
                load_locations, but with a couple
                extra fields ('cleaned text' and 'cleaned span').
        """

        for i, loc in enumerate(locs):
            loc_txt = loc['text']

            loc_txt = clean_string(loc_txt)
            string = clean_string(string)

            loc['cleaned text'] = loc_txt

            spans = [x.span() for x in re.finditer(re.escape(loc_txt), string)]
            if spans:
                # The string may have occurred multiple times, and since the
                # spans don't line up perfectly we can't know which one is the
                # "correct" one. Best we can do is find the python span closest
                # to the expected javascript span.
                closest = np.argmin(np.abs(
                    np.array([x[0] for x in spans]) - loc['start']
                ))
                loc['cleaned span'] = spans[closest]

            locs[i] = loc

        return locs

    df['locations'] = df.apply(
        lambda r: find_loc_in_string(r['locations'], r['bodytext']),
        axis=1
    )

    num_no_match = df['locations'].apply(
        lambda locs: any([('cleaned span' not in loc) for loc in locs])
    ).sum()
    if num_no_match:
        warnings.warn(('{} location strings were not found in'
                       ' the bodytext.').format(num_no_match),
                      RuntimeWarning)

    model_locations_df = load_model_locations(data_folder)
    model_locations_df = model_locations_df.set_index('article_id')
    model_locations_gb = model_locations_df.groupby('article_id')
    model_locations_text = model_locations_gb['text'].apply(list)
    df['model_location_text'] = model_locations_text

    categories_df = load_categories(data_folder)
    categories_df.set_index('id', drop=True, inplace=True)

    # tags_df['category_id'] = tags_df['category_id'].astype(str)
    tags_df['category_abbreviation'] = (categories_df
                                        ['abbreviation']
                                        [tags_df['category_id']]
                                        .values)
    model_tags_df['category_abbreviation'] = (categories_df
                                              ['abbreviation']
                                              [model_tags_df['category_id']]
                                              .values)

    if np.setdiff1d(tags_df['article_id'].values, df.index.values).size:
        warnings.warn('Tags were found for article IDs that do not exist.',
                      RuntimeWarning)

    def update_df_with_categories(article_ids, cat_abbreviations, vals,
                                  is_model):
        # for some reason, some articles that are tagged don't show up
        # in the articles CSV. filter those out.
        existing_ids_filter = np.isin(article_ids, df.index.values)

        article_ids = article_ids[existing_ids_filter]
        cat_abbreviations = cat_abbreviations[existing_ids_filter]
        vals = vals[existing_ids_filter]

        for i in range(categories_df.shape[0]):
            cat_name = categories_df.loc[i+1, 'abbreviation']
            if is_model:
                cat_name += '_model'
            df[cat_name] = 0
            if not is_model:
                df[cat_name] = df[cat_name].astype('int8')
            matches = cat_abbreviations == cat_name
            if not matches.sum():
                continue
            df.loc[article_ids[matches], cat_name] = vals[matches]

    update_df_with_categories(
        model_tags_df['article_id'].values,
        model_tags_df['category_abbreviation'].values + '_model',
        model_tags_df['relevance'].values,
        is_model=True
    )
    update_df_with_categories(
        tags_df['article_id'].values,
        tags_df['category_abbreviation'].values,
        np.ones((tags_df['article_id'].values.shape), dtype='int8'),
        is_model=False
    )

    df.loc[df['bodytext'].isnull(), 'bodytext'] = ''

    return df


def subsample_and_resave(out_folder, n=5, input_folder=__data_folder,
                         random_seed=5):
    """
    Subsamples the CSV data files so that we have at least
    `n` articles from each type-of-crime tag as determined
    by the human coding. Saves the subsampled CSV data
    into `out_folder`. If there are fewer than `n` articles
    tagged with a type-of-crime, then we will use all of
    the articles with that tag.

    Inputs
    ------
    out_folder : str
        Path to folder where data should be saved. Should already exist.
    n : int
        How many examples from each category should we have?
    input_folder : str
        Path to where the full CSV files are saved.
    random_seed : None or int
        np.random.RandomState() will be seeded with this value
        in order to perform the random subsampling.
    """
    out_folder = str(Path(out_folder).expanduser().absolute())
    input_folder = str(Path(input_folder).expanduser().absolute())
    if out_folder == input_folder:
        raise RuntimeError('out_folder cannot match input_folder.')

    random_state = np.random.RandomState(random_seed)

    df = load_data(input_folder)
    chosen_indexes = []
    for crime_type in df.loc[:, 'OEMC':].columns:
        is_type = df[crime_type].astype(bool)
        n_samps = min(n, is_type.sum())
        chosen_indexes += (df.loc[is_type, :]
                           .sample(n_samps, random_state=random_state)
                           .index
                           .tolist())
    del df

    chosen_indexes = sorted(list(set(chosen_indexes)))

    # newsarticles_article.csv
    articles_df = load_articles(input_folder)
    sample = (articles_df
              .reset_index()
              .set_index('id')
              .loc[chosen_indexes, 'index'])
    articles_df = articles_df.loc[sample, :]
    # garble garble
    articles_df['bodytext'] = articles_df['bodytext'].astype(str).apply(
        lambda x: codecs.encode(x, 'rot-13')
    )
    articles_df.to_csv(os.path.join(out_folder, 'newsarticles_article.csv'),
                       header=None, index=False)
    del articles_df

    # newsarticles_category.csv
    shutil.copyfile(os.path.join(input_folder, 'newsarticles_category.csv'),
                    os.path.join(out_folder, 'newsarticles_category.csv'))

    # newsarticles_usercoding.csv
    uc_column_names = ['id', 'date', 'relevant',
                       'article_id', 'user_id', 'locations']

    uc_df = pd.read_csv(os.path.join(input_folder,
                                     'newsarticles_usercoding.csv'),
                        header=None,
                        names=uc_column_names)

    sample = np.where(uc_df['article_id'].isin(chosen_indexes))[0]
    uc_df.loc[sample, :].to_csv(
        os.path.join(out_folder, 'newsarticles_usercoding.csv'),
        header=None, index=False
    )

    uc_tags_column_names = ['id', 'usercoding_id', 'category_id']

    # newsarticles_usercoding_categories.csv
    uc_tags_df = pd.read_csv(
        os.path.join(input_folder,
                     'newsarticles_usercoding_categories.csv'),
        header=None,
        names=uc_tags_column_names,
        dtype={'id': int, 'usercoding_id': int, 'category_id': int}
    )
    sample = np.where(uc_df
                      .set_index('id')
                      .loc[uc_tags_df['usercoding_id'], 'article_id']
                      .isin(chosen_indexes)
                      )[0]
    uc_tags_df = uc_tags_df.loc[sample, :]
    uc_tags_df.to_csv(
        os.path.join(out_folder, 'newsarticles_usercoding_categories.csv'),
        header=None, index=False
    )

    # newsarticles_trainedcoding
    tc_names = ['id', 'date', 'model_info', 'relevance', 'article_id']
    tc = pd.read_csv(
        'tagnews/data/newsarticles_trainedcoding.csv',
        names=tc_names
    )
    tc = tc.loc[tc['article_id'].isin(chosen_indexes)]
    tc.to_csv(
        os.path.join(out_folder, 'newsarticles_trainedcoding.csv'),
        header=False, index=False
    )

    # newsarticles_trainedcategoryrelevance
    tcr_names = ['id', 'relevance', 'category_id', 'coding_id']
    tcr = pd.read_csv(
        'tagnews/data/newsarticles_trainedcategoryrelevance.csv',
        names=tcr_names
    )
    tcr = tcr.loc[tcr['coding_id'].isin(tc['id'])]
    tcr.to_csv(
        os.path.join(out_folder, 'newsarticles_trainedcategoryrelevance.csv'),
        header=False, index=False
    )

    # newsarticles_trainedlocation
    tl_names = ['id', 'text', 'latitude', 'longitude', 'coding_id']
    tl = pd.read_csv(
        'tagnews/data/newsarticles_trainedlocation.csv',
        names=tl_names
    )
    tl = tl.loc[tl['coding_id'].isin(tc['id'])]
    tl.to_csv(
        os.path.join(out_folder, 'newsarticles_trainedlocation.csv'),
        header=False, index=False
    )


def load_crime_data(data_folder=__data_folder):
    crimes = pd.read_csv(os.path.join(data_folder, 'Crimes.csv'))
    crimes = crimes[crimes['Year'] > 2010]

    crime_string = pd.Series('', crimes.index)

    # ['ID', 'Case Number', 'Date', 'Block', 'IUCR', 'Primary Type',
    #  'Description', 'Location Description', 'Arrest', 'Domestic', 'Beat',
    #  'District', 'Ward', 'Community Area', 'FBI Code', 'X Coordinate',
    #  'Y Coordinate', 'Year', 'Updated On', 'Latitude', 'Longitude',
    #  'Location']

    # TODO: synonyms on this for month name, weekday name,
    # time of day (e.g. afternoon), etc.
    crime_string += crimes['Date'] + ' '

    # TODO: synonyms?
    crime_string += crimes['Primary Type'] + ' '

    # TODO: synonyms?
    crime_string += crimes['Description'] + ' '

    # TODO: synonyms?
    crime_string += crimes['Location Description'] + ' '

    # TODO: synonyms?
    iucr = pd.read_csv(os.path.join(data_folder, 'IUCR.csv'))
    iucr.set_index('IUCR', drop=True, inplace=True)
    idx = iucr.index
    idx_values = idx.values
    idx_values[idx.str.len() == 3] = '0' + idx_values[idx.str.len() == 3]
    crime_string += (iucr.loc[crimes['IUCR'], 'PRIMARY DESCRIPTION']
                     .fillna('')
                     .values
                     + ' ')
    crime_string += (iucr.loc[crimes['IUCR'], 'SECONDARY DESCRIPTION']
                     .fillna('')
                     .values
                     + ' ')
    community_areas = pd.read_csv(os.path.join(data_folder, 'CommAreas.csv'))
    community_areas.set_index('AREA_NUM_1', inplace=True, drop=True)
    crime_string += (community_areas.loc[crimes['Community Area'], 'COMMUNITY']
                     .fillna('')
                     .values
                     + ' ')

    return crimes, crime_string


def load_ner_data(data_folder=__data_folder):
    """
    Loads ner.csv from the specified data folder.

    The column 'stag' is a binary value indicating whether or not
    the row corresponds to the entity "geo". Typically, you will
    want to use column 'word' to predict the column 'stag'.
    """
    df = pd.read_csv(os.path.join(data_folder, 'ner.csv'),
                     encoding="ISO-8859-1",
                     error_bad_lines=False,
                     index_col=0)

    df.dropna(subset=['word', 'tag'], inplace=True)
    df.reset_index(inplace=True, drop=True)
    df['stag'] = (df['tag'] == 'B-geo') | (df['tag'] == 'I-geo')
    df['all_tags'] = df['tag']
    df['tag'] = df['stag']
    df = df[['word', 'all_tags', 'tag']]

    return df
