"""### Tools to download/upload data on the web

Here are a collection of tools to work with data on the internet. Thus,
this module mostly requires an internet connection. Which on some
supercomputers would be turned off during a job run. In scripts, make sure to
use these to preprocess before submitting jobs.
"""
__author__ = 'Qusai Al Shidi'
__email__ = 'qusai@umich.edu'

import datetime as dt
import ftplib
from functools import lru_cache
import gzip
from operator import itemgetter
import os.path
import shutil
import urllib
import urllib.request
import warnings
from dateutil import rrule
from filecache import filecache
import numpy as np
from .tools import _nearest, carrington_rotation_number

# Global defines
# This is straight from the format guide on spdf with nicer names as second col
# spdf_name nice_name type
OMNI_HIRES_COLS = (('ID for IMF spacecraft', 'id_imf'),
                   ('ID for SW Plasma spacecraft', 'id_sw'),
                   ('# of points in IMF averages', 'num_avg_imf'),
                   ('# of points in Plasma averages', 'num_avg_sw'),
                   ('Percent interp', 'interp'),
                   ('Timeshift, sec', 'timeshift'),
                   ('RMS, Timeshift', 'rms_timeshift'),
                   ('RMS, Phase front normal', 'rms_phase'),
                   ('Time btwn observations, sec', 'dt'),
                   ('Field magnitude average, nT', 'b'),
                   ('Bx, nT (GSE, GSM)', 'bx'),
                   ('By, nT (GSE)', 'by_gse'),
                   ('Bz, nT (GSE)', 'bz_gse'),
                   ('By, nT (GSM)', 'by'),
                   ('Bz, nT (GSM)', 'bz'),
                   ('RMS SD B scalar, nT', 'rms_sd_b'),
                   ('RMS SD field vector, nT', 'rms_sd_field'),
                   ('Flow speed, km/s', 'v'),
                   ('Vx Velocity, km/s, GSE', 'vx_gse'),
                   ('Vy Velocity, km/s, GSE', 'vy_gse'),
                   ('Vz Velocity, km/s, GSE', 'vz_gse'),
                   ('Proton Density, n/cc', 'density'),
                   ('Temperature, K', 'temperature'),
                   ('Flow pressure, nPa', 'pressure'),
                   ('Electric field, mV/m', 'e'),
                   ('Plasma beta', 'beta'),
                   ('Alfven mach number', 'alfven_mach'),
                   ('X(s/c), GSE, Re', 'x_gse'),
                   ('Y(s/c), GSE, Re', 'y_gse'),
                   ('Z(s/c), GSE, Re', 'z_gse'),
                   ('BSN location, Xgse, Re', 'bsn_x_gse'),
                   ('BSN location, Ygse, Re', 'bsn_y_gse'),
                   ('BSN location, Zgse, Re', 'bsn_z_gse'),
                   ('AE-index, nT', 'ae'),
                   ('AL-index, nT', 'al'),
                   ('AU-index, nT', 'au'),
                   ('SYM/D index, nT', 'sym_d'),
                   ('SYM/H index, nT', 'sym_h'),
                   ('ASY/D index, nT', 'asy_d'),
                   ('ASY/H index, nT', 'asy_h'),
                   ('PC(N) index', 'pc_n'),
                   ('Magnetosonic mach number', 'mach'),
                   )

OMNI_LORES_COLS = (('Bartels rotation number', 'bartels'),
                   ('ID for IMF spacecraft', 'id_imf'),
                   ('ID for SW plasma spacecraft', 'id_sw'),
                   ('# of points in the IMF averages', 'num_avg_imf'),
                   ('# of points in the plasma averages', 'num_avg_sw'),
                   ('Field Magnitude Average |B|', 'b_avg'),
                   ('Magnitude of Average Field Vector', 'b'),
                   ('Lat.Angle of Aver. Field Vector', 'b_angle_lat'),
                   ('Long.Angle of Aver.Field Vector', 'b_angle_lon'),
                   ('Bx GSE, GSM', 'bx'),
                   ('By GSE', 'by_gse'),
                   ('Bz GSE', 'bz_gse'),
                   ('By GSM', 'by'),
                   ('Bz GSM', 'bz'),
                   ('sigma|B|', 'sigma_b_avg'),
                   ('sigma B', 'sigma_b'),
                   ('sigma Bx', 'sigma_bx'),
                   ('sigma By', 'sigma_by'),
                   ('sigma Bz', 'sigma_bz'),
                   ('Proton temperature', 'temperature'),
                   ('Proton Density', 'density'),
                   ('Plasma (Flow) speed', 'v'),
                   ('Plasma Flow Long. Angle', 'v_angle_lat'),
                   ('Plasma  Flow Lat. Angle', 'v_angle_lon'),
                   ('Na/Np', 'na_np'),
                   ('Flow Pressure', 'dyn_p'),
                   ('sigma T', 'sigma_t'),
                   ('sigma N', 'sigma_n'),
                   ('sigma V', 'sigma_v'),
                   ('sigma phi V', 'sigma_phi_v'),
                   ('sigma theta V', 'sigma_theta_v'),
                   ('sigma-Na/Np', 'sigma_na_np'),
                   ('Electric field', 'e'),
                   ('Plasma beta', 'beta'),
                   ('Alfven mach number', 'alfven_mach'),
                   ('Kp', 'kp'),
                   ('R', 'sunspot'),
                   ('DST Index', 'dst'),
                   ('AE-index', 'ae'),
                   ('Proton flux number/cmsq sec sr >1 Mev', 'p_flux_1'),
                   ('Proton flux number/cmsq sec sr >2 Mev', 'p_flux_2'),
                   ('Proton flux number/cmsq sec sr >4 Mev', 'p_flux_4'),
                   ('Proton flux number/cmsq sec sr >10 Mev', 'p_flux_10'),
                   ('Proton flux number/cmsq sec sr >30 Mev', 'p_flux_30'),
                   ('Proton flux number/cmsq sec sr >60 Mev', 'p_flux_60'),
                   ('Flag(***)', 'flag'),
                   ('ap-index', 'ap'),
                   ('f10.7_index', 'f10_7'),
                   ('PC(N) index', 'pc_n'),
                   ('AL-index', 'al'),
                   ('AU-index', 'au'),
                   ('Magnetosonic mach number', 'mach'),
                   )


@lru_cache(maxsize=4)
def get_omni_data(time_from, time_to, **kwargs):
    """Retrieve omni solar wind data over http.

    This will download omni data from https://spdf.gsfc.nasa.gov/pub/data/omni
    and put it into a dictionary. If your data is large, then make a csv and
    use swmfpy.io.read_omni_data().

    Note that calling this more than once with the same arguments will point to
    your original dictionary that it created. This is to speed up code that
    calls this multiple times as it requires internet access and download.
    If you mutate your original try doing an omni_dict2 = omni_dict1.copy()
    and mutate the other one.

    Args:
        time_from (datetime.datetime): The start time of the solar wind
                                       data that you want to receive.
        time_to (datetime.datetime): The end time of the solar wind data
                                     you want to receive.
        **kwargs:
            original_colnames (bool): Use the original column names from the
                                      spdf specification. The alternative is
                                      nicer and shorter names. Defaults to
                                      False.
            resolution (str): (default: 'high') Here you can choose 'high' or
                              'low' resolution omni data. Some columns appear
                              in one but not the other.

    Returns:
        dict: This will be a list of *all* columns
              available in the omni data set.

    Examples:
        ```python
        import datetime
        import swmfpy.web

        storm_start = datetime.datetime(year=2000, month=1, day=1)
        storm_end = datetime.datetime(year=2000, month=2, day=15)
        data = swmfpy.web.get_omni_data(time_from=storm_start,
                                        time_to=storm_end)
        # or for low res
        data = swmfpy.web.get_omni_data(time_from=storm_start,
                                        time_to=storm_end,
                                        resolution='low')
        ```

    Note:
        This function creates a filecache of omni data to speed up multiple
        uses of it. The cache is in the directory you run the python
        interpreter in.
    """
    # Author: Qusai Al Shidi
    # Email: qusai@umich.edu

    # Initialize
    return_data = {}
    return_data['times'] = []
    name_type = int(not kwargs.get('original_colnames', False))
    name_type = itemgetter(name_type)
    resolution = kwargs.get('resolution', 'high')
    # omni cols do not include time cols
    omni = {
        'high': {
            'urls': _urls_omni_hires,
            'cols': list(map(name_type, OMNI_HIRES_COLS)),
            'parsetime': lambda cols: (dt.datetime(int(cols[0]), 1, 1,
                                                   int(cols[2]),
                                                   int(cols[3]))
                                       + dt.timedelta(int(cols[1])-1)),
            'ntimecols': 4,  # Minute
            },
        'low': {
            'urls': _urls_omni_lores,
            'cols': list(map(name_type, OMNI_LORES_COLS)),
            'parsetime': lambda cols: (dt.datetime(int(cols[0]), 1, 1,
                                                   int(cols[2]))
                                       + dt.timedelta(int(cols[1])-1)),
            'ntimecols': 3,  # Hourly
            },
        }
    omni = omni[resolution]  # To save line space
    for col_name in omni['cols']:
        return_data[col_name] = []

    # Iterate by url to save RAM
    for url in omni['urls'](time_from, time_to):

        # Parse omni data
        for line in _download_omni_data(url):
            cols = line.decode('ascii').split()

            time = omni['parsetime'](cols)
            if time_from <= time <= time_to:
                return_data['times'] += [time]
                # Assign the data from after the time columns
                for col_name, value in zip(omni['cols'],
                                           cols[omni['ntimecols']:]):
                    if _bad_omni_num(value):
                        return_data[col_name] += [None]
                    else:
                        return_data[col_name] += [float(value)]

    for col_name in omni['cols']:
        return_data[col_name] = np.array(return_data[col_name],
                                         dtype=float)
    return return_data


@filecache
def _download_omni_data(url):
    """Downloads omni data and returns as list
    """
    return list(urllib.request.urlopen(url))


def _urls_omni_hires(time_from, time_to):
    """Returns hires omni urls from time_from to time_to
    """
    prefix = 'https://spdf.gsfc.nasa.gov/pub/data/omni/'
    prefix += 'high_res_omni/monthly_1min/'
    for date in rrule.rrule(rrule.MONTHLY,
                            dtstart=time_from,
                            until=dt.datetime(time_to.year,
                                              time_to.month,
                                              1)+dt.timedelta(32)):
        suffix = 'omni_min'
        suffix += str(date.year) + str(date.month).zfill(2)
        suffix += '.asc'
        yield prefix+suffix


def _urls_omni_lores(time_from, time_to):
    """Returns lores omni urls from time_from to time_to
    """
    prefix = 'https://spdf.gsfc.nasa.gov/pub/data/omni/'
    prefix += 'low_res_omni/'
    for year in range(time_from.year, time_to.year+1):
        suffix = 'omni2_'
        suffix += str(year)
        suffix += '.dat'
        yield prefix+suffix


def _bad_omni_num(value_string):
    """Returns true if bad or false if not. Bad numbers usually just have 9s
       in omni.
    """
    return all([char in ('9', '.') for char in value_string])


def download_magnetogram_hmi(mag_time, hmi_map='hmi.B_720s', **kwargs):
    """Downloads HMI vector magnetogram fits files.

    This will download vector magnetogram FITS files from
    Joint Science Operations Center (JSOC) near a certain hour.

    Args:
        mag_time (datetime.datetime): Time after which to find
                                      vector magnetograms.
        hmi_map (str): JSOC prefix for hmi maps. Currently can only do
                       'hmi.B_720s' and 'hmi.b_synoptic.small'.

    **kwargs:
        download_dir (str): Relative directory to download to.
        verbose (bool): (default False) print out the files it's downloading.

    Returns:
        str: list of filenames downloaded.

    Raises:
        ImportError: If module `drms` is not found.
        FileNotFoundError: If the JSOC doesn't have the magnetograms for that
                           time.

    Examples:
        ```python
        from swmfpy.web import download_magnetogram_hmi
        import datetime as dt

        # I am interested in the hmi vector magnetogram from 2014, 2, 18
        time_mag = dt.datetime(2014, 2, 18, 10)  # Around hour 10

        # Calling it will download
        filenames = download_magnetogram_hmi(mag_time=time_mag,
                                             hmi_map='B_720s',
                                             download_dir='mydir/')

        # To see my list
        print('The magnetograms I downloaded are:', filenames)

        # You may call and ignore the file list
        download_magnetogram_hmi(mag_time=time_mag,
                                 hmi_map='b_synoptic_small',
                                 download_dir='mydir')
        ```
    """
    import drms

    get_urls = {
        'hmi.B_720s': _get_urls_hmi_b720,
        'hmi.b_synoptic_small': _get_urls_hmi_b_synoptic_small,
        }
    client = drms.Client()

    urls = get_urls[hmi_map](client, mag_time)

    # Download data
    if kwargs.get('verbose', False):
        print('Starting download of magnetograms:\n')
    return_name = ''
    download_dir = kwargs.get('download_dir', '')
    if not download_dir.endswith('/') and download_dir != '':
        download_dir += '/'
    for data_time, mag_url in urls:
        if mag_url == 'BadSegLink':  # JSOC will return this if not found
            raise FileNotFoundError('Could not find those HMI magnetograms.')
        filename = 'hmi_' + str(data_time).replace(' ', '_')  # Add timestamp
        filename += '_' + mag_url.split('/')[-1]  # Last is filename
        url = 'http://jsoc.stanford.edu' + mag_url
        if kwargs.get('verbose', False):
            print(f'Downloading from {url} to {download_dir+filename}.')
        with urllib.request.urlopen(url) as fits_file:
            with open(download_dir+filename, 'wb') as local_file:
                local_file.write(fits_file.read())
        if kwargs.get('verbose', False):
            print(f'Done writing {download_dir+filename}.\n')
        return_name = download_dir+filename

    if kwargs.get('verbose', False):
        print('Completed downloads.\n')

    return return_name


def _get_urls_hmi_b_synoptic_small(client, mag_time):
    """Returns for #download_magnetogram_hmi needed urls

    Args:
        client (drms.Client): To query and return urls.
        mag_time (datetime.datetime): To find nearest magnetogram.

    Returns:
        generator that yields (datetime.datetime, str): Time of magnetogram,
            suffix url of magnetogram
    """
    import drms

    cr_number = carrington_rotation_number(mag_time)
    query_string = f'hmi.b_synoptic_small[{cr_number}]'
    components = ['Bp', 'Bt', 'Br']
    data = client.query(query_string, seg=components)
    # Generator to find the nearest time
    prefix_str = 'CR' + str(cr_number) + '_' + str(mag_time)
    urls = ((prefix_str, data[component][0]) for component in components)
    return urls


def _get_urls_hmi_b720(client, mag_time):
    """Returns for #download_magnetogram_hmi needed urls for hmi.B_720s

    Args:
        client (drms.Client): To query and return urls.
        mag_time (datetime.datetime): To find nearest magnetogram.

    Returns:
        generator that yields (datetime.datetime, str): Time of magnetogram,
            suffix url of magnetogram
    """
    import drms
    query_string = 'hmi.B_720s'
    query_string += f'[{mag_time.year}.'
    query_string += f'{str(mag_time.month).zfill(2)}.'
    query_string += f'{str(mag_time.day).zfill(2)}_'
    query_string += f'{str(mag_time.hour).zfill(2)}'
    query_string += '/1h]'
    data = client.query(query_string, key='T_REC', seg='field')
    times = drms.to_datetime(data[0].T_REC)
    nearest_time = _nearest(mag_time, times)
    # Generator to find the nearest time
    urls = ((data_time, mag_url) for (data_time, mag_url)
            in zip(times, data[1].field) if data_time == nearest_time)
    return urls


def download_magnetogram_adapt(time, map_type='fixed', **kwargs):
    """This routine downloads GONG ADAPT magnetograms.

    Downloads ADAPT magnetograms from ftp://gong2.nso.edu/adapt/maps/gong/
    to a local directory. It will download all maps with the regex file
    pattern: adapt4[0,1]3*yyyymmddhh

    Args:
        time (datetime.datetime): Time in which you want the magnetogram.
        map_type (str): (default: 'fixed')
                        Choose either 'fixed' or 'central' for
                        the map type you want.

    **kwargs:
        download_dir (str): (default is current dir) Relative directory
                            where you want the maps to be downloaded.

    Returns:
        str: First unzipped filename found.

    Raises:
        NotADirectoryError: If the adapt maps directory
                            is not found on the server.
        ValueError: If map_type is not recognized.
                    (i.e. not 'fixed' or 'central')
        FileNotFoundError: If maps were not found.

    Examples:
        ```python
        import datetime as dt

        # Use datetime objects for the time
        time_flare = dt.datetime(2018, 2, 12, hour=10)
        swmfpy.web.download_magnetogram_adapt(time=time_flare,
                                              map_type='central',
                                              download_dir='./mymaps/')
        ```
    """
    # Author: Zhenguang Huang
    # Email: zghuang@umich.edu

    if map_type == 'fixed':
        map_id = '0'
    elif map_type == 'central':
        map_id = '1'
    else:
        raise ValueError('Not recognized type of ADAPT map')

    # Go to the the ADAPT ftp server
    ftp = ftplib.FTP('gong2.nso.edu')
    ftp.login()

    # Only ADAPT GONG is considered
    ftp.cwd('adapt/maps/gong')

    # Go to the specific year
    try:
        ftp.cwd(str(time.year))
    except ftplib.all_errors:
        ftp.quit()
        raise NotADirectoryError('Cannot go to the specific year directory')

    # ADAPT maps only contains the hours for even numbers
    if time.hour % 2 != 0:
        print('Warning: Hour must be an even number.',
              'The entered hour value is changed to',
              time.hour//2*2)
    # Only consider the public (4) Carrington Fixed (0) GONG (3) ADAPT maps
    file_pattern = 'adapt4' + map_id + '3*' \
        + str(time.year).zfill(4) \
        + str(time.month).zfill(2) \
        + str(time.day).zfill(2) \
        + str(time.hour//2*2).zfill(2) + '*'
    # adapt4[0,1]3*yyyymmddhh

    filenames = ftp.nlst(file_pattern)

    if len(filenames) < 1:
        raise FileNotFoundError('Could not find a file that matches'
                                + 'the pattern.')

    directory = kwargs.get('download_dir', './')
    if directory[-1] != '/':
        directory += '/'

    for filename in filenames:
        # Only try to download if the file does not exist
        if os.path.isfile(directory+filename):
            warnings.warn(f'{filename} exists, not downloading',
                          RuntimeWarning)
        else:
            # open the file locally
            with open(directory + filename, 'wb') as fhandle:
                # try to download the magnetogram
                try:
                    ftp.retrbinary('RETR ' + filename, fhandle.write)
                except ftplib.all_errors:
                    ftp.quit()
                    raise FileNotFoundError('Cannot download ', filename)

        # unzip the file
        if '.gz' in filename:
            filename_unzip = filename.replace('.gz', '')
            with gzip.open(directory + filename, 'rb') as s_file:
                with open(directory + filename_unzip, 'wb') as d_file:
                    shutil.copyfileobj(s_file, d_file, 65536)

    # close the connection
    ftp.quit()

    # return first file name if all goes well
    return_names = filenames
    for index, filename in enumerate(return_names):
        if '.gz' in filename:
            return_names[index] = filename.replace('.gz', '')
    return return_names
