"""
    ulmo.cpc.drought.core
    ~~~~~~~~~~~~~~~~~~~~~

    This module provides direct access to `Climate Predicition Center`_ `Weekly
    Drought Index`_ dataset.

    .. _Climate Prediction Center: http://www.cpc.ncep.noaa.gov/
    .. _Weekly Drought Index: http://www.cpc.ncep.noaa.gov/products/analysis_monitoring/cdus/palmer_drought/
"""
from __future__ import division
from builtins import str
from builtins import range
from past.utils import old_div

import datetime
import os
import requests

import numpy as np
import pandas

from ulmo import util

# directory where drought data will be stashed
CPC_DROUGHT_DIR = os.path.join(util.get_ulmo_dir(), 'cpc/drought')

# state codes (note: these are not FIPS codes)
STATE_CODES = {
    'AL': 1,
    'AZ': 2,
    'AR': 3,
    'CA': 4,
    'CO': 5,
    'CT': 6,
    'DE': 7,
    'FL': 8,
    'GA': 9,
    'IA': 13,
    'ID': 10,
    'IL': 11,
    'IN': 12,
    'KS': 14,
    'KY': 15,
    'LA': 16,
    'MA': 19,
    'MD': 18,
    'ME': 17,
    'MI': 20,
    'MN': 21,
    'MO': 23,
    'MS': 22,
    'MT': 24,
    'NC': 31,
    'ND': 32,
    'NE': 25,
    'NH': 27,
    'NJ': 28,
    'NM': 29,
    'NV': 26,
    'NY': 30,
    'OH': 33,
    'OK': 34,
    'OR': 35,
    'PA': 36,
    'PR': 66,
    'RI': 37,
    'SC': 38,
    'SD': 39,
    'TN': 40,
    'TX': 41,
    'UT': 42,
    'VA': 44,
    'VT': 43,
    'WA': 45,
    'WI': 47,
    'WV': 46,
    'WY': 48,
}

def get_data(state=None, climate_division=None, start=None, end=None,
             as_dataframe=False):
    """Retreives data.


    Parameters
    ----------
    state : ``None`` or str
        If specified, results will be limited to the state corresponding to the
        given 2-character state code.
    climate_division : ``None`` or int
        If specified, results will be limited to the climate division.
    start : ``None`` or date (see :ref:`dates-and-times`)
        Results will be limited to those after the given date. Default is the
        start of the current calendar year.
    end : ``None`` or date (see :ref:`dates-and-times`)
        If specified, results will be limited to data before this date.
    as_dataframe: bool
        If ``False`` (default), a dict with a nested set of dicts will be
        returned with data indexed by state, then climate division. If ``True``
        then a pandas.DataFrame object will be returned.  The pandas dataframe
        is used internally, so setting this to ``True`` is a little bit faster
        as it skips a serialization step.


    Returns
    -------
    data : dict or pandas.Dataframe
        A dict or pandas.DataFrame representing the data. See the
        ``as_dataframe`` parameter for more.
    """
    if not start is None:
        start_date = util.convert_date(start)
    else:
        start_date = None
    if not end is None:
        end_date = util.convert_date(end)
    else:
        end_date = None

    if not end_date:
        end_date = datetime.date.today()
    if not start_date:
        start_date = datetime.date(end_date.year, 1, 1)

    start_year, start_week = _week_number(start_date)
    end_year, end_week = _week_number(end_date)

    if state:
        state_code = STATE_CODES.get(state.upper())
    else:
        state_code = None

    data = None
    for year in range(start_year, end_year + 1):
        url, current_year_flag = _get_data_url(year)
        format_type = _get_data_format(year)
        with _open_data_file(url) as data_file:
            year_data = _parse_data_file(data_file, format_type, year, current_year_flag)

        if state_code:
            year_data = year_data[year_data['state_code'] == state_code]
        if climate_division:
            year_data = year_data[year_data['climate_division'] == climate_division]

        year_data = _reindex_data(year_data)

        if data is None:
            data = year_data
        else:
            # some data are duplicated (e.g. final data from 2011 stretches into
            # prelim data of 2012), so just take those that are new
            append_index = year_data.index.difference(data.index)
            if len(append_index):
                data = data.append(year_data.loc[append_index])

    # restrict results to date range
    period_index = pandas.PeriodIndex(data['period'])
    periods_in_range = (period_index >= start_date) & (period_index <= end_date)
    data = data[periods_in_range]

    # this does what data.reset_index() should do, but at least as of 0.10.1, that sets
    # will cast period objects to ints
    data.index = np.arange(len(data))
    if as_dataframe:
        return data
    else:
        return _as_data_dict(data)


def _as_data_dict(dataframe):
    data_dict = {}
    for state in dataframe['state'].unique():
        state_dict = {}
        state_dataframe = dataframe[dataframe['state'] == state]
        for name, group in state_dataframe.groupby(['state', 'climate_division']):
            s, climate_division = name
            climate_division_data = group.T.drop(['state', 'climate_division'])
            values = [
                _value_dict(value)
                for k, value in climate_division_data.iteritems()
            ]
            state_dict[climate_division] = values
        data_dict[state] = state_dict
    return data_dict


def _convert_state_codes(dataframe):
    """adds state abbreviations to a dataframe, based on state codes"""
    state_codes = pandas.DataFrame(
        np.array([i for i in STATE_CODES.items()],
                 dtype=np.dtype([('state', '|U2'), ('code', int)])))
    merged = pandas.merge(dataframe, state_codes,
            left_on='state_code', right_on='code', how='left')
    column_names = dataframe.columns.tolist()
    column_names.remove('state_code')
    column_names.insert(0, 'state')
    return merged[column_names]


def _convert_week_numbers(dataframe):
    """convert a dataframe's week numbers to period objects"""
    weeks = [key for key, group in dataframe.groupby(['year', 'week'])]
    periods = [(week[0], week[1], _period_for_week(*week)) for week in weeks]
    period_dataframe = pandas.DataFrame(periods, columns=['year', 'week', 'period'])
    merged = pandas.merge(dataframe, period_dataframe,
            left_on=['year', 'week'], right_on=['year', 'week'])
    column_names = dataframe.columns.tolist()
    column_names.remove('week')
    column_names.remove('year')
    column_names.insert(2, 'period')
    return merged[column_names]


def _first_sunday(year):
    """returns the first Sunday of a growing season, which is the first Sunday
    after the first Wednesday in March
    """
    first_day = datetime.date(year, 3, 1)
    if first_day.weekday() == 6:
        return first_day
    elif first_day.weekday() <= 2:
        return first_day - pandas.tseries.offsets.Week(weekday=6)
    else:
        return first_day + pandas.tseries.offsets.Week(weekday=6)


def _get_data_format(year):
    if year >= 2001:
        return 'format5'
    elif 1997 <= year <= 2000:
        return 'format4'
    else:
        return 'format2'


def _get_data_url(year):
    current_year, current_week = _week_number(datetime.date.today())
    if year == current_year:
        return ('https://ftp.cpc.ncep.noaa.gov/htdocs/temp4/current.data', True)
    elif year == current_year - 1:
        url = ('https://ftp.cpc.ncep.noaa.gov/htdocs/temp2/palmer%s-PRELIM' % str(year)[-2:],
                False)
        if not _url_exists(url[0]):
            url = ('https://ftp.cpc.ncep.noaa.gov/htdocs/temp4/current.data', True)
        return url
    elif year <= 1985:
        return ('https://ftp.cpc.ncep.noaa.gov/htdocs/temp2/palmer73-85', False)
    else:
        url = ('https://ftp.cpc.ncep.noaa.gov/htdocs/temp2/palmer%s' % str(year)[-2:], False)
        if not _url_exists(url[0]):
            url = ('https://ftp.cpc.ncep.noaa.gov/htdocs/temp2/palmer%s-PRELIM' % str(year)[-2:],
                    False)
        return url


def _open_data_file(url):
    """returns an open file handle for a data file; downloading if necessary or otherwise using a previously downloaded file"""
    file_name = url.rsplit('/', 1)[-1]
    file_path = os.path.join(CPC_DROUGHT_DIR, file_name)
    return util.open_file_for_url(url, file_path, check_modified=True, use_bytes=True)


def _parse_data_file(data_file, palmer_format, year, current_year_flag):
    """
    based on the fortran format strings:
        format2: FORMAT(I4,3I2,F4.1,F4.0,10F6.2,4F6.4,F6.3,10F6.2,F4.0,12F6.2)
        format4: FORMAT(2I4,I2,F4.1,F4.0,10F6.2,4F6.4,F6.3,10F6.2,F4.0,12F6.2)
        format5: FORMAT(2I4,I2,F5.2,F5.1,10F6.2,4F6.4,F6.3,10F6.2,F4.0,12F6.2)
    """
    if palmer_format == 'format5':
        delim_sequence = (2, 2, 4, 2, 5, 5) + 10*(6,) + 4*(6,) + (6,) + 10*(6,) + (4,) + 12*(6,)
        use_columns = (0, 1, 2, 3, 4, 5, 9, 15, 28, 29, 37, 40, 41)
    elif palmer_format == 'format4':
        delim_sequence = (2, 2, 4, 2, 4, 4) + 10*(6,) + 4*(6,) + (6,) + 10*(6,) + (4,) + 12*(6,)
        use_columns = (0, 1, 2, 3, 4, 5, 9, 15, 28, 29, 37, 40, 41)
    elif palmer_format == 'format2':
        delim_sequence = (2, 2, 2, 2, 2, 4, 4) + 10*(6,) + 4*(6,) + (6,) + 10*(6,) + (4,) + 12*(6,)
        use_columns = (0, 1, 2, 3, 5, 6, 10, 16, 29, 30, 38, 41, 42)
    else:
        raise NotImplementedError("we have not implemented the format for given date range")

    dtype = [
        ('state_code', 'i1'),
        ('climate_division', 'i1'),
        ('year', 'i4'),
        ('week', 'i4'),
        ('precipitation', 'f8'),
        ('temperature', 'f8'),
        ('potential_evap', 'f8'),
        ('runoff', 'f8'),
        ('soil_moisture_upper', 'f8'),
        ('soil_moisture_lower', 'f8'),
        ('pdsi', 'f8'),
        ('cmi', 'f8')
    ]

    decodef = lambda x: x.decode("utf-8")
    data_array = np.genfromtxt(data_file, dtype=dtype, delimiter=delim_sequence, usecols=use_columns)
    if not current_year_flag:
        data_array['year'] = year
    dataframe = pandas.DataFrame(data_array)
    return dataframe


def _periods_for_range(start_date, end_date):
    return pandas.period_range(start_date, end_date, freq='W-SAT')


def _period_for_week(year, week_number):
    """returns a pandas.Period for a given growing season year and week number"""
    first_sunday = _first_sunday(year)
    return pandas.Period(first_sunday, freq='W-SAT') + week_number - 1


def _reindex_data(dataframe):
    dataframe = _convert_week_numbers(dataframe)
    dataframe = _convert_state_codes(dataframe)
    return dataframe.set_index(
        ['state', 'climate_division', 'period'], drop=False)


def _url_exists(url):
    return requests.head(url).status_code == 200


def _value_dict(value):
    value_dict = value.to_dict()
    value_dict['period'] = str(value_dict['period'])
    return value_dict


def _week_number(date):
    """returns the growing season week number for a given datetime.date"""
    first_sunday = _first_sunday(date.year)
    date_ts = pandas.Timestamp(date)
    first_sunday_ts = pandas.Timestamp(first_sunday)
    if date_ts < first_sunday_ts:
        first_sunday_ts = pandas.Timestamp(_first_sunday(date.year - 1))
    days_since_first_sunday = (date_ts - first_sunday_ts).days
    return (first_sunday_ts.year, (old_div(days_since_first_sunday, 7)) + 1)