"""
Import functions for NOAA SURFRAD Data.
"""
import io
from urllib.request import urlopen, Request
import pandas as pd
import numpy as np

SURFRAD_COLUMNS = [
    'year', 'jday', 'month', 'day', 'hour', 'minute', 'dt', 'zen',
    'dw_solar', 'dw_solar_flag', 'uw_solar', 'uw_solar_flag', 'direct_n',
    'direct_n_flag', 'diffuse', 'diffuse_flag', 'dw_ir', 'dw_ir_flag',
    'dw_casetemp', 'dw_casetemp_flag', 'dw_dometemp', 'dw_dometemp_flag',
    'uw_ir', 'uw_ir_flag', 'uw_casetemp', 'uw_casetemp_flag', 'uw_dometemp',
    'uw_dometemp_flag', 'uvb', 'uvb_flag', 'par', 'par_flag', 'netsolar',
    'netsolar_flag', 'netir', 'netir_flag', 'totalnet', 'totalnet_flag',
    'temp', 'temp_flag', 'rh', 'rh_flag', 'windspd', 'windspd_flag',
    'winddir', 'winddir_flag', 'pressure', 'pressure_flag']

# Dictionary mapping surfrad variables to pvlib names
VARIABLE_MAP = {
    'zen': 'solar_zenith',
    'dw_solar': 'ghi',
    'dw_solar_flag': 'ghi_flag',
    'direct_n': 'dni',
    'direct_n_flag': 'dni_flag',
    'diffuse': 'dhi',
    'diffuse_flag': 'dhi_flag',
    'temp': 'temp_air',
    'temp_flag': 'temp_air_flag',
    'windspd': 'wind_speed',
    'windspd_flag': 'wind_speed_flag',
    'winddir': 'wind_direction',
    'winddir_flag': 'wind_direction_flag',
    'rh': 'relative_humidity',
    'rh_flag': 'relative_humidity_flag'
}


def read_surfrad(filename, map_variables=True):
    """Read in a daily NOAA SURFRAD file.  The SURFRAD network is
    described in [1]_.

    Parameters
    ----------
    filename: str
        Filepath or URL. URL can be either FTP or HTTP.
    map_variables: bool
        When true, renames columns of the Dataframe to pvlib variable names
        where applicable. See variable :const:`VARIABLE_MAP`.

    Returns
    -------
    Tuple of the form (data, metadata).

    data: Dataframe
        Dataframe with the fields found below.
    metadata: dict
        Site metadata included in the file.

    Notes
    -----
    Metadata dictionary includes the following fields:

    ===============  ======  ===============
    Key              Format  Description
    ===============  ======  ===============
    station          String  site name
    latitude         Float   site latitude
    longitude        Float   site longitude
    elevation        Int     site elevation
    surfrad_version  Int     surfrad version
    tz               String  Timezone (UTC)
    ===============  ======  ===============

    Dataframe includes the following fields:

    =======================  ======  ==========================================
    raw, mapped              Format  Description
    =======================  ======  ==========================================
    **Mapped field names are returned when the map_variables argument is True**
    ---------------------------------------------------------------------------
    year                     int     year as 4 digit int
    jday                     int     day of year 1-365(or 366)
    month                    int     month (1-12)
    day                      int     day of month(1-31)
    hour                     int     hour (0-23)
    minute                   int     minute (0-59)
    dt                       float   decimal time i.e. 23.5 = 2330
    zen, solar_zenith        float   solar zenith angle (deg)
    **Fields below have associated qc flags labeled <field>_flag.**
    ---------------------------------------------------------------------------
    dw_solar, ghi            float   downwelling global solar(W/m^2)
    uw_solar                 float   updownwelling global solar(W/m^2)
    direct_n, dni            float   direct normal solar (W/m^2)
    diffuse, dhi             float   downwelling diffuse solar (W/m^2)
    dw_ir                    float   downwelling thermal infrared (W/m^2)
    dw_casetemp              float   downwelling IR case temp (K)
    dw_dometemp              float   downwelling IR dome temp (K)
    uw_ir                    float   upwelling thermal infrared (W/m^2)
    uw_casetemp              float   upwelling IR case temp (K)
    uw_dometemp              float   upwelling IR case temp (K)
    uvb                      float   global uvb (miliWatts/m^2)
    par                      float   photosynthetically active radiation(W/m^2)
    netsolar                 float   net solar (dw_solar - uw_solar) (W/m^2)
    netir                    float   net infrared (dw_ir - uw_ir) (W/m^2)
    totalnet                 float   net radiation (netsolar+netir) (W/m^2)
    temp, temp_air           float   10-meter air temperature (?C)
    rh, relative_humidity    float   relative humidity (%)
    windspd, wind_speed      float   wind speed (m/s)
    winddir, wind_direction  float   wind direction (deg, clockwise from north)
    pressure                 float   station pressure (mb)
    =======================  ======  ==========================================

    See README files located in the station directories in the SURFRAD
    data archives [2]_ for details on SURFRAD daily data files. In addition to
    the FTP server, the SURFRAD files are also available via HTTP access [3]_.

    References
    ----------
    .. [1] NOAA Earth System Research Laboratory Surface Radiation Budget
       Network
       `SURFRAD Homepage <https://www.esrl.noaa.gov/gmd/grad/surfrad/>`_
    .. [2] NOAA SURFRAD Data Archive
       `SURFRAD Archive <ftp://aftp.cmdl.noaa.gov/data/radiation/surfrad/>`_
    .. [3] `NOAA SURFRAD HTTP Index
       <https://gml.noaa.gov/aftp/data/radiation/surfrad/>`_
    """
    if str(filename).startswith('ftp') or str(filename).startswith('http'):
        req = Request(filename)
        response = urlopen(req)
        file_buffer = io.StringIO(response.read().decode(errors='ignore'))
    else:
        file_buffer = open(str(filename), 'r')

    # Read and parse the first two lines to build the metadata dict.
    station = file_buffer.readline()
    file_metadata = file_buffer.readline()

    metadata_list = file_metadata.split()
    metadata = {}
    metadata['name'] = station.strip()
    metadata['latitude'] = float(metadata_list[0])
    metadata['longitude'] = float(metadata_list[1])
    metadata['elevation'] = float(metadata_list[2])
    metadata['surfrad_version'] = int(metadata_list[-1])
    metadata['tz'] = 'UTC'

    data = pd.read_csv(file_buffer, delim_whitespace=True,
                       header=None, names=SURFRAD_COLUMNS)
    file_buffer.close()

    data = format_index(data)
    missing = data == -9999.9
    data = data.where(~missing, np.NaN)

    if map_variables:
        data.rename(columns=VARIABLE_MAP, inplace=True)
    return data, metadata


def format_index(data):
    """Create UTC localized DatetimeIndex for the dataframe.

    Parameters
    ----------
    data: Dataframe
        Must contain columns 'year', 'jday', 'hour' and
        'minute'.

    Return
    ------
    data: Dataframe
        Dataframe with a DatetimeIndex localized to UTC.
    """
    year = data.year.apply(str)
    jday = data.jday.apply(lambda x: '{:03d}'.format(x))
    hours = data.hour.apply(lambda x: '{:02d}'.format(x))
    minutes = data.minute.apply(lambda x: '{:02d}'.format(x))
    index = pd.to_datetime(year + jday + hours + minutes, format="%Y%j%H%M")
    data.index = index
    data = data.tz_localize('UTC')
    return data
