# -*- coding: utf-8 -*-

import warnings

from pandas.compat import reduce, lrange, string_types
import pandas as pd
import numpy as np

from pandas_datareader.base import _BaseReader

# This list of country codes was pulled from wikipedia during October 2014.
# While some exceptions do exist, it is the best proxy for countries supported
# by World Bank.  It is an aggregation of the 2-digit ISO 3166-1 alpha-2, and
# 3-digit ISO 3166-1 alpha-3, codes, with 'all', 'ALL', and 'All' appended ot
# the end.

country_codes = ['AD', 'AE', 'AF', 'AG', 'AI', 'AL', 'AM', 'AO', 'AQ', 'AR',
                 'AS', 'AT', 'AU', 'AW', 'AX', 'AZ', 'BA', 'BB', 'BD', 'BE',
                 'BF', 'BG', 'BH', 'BI', 'BJ', 'BL', 'BM', 'BN', 'BO', 'BQ',
                 'BR', 'BS', 'BT', 'BV', 'BW', 'BY', 'BZ', 'CA', 'CC', 'CD',
                 'CF', 'CG', 'CH', 'CI', 'CK', 'CL', 'CM', 'CN', 'CO', 'CR',
                 'CU', 'CV', 'CW', 'CX', 'CY', 'CZ', 'DE', 'DJ', 'DK', 'DM',
                 'DO', 'DZ', 'EC', 'EE', 'EG', 'EH', 'ER', 'ES', 'ET', 'FI',
                 'FJ', 'FK', 'FM', 'FO', 'FR', 'GA', 'GB', 'GD', 'GE', 'GF',
                 'GG', 'GH', 'GI', 'GL', 'GM', 'GN', 'GP', 'GQ', 'GR', 'GS',
                 'GT', 'GU', 'GW', 'GY', 'HK', 'HM', 'HN', 'HR', 'HT', 'HU',
                 'ID', 'IE', 'IL', 'IM', 'IN', 'IO', 'IQ', 'IR', 'IS', 'IT',
                 'JE', 'JM', 'JO', 'JP', 'KE', 'KG', 'KH', 'KI', 'KM', 'KN',
                 'KP', 'KR', 'KW', 'KY', 'KZ', 'LA', 'LB', 'LC', 'LI', 'LK',
                 'LR', 'LS', 'LT', 'LU', 'LV', 'LY', 'MA', 'MC', 'MD', 'ME',
                 'MF', 'MG', 'MH', 'MK', 'ML', 'MM', 'MN', 'MO', 'MP', 'MQ',
                 'MR', 'MS', 'MT', 'MU', 'MV', 'MW', 'MX', 'MY', 'MZ', 'NA',
                 'NC', 'NE', 'NF', 'NG', 'NI', 'NL', 'NO', 'NP', 'NR', 'NU',
                 'NZ', 'OM', 'PA', 'PE', 'PF', 'PG', 'PH', 'PK', 'PL', 'PM',
                 'PN', 'PR', 'PS', 'PT', 'PW', 'PY', 'QA', 'RE', 'RO', 'RS',
                 'RU', 'RW', 'SA', 'SB', 'SC', 'SD', 'SE', 'SG', 'SH', 'SI',
                 'SJ', 'SK', 'SL', 'SM', 'SN', 'SO', 'SR', 'SS', 'ST', 'SV',
                 'SX', 'SY', 'SZ', 'TC', 'TD', 'TF', 'TG', 'TH', 'TJ', 'TK',
                 'TL', 'TM', 'TN', 'TO', 'TR', 'TT', 'TV', 'TW', 'TZ', 'UA',
                 'UG', 'UM', 'US', 'UY', 'UZ', 'VA', 'VC', 'VE', 'VG', 'VI',
                 'VN', 'VU', 'WF', 'WS', 'YE', 'YT', 'ZA', 'ZM', 'ZW',
                 'ABW', 'AFG', 'AGO', 'AIA', 'ALA', 'ALB', 'AND', 'ARE',
                 'ARG', 'ARM', 'ASM', 'ATA', 'ATF', 'ATG', 'AUS', 'AUT',
                 'AZE', 'BDI', 'BEL', 'BEN', 'BES', 'BFA', 'BGD', 'BGR',
                 'BHR', 'BHS', 'BIH', 'BLM', 'BLR', 'BLZ', 'BMU', 'BOL',
                 'BRA', 'BRB', 'BRN', 'BTN', 'BVT', 'BWA', 'CAF', 'CAN',
                 'CCK', 'CHE', 'CHL', 'CHN', 'CIV', 'CMR', 'COD', 'COG',
                 'COK', 'COL', 'COM', 'CPV', 'CRI', 'CUB', 'CUW', 'CXR',
                 'CYM', 'CYP', 'CZE', 'DEU', 'DJI', 'DMA', 'DNK', 'DOM',
                 'DZA', 'ECU', 'EGY', 'ERI', 'ESH', 'ESP', 'EST', 'ETH',
                 'FIN', 'FJI', 'FLK', 'FRA', 'FRO', 'FSM', 'GAB', 'GBR',
                 'GEO', 'GGY', 'GHA', 'GIB', 'GIN', 'GLP', 'GMB', 'GNB',
                 'GNQ', 'GRC', 'GRD', 'GRL', 'GTM', 'GUF', 'GUM', 'GUY',
                 'HKG', 'HMD', 'HND', 'HRV', 'HTI', 'HUN', 'IDN', 'IMN',
                 'IND', 'IOT', 'IRL', 'IRN', 'IRQ', 'ISL', 'ISR', 'ITA',
                 'JAM', 'JEY', 'JOR', 'JPN', 'KAZ', 'KEN', 'KGZ', 'KHM',
                 'KIR', 'KNA', 'KOR', 'KWT', 'LAO', 'LBN', 'LBR', 'LBY',
                 'LCA', 'LIE', 'LKA', 'LSO', 'LTU', 'LUX', 'LVA', 'MAC',
                 'MAF', 'MAR', 'MCO', 'MDA', 'MDG', 'MDV', 'MEX', 'MHL',
                 'MKD', 'MLI', 'MLT', 'MMR', 'MNE', 'MNG', 'MNP', 'MOZ',
                 'MRT', 'MSR', 'MTQ', 'MUS', 'MWI', 'MYS', 'MYT', 'NAM',
                 'NCL', 'NER', 'NFK', 'NGA', 'NIC', 'NIU', 'NLD', 'NOR',
                 'NPL', 'NRU', 'NZL', 'OMN', 'PAK', 'PAN', 'PCN', 'PER',
                 'PHL', 'PLW', 'PNG', 'POL', 'PRI', 'PRK', 'PRT', 'PRY',
                 'PSE', 'PYF', 'QAT', 'REU', 'ROU', 'RUS', 'RWA', 'SAU',
                 'SDN', 'SEN', 'SGP', 'SGS', 'SHN', 'SJM', 'SLB', 'SLE',
                 'SLV', 'SMR', 'SOM', 'SPM', 'SRB', 'SSD', 'STP', 'SUR',
                 'SVK', 'SVN', 'SWE', 'SWZ', 'SXM', 'SYC', 'SYR', 'TCA',
                 'TCD', 'TGO', 'THA', 'TJK', 'TKL', 'TKM', 'TLS', 'TON',
                 'TTO', 'TUN', 'TUR', 'TUV', 'TWN', 'TZA', 'UGA', 'UKR',
                 'UMI', 'URY', 'USA', 'UZB', 'VAT', 'VCT', 'VEN', 'VGB',
                 'VIR', 'VNM', 'VUT', 'WLF', 'WSM', 'YEM', 'ZAF', 'ZMB',
                 'ZWE', 'all', 'ALL', 'All']


class WorldBankReader(_BaseReader):

    """
    Download data series from the World Bank's World Development Indicators

    Parameters
    ----------

    symbols: WorldBank indicator string or list of strings
        taken from the ``id`` field in ``WDIsearch()``

    countries: string or list of strings.
        ``all`` downloads data for all countries
        2 or 3 character ISO country codes select individual
        countries (e.g.``US``,``CA``) or (e.g.``USA``,``CAN``).  The codes
        can be mixed.

        The two ISO lists of countries, provided by wikipedia, are hardcoded
        into pandas as of 11/10/2014.

    start: Timestamp or int
        First year of the data series. Month and day are ignored.

    end: Timestamp or int
        Last year of the data series (inclusive). Month and day are ignored.

    errors: str {'ignore', 'warn', 'raise'}, default 'warn'
        Country codes are validated against a hardcoded list.  This controls
        the outcome of that validation, and attempts to also apply
        to the results from world bank.

        errors='raise', will raise a ValueError on a bad country code.

    Returns
    -------

    ``pandas`` DataFrame with columns: country, iso_code, year,
    indicator value.
    """

    _format = 'json'

    def __init__(self, symbols=None, countries=None,
                 start=None, end=None,
                 retry_count=3, pause=0.001, session=None, errors='warn'):

        if symbols is None:
            symbols = ['NY.GDP.MKTP.CD', 'NY.GNS.ICTR.ZS']
        elif isinstance(symbols, string_types):
            symbols = [symbols]

        super(WorldBankReader, self).__init__(symbols=symbols,
                                              start=start, end=end,
                                              retry_count=retry_count,
                                              pause=pause, session=session)

        if countries is None:
            countries = ['MX', 'CA', 'US']
        elif isinstance(countries, string_types):
            countries = [countries]

        bad_countries = np.setdiff1d(countries, country_codes)
        # Validate the input
        if len(bad_countries) > 0:
            tmp = ", ".join(bad_countries)
            if errors == 'raise':
                raise ValueError("Invalid Country Code(s): %s" % tmp)
            if errors == 'warn':
                warnings.warn('Non-standard ISO '
                              'country codes: %s' % tmp, UserWarning)

        self.countries = countries
        self.errors = errors

    @property
    def url(self):
        countries = ';'.join(self.countries)
        return ("http://api.worldbank.org/countries/" +
                countries + "/indicators/")

    @property
    def params(self):
        return {'date': '{0}:{1}'.format(self.start.year, self.end.year),
                'per_page': 25000, 'format': 'json'}

    def read(self):
        try:
            return self._read()
        finally:
            self.close()

    def _read(self):
        data = []
        for indicator in self.symbols:
            # Build URL for api call
            try:
                df = self._read_one_data(self.url + indicator, self.params)
                df.columns = ['country', 'iso_code', 'year', indicator]
                data.append(df)

            except ValueError as e:
                msg = str(e) + ' Indicator: ' + indicator
                if self.errors == 'raise':
                    raise ValueError(msg)
                elif self.errors == 'warn':
                    warnings.warn(msg)

        # Confirm we actually got some data, and build Dataframe
        if len(data) > 0:
            out = reduce(lambda x, y: x.merge(y, how='outer'), data)
            out = out.drop('iso_code', axis=1)
            out = out.set_index(['country', 'year'])
            out = out.apply(pd.to_numeric, errors='ignore')

            return out
        else:
            msg = "No indicators returned data."
            raise ValueError(msg)

    def _read_lines(self, out):
        # Check to see if there is a possible problem
        possible_message = out[0]

        if 'message' in possible_message.keys():
            msg = possible_message['message'][0]
            try:
                msg = msg['key'].split() + ["\n "] + msg['value'].split()
                wb_err = ' '.join(msg)
            except:
                wb_err = ""
                if 'key' in msg.keys():
                    wb_err = msg['key'] + "\n "
                if 'value' in msg.keys():
                    wb_err += msg['value']

            msg = "Problem with a World Bank Query \n %s." % wb_err
            raise ValueError(msg)

        if 'total' in possible_message.keys():
            if possible_message['total'] == 0:
                msg = "No results found from world bank."
                raise ValueError(msg)

        # Parse JSON file
        data = out[1]
        country = [x['country']['value'] for x in data]
        iso_code = [x['country']['id'] for x in data]
        year = [x['date'] for x in data]
        value = [x['value'] for x in data]
        # Prepare output
        df = pd.DataFrame([country, iso_code, year, value]).T
        return df

    def get_countries(self):
        """Query information about countries

        Provides information such as:
            country code, region, income level,
            capital city, latitude, and longitude
        """
        url = 'http://api.worldbank.org/countries/?per_page=1000&format=json'

        resp = self._get_response(url)
        data = resp.json()[1]

        data = pd.DataFrame(data)
        data.adminregion = [x['value'] for x in data.adminregion]
        data.incomeLevel = [x['value'] for x in data.incomeLevel]
        data.lendingType = [x['value'] for x in data.lendingType]
        data.region = [x['value'] for x in data.region]
        data.latitude = [float(x) if x != "" else np.nan
                         for x in data.latitude]
        data.longitude = [float(x) if x != "" else np.nan
                          for x in data.longitude]
        data = data.rename(columns={'id': 'iso3c', 'iso2Code': 'iso2c'})
        return data

    def get_indicators(self):
        """Download information about all World Bank data series"""
        global _cached_series
        if isinstance(_cached_series, pd.DataFrame):
            return _cached_series.copy()

        url = 'http://api.worldbank.org/indicators?per_page=50000&format=json'

        resp = self._get_response(url)
        data = resp.json()[1]

        data = pd.DataFrame(data)
        # Clean fields
        data.source = [x['value'] for x in data.source]

        def encode_ascii(x):
            return x.encode('ascii', 'ignore')

        data.sourceOrganization = data.sourceOrganization.apply(encode_ascii)
        # Clean topic field

        def get_value(x):
            try:
                return x['value']
            except:
                return ''

        def get_list_of_values(x):
            return [get_value(y) for y in x]

        data.topics = data.topics.apply(get_list_of_values)
        data.topics = data.topics.apply(lambda x: ' ; '.join(x))

        # Clean output
        data = data.sort_values(by='id')
        data.index = pd.Index(lrange(data.shape[0]))

        # cache
        _cached_series = data.copy()

        return data

    def search(self, string='gdp.*capi', field='name', case=False):
        """
        Search available data series from the world bank

        Parameters
        ----------

        string: string
            regular expression
        field: string
            id, name, source, sourceNote, sourceOrganization, topics
            See notes below
        case: bool
            case sensitive search?

        Notes
        -----

        The first time this function is run it will download and cache the full
        list of available series. Depending on the speed of your network
        connection, this can take time. Subsequent searches will use the cached
        copy, so they should be much faster.

        id : Data series indicator (for use with the ``indicator`` argument of
        ``WDI()``) e.g. NY.GNS.ICTR.GN.ZS"
        name: Short description of the data series
        source: Data collection project
        sourceOrganization: Data collection organization
        note:
        sourceNote:
        topics:
        """
        indicators = self.get_indicators()
        data = indicators[field]
        idx = data.str.contains(string, case=case)
        out = indicators.loc[idx].dropna()
        return out


def download(country=None, indicator=None, start=2003, end=2005,
             errors='warn', **kwargs):
    """
    Download data series from the World Bank's World Development Indicators

    Parameters
    ----------

    indicator: string or list of strings
        taken from the ``id`` field in ``WDIsearch()``

    country: string or list of strings.
        ``all`` downloads data for all countries
        2 or 3 character ISO country codes select individual
        countries (e.g.``US``,``CA``) or (e.g.``USA``,``CAN``).  The codes
        can be mixed.

        The two ISO lists of countries, provided by wikipedia, are hardcoded
        into pandas as of 11/10/2014.

    start: int
        First year of the data series

    end: int
        Last year of the data series (inclusive)

    errors: str {'ignore', 'warn', 'raise'}, default 'warn'
        Country codes are validated against a hardcoded list.  This controls
        the outcome of that validation, and attempts to also apply
        to the results from world bank.

        errors='raise', will raise a ValueError on a bad country code.

    kwargs:
        keywords passed to WorldBankReader

    Returns
    -------

    ``pandas`` DataFrame with columns: country, iso_code, year,
    indicator value.

    """
    return WorldBankReader(symbols=indicator, countries=country,
                           start=start, end=end, errors=errors,
                           **kwargs).read()


def get_countries(**kwargs):
    """Query information about countries

    Provides information such as:
        country code, region, income level,
        capital city, latitude, and longitude

    Parameters
    ----------

    kwargs:
        keywords passed to WorldBankReader

    """
    return WorldBankReader(**kwargs).get_countries()


def get_indicators(**kwargs):
    '''Download information about all World Bank data series

    Parameters
    ----------

    kwargs:
        keywords passed to WorldBankReader

    '''
    return WorldBankReader(**kwargs).get_indicators()


_cached_series = None


def search(string='gdp.*capi', field='name', case=False, **kwargs):
    """
    Search available data series from the world bank

    Parameters
    ----------

    string: string
        regular expression
    field: string
        id, name, source, sourceNote, sourceOrganization, topics
        See notes below
    case: bool
        case sensitive search?
    kwargs:
        keywords passed to WorldBankReader

    Notes
    -----

    The first time this function is run it will download and cache the full
    list of available series. Depending on the speed of your network
    connection, this can take time. Subsequent searches will use the cached
    copy, so they should be much faster.

    id : Data series indicator (for use with the ``indicator`` argument of
    ``WDI()``) e.g. NY.GNS.ICTR.GN.ZS"
    name: Short description of the data series
    source: Data collection project
    sourceOrganization: Data collection organization
    note:
    sourceNote:
    topics:
    """

    return WorldBankReader(**kwargs).search(string=string, field=field,
                                            case=case)
