"""
Lookup table and function to find countries that have a specific language as their official or de-facto languages.

LANG_TO_COUNTRY_LOOKUP: A dict keyed to lang code (ISO-639-1) with values of a set of country codes (ISO-3166-1 alpha-2)
get_country_set_for_language(lang_code): Convenience function for the lookup table


Lookup table generated using the following code:

import iso3166
import babel.languages
from pprint import pprint

lang_lookup = {}
for code in iso3166.codes:
    for lang in babel.languages.get_official_languages(code, de_facto=True):
        if len(lang) > 2:
            continue
        lookup_value = lang_lookup.setdefault(lang, set())
        lookup_value.add(code)
        lang_lookup[lang] = lookup_value

pprint(lang_lookup, indent=4, compact=True, sort_dicts=False)
"""

LANG_TO_COUNTRY_LOOKUP = {
    'fa': {'IR', 'AF'},
    'ps': {'AF'},
    'sv': {'FI', 'AX', 'SE'},
    'sq': {'AL'},
    'ar': {
        'AE',
        'BH',
        'DJ',
        'DZ',
        'EG',
        'EH',
        'ER',
        'IL',
        'IQ',
        'JO',
        'KM',
        'KW',
        'LB',
        'LY',
        'MA',
        'MR',
        'OM',
        'PS',
        'QA',
        'SA',
        'SD',
        'SO',
        'SY',
        'TD',
        'TN',
        'YE',
    },
    'fr': {
        'BE',
        'BF',
        'BI',
        'BJ',
        'BL',
        'CA',
        'CD',
        'CF',
        'CG',
        'CH',
        'CI',
        'CM',
        'DJ',
        'DZ',
        'FR',
        'GA',
        'GF',
        'GN',
        'GP',
        'GQ',
        'HT',
        'KM',
        'LU',
        'MA',
        'MC',
        'MF',
        'MG',
        'ML',
        'MQ',
        'MU',
        'NC',
        'NE',
        'PF',
        'PM',
        'RE',
        'RW',
        'SC',
        'SN',
        'SY',
        'TD',
        'TG',
        'TN',
        'VU',
        'WF',
        'YT',
    },
    'sm': {'AS', 'WS'},
    'en': {
        'AG',
        'AI',
        'AS',
        'AU',
        'BB',
        'BI',
        'BM',
        'BS',
        'BW',
        'BZ',
        'CA',
        'CC',
        'CK',
        'CM',
        'CX',
        'DM',
        'ER',
        'FJ',
        'FK',
        'FM',
        'GB',
        'GD',
        'GG',
        'GH',
        'GI',
        'GM',
        'GU',
        'GY',
        'HK',
        'IE',
        'IM',
        'IN',
        'IO',
        'JE',
        'JM',
        'KE',
        'KI',
        'KN',
        'KY',
        'LC',
        'LR',
        'LS',
        'MG',
        'MH',
        'MP',
        'MS',
        'MT',
        'MU',
        'MW',
        'NA',
        'NF',
        'NG',
        'NR',
        'NU',
        'NZ',
        'PG',
        'PH',
        'PK',
        'PN',
        'PR',
        'PW',
        'RW',
        'SB',
        'SC',
        'SD',
        'SG',
        'SH',
        'SL',
        'SS',
        'SX',
        'SZ',
        'TC',
        'TK',
        'TO',
        'TT',
        'TV',
        'TZ',
        'UG',
        'UM',
        'US',
        'VC',
        'VG',
        'VI',
        'VU',
        'WS',
        'ZA',
        'ZM',
        'ZW',
    },
    'ca': {'AD'},
    'pt': {'GQ', 'BR', 'GW', 'MZ', 'MO', 'PT', 'ST', 'AO', 'CV', 'TL'},
    'es': {
        'AR',
        'BO',
        'CL',
        'CO',
        'CR',
        'CU',
        'DO',
        'EC',
        'ES',
        'GQ',
        'GT',
        'HN',
        'MX',
        'NI',
        'PA',
        'PE',
        'PR',
        'PY',
        'SV',
        'UY',
        'VE',
    },
    'hy': {'AM'},
    'nl': {'CW', 'NL', 'AW', 'SX', 'SR', 'BE', 'BQ'},
    'de': {'LU', 'AT', 'DE', 'LI', 'BE', 'CH'},
    'az': {'AZ'},
    'bn': {'BD'},
    'be': {'BY'},
    'ru': {'RU', 'KZ', 'KG', 'UA', 'BY'},
    'dz': {'BT'},
    'qu': {'PE', 'EC', 'BO'},
    'ay': {'BO'},
    'bs': {'BA'},
    'hr': {'BA', 'HR'},
    'sr': {'RS', 'BA'},
    'tn': {'BW'},
    'ms': {'BN', 'MY', 'SG'},
    'bg': {'BG'},
    'rn': {'BI'},
    'km': {'KH'},
    'sg': {'CF'},
    'zh': {'CN', 'SG'},
    'el': {'GR', 'CY'},
    'tr': {'TR', 'CY'},
    'cs': {'CZ'},
    'da': {'DK'},
    'ti': {'ER'},
    'et': {'EE'},
    'am': {'ET'},
    'fo': {'FO'},
    'fj': {'FJ'},
    'fi': {'FI'},
    'ty': {'PF'},
    'ka': {'GE'},
    'kl': {'GL'},
    'ch': {'GU'},
    'ht': {'HT'},
    'it': {'VA', 'IT', 'SM', 'CH'},
    'hu': {'HU'},
    'is': {'IS'},
    'hi': {'IN'},
    'id': {'ID'},
    'ga': {'IE'},
    'gv': {'IM'},
    'he': {'IL'},
    'ja': {'JP'},
    'kk': {'KZ'},
    'sw': {'UG', 'TZ', 'KE'},
    'ko': {'KR', 'KP'},
    'ky': {'KG'},
    'lo': {'LA'},
    'lv': {'LV'},
    'st': {'LS'},
    'lt': {'LT'},
    'lb': {'LU'},
    'mk': {'MK'},
    'mg': {'MG'},
    'ny': {'MW'},
    'dv': {'MV'},
    'mt': {'MT'},
    'mh': {'MH'},
    'ro': {'RO', 'MD'},
    'mn': {'MN'},
    'my': {'MM'},
    'na': {'NR'},
    'ne': {'NP'},
    'mi': {'NZ'},
    'yo': {'NG'},
    'no': {'NO'},
    'nb': {'NO', 'SJ'},
    'nn': {'NO'},
    'ur': {'PK'},
    'ho': {'PG'},
    'gn': {'PY'},
    'pl': {'PL'},
    'rw': {'RW'},
    'wo': {'SN'},
    'ta': {'LK', 'SG'},
    'sk': {'SK'},
    'sl': {'SI'},
    'so': {'SO'},
    'si': {'LK'},
    'ss': {'SZ'},
    'tg': {'TJ'},
    'th': {'TH'},
    'to': {'TO'},
    'tk': {'TM'},
    'uk': {'UA'},
    'uz': {'UZ'},
    'bi': {'VU'},
    'vi': {'VN'},
    'sn': {'ZW'},
    'nd': {'ZW'},
}


def get_country_set_for_language(lang_code):
    """
    Get a set of ISO-3166-1 alpha-2  country codes where a language is officially spoken or is de-facto official
    :param lang_code: ISO-639-1 language code
    :type lang_code: str
    :return: Set of ISO-3166-1 alpha-2 country codes
    """
    return LANG_TO_COUNTRY_LOOKUP.get(lang_code, set())
