#!/usr/bin/env python
"""
Created by Josh Laughner

(Originally named egi_name.py)
"""

from argparse import ArgumentParser
from datetime import datetime as dtime
import sys


def format_lon(lon, prec=2, zero_pad=False):
    """
    Convert longitude between string and numeric representations.

    If ``lon`` is a number, then it is converted to a string. The string will be the absolute value with "W" or "E" at
    the end to indicate west (<= 0) or east (> 0). If given a string in that format, it converts it to a number.

    :param lon: the longitude to convert
    :type lon: float or str

    :param prec: the precision after the decimal point to use. Only has an effect when converting float to string.
    :type prec: int

    :param zero_pad: set to ``True`` to zero pad the longitude string so that there are 3 digits before the decimal
     place. Only has an effect when converting float to string.

    :return: the formatted longitude string or the float representation of the longitude, with west being negative.
    :rtype: str or float
    """
    def to_str(lon):
        ew = 'E' if lon >= 0 else 'W'
        # In Python float format specification, "0X.Y" means to zero pad so that there's X total characters and Y after
        # the decimal point. We want lon zero padded to have three numbers before the decimal point, so the total width
        # needs to be the precision + 4 if there will be a decimal point otherwise just 3.
        width = prec + 4 if prec > 0 else prec + 3
        pad = '0{}'.format(width) if zero_pad else ''
        fmt_str = '{{:{padding}.{prec}f}}{{}}'.format(padding=pad, prec=prec)
        return fmt_str.format(abs(lon), ew)

    def to_float(lon):
        if lon[-1] == 'E':
            sign = 1
        elif lon[-1] == 'W':
            sign = -1
        else:
            raise ValueError('A longitude string must end in "E" or "W"')

        return float(lon[:-1]) * sign

    if isinstance(lon, str):
        return to_float(lon)
    else:
        if lon > 180:
            lon -= 360
        return to_str(lon)


def format_lat(lat, prec=2, zero_pad=False):
    """
    Convert latitude between string and numeric representations.

    If ``lat`` is a number, then it is converted to a string. The string will be the absolute value with "N" or "S" at
    the end to indicate south (<= 0) or north (> 0). If given a string in that format, it converts it to a number.

    :param lat: the latitude to convert
    :type lat: float or str

    :param prec: the precision after the decimal point to use. Only has an effect when converting float to string.
    :type prec: int

    :param zero_pad: set to ``True`` to zero pad the latitude string so that there are 2 digits before the decimal
     place. Only has an effect when converting float to string.

    :return: the formatted latitude string or the float representation of the latitude, with south being negative.
    :rtype: str or float
    """
    def to_str(lat):
        ns = 'N' if lat >= 0 else 'S'
        # In Python float format specification, "0X.Y" means to zero pad so that there's X total characters and Y after
        # the decimal point. We want lat zero padded to have two numbers before the decimal point, so the total width
        # needs to be the precision + 3 if there will be a decimal point otherwise just 2.
        width = prec + 3 if prec > 0 else prec + 2
        pad = '0{}'.format(width) if zero_pad else ''
        fmt_str = '{{:{padding}.{prec}f}}{{}}'.format(padding=pad, prec=prec)
        return fmt_str.format(abs(lat), ns)

    def to_float(lat):
        if lat[-1] == 'N':
            sign = 1
        elif lat[-1] == 'S':
            sign = -1
        else:
            raise ValueError('A latitude string must end in "N" or "S"')

        return float(lat[:-1]) * sign

    if isinstance(lat, str):
        return to_float(lat)
    else:
        return to_str(lat)


def make_egi_tar_file_name(sites, lats, lons, start_date):
    """
    Create the expected .tar file name for an EGI priors .tar file.

    Note that this will just be the filename, not the FTP address or directory.

    :param sites: a single site ID or a sequence (list or tuple) of site IDs.
    :type sites: str or Sequence[str]

    :param lats: a single latitude (in the range -180 to 360) or a list/tuple of latitudes
    :type lats: float or Sequence[float]

    :param lons: a single longitude (in the range -180 to 360) or a list/tuple of longitudes
    :type lons: float or Sequence[float]

    :param start_date: the first date of data requested
    :type start_date: :class:`datetime.datetime`

    :return: the tar filename
    :rtype: str
    """
    def make_seq(v):
        if isinstance(v, (tuple, list)):
            return v, len(v)
        else:
            return [v], 1

    sites, nsites = make_seq(sites)
    lons, nlons = make_seq(lons)
    lats, nlats = make_seq(lats)

    if nsites != nlons or nsites != nlats:
        print('Mismatch in number of sites ({ns}), lons ({nlons}), and/or lats ({nlats}). '
              'The shortest number will define how many are included in the file name.'
              .format(ns=nsites, nlons=nlons, nlats=nlats),
              file=sys.stderr)

    locs = []
    for site, lat, lon in zip(sites, lats, lons):
        locs.append('{}{}{}'.format(site, format_lat(lat, prec=0), format_lon(lon, prec=0)))
    return 'EGI_{loc}_{date}.tgz'.format(loc='_'.join(locs), date=start_date.strftime('%Y%m%d'))


def parse_args():
    def cslstr(v):
        return tuple(s.strip() for s in v.split(','))

    def cslfloat(v):
        return tuple(float(f) for f in v.split(','))

    def datetime(v):
        return dtime.strptime(v, '%Y%m%d')

    # the prefix_chars argument is necessary to avoid intepreting negative lat/lons as flags/options
    p = ArgumentParser(description='Generate the expected name for an EGI priors tarball', prefix_chars='!')
    p.add_argument('sites', type=cslstr, help='Site or comma-separated list of sites requested in the tarball')
    p.add_argument('lats', type=cslfloat, help='Latitude or comma-separated list of latitudes corresponding to the '
                                               'sites. South must be given as negative.')
    p.add_argument('lons', type=cslfloat, help='Longitude or comma-separated list of longitudes corresponding to the '
                                               'sites. West must be given as negative.')
    p.add_argument('start_date', type=datetime, help='The start date requested in the input file, in YYYYMMDD format.')

    args = vars(p.parse_args())
    return args


def main():
    args = parse_args()
    print(make_egi_tar_file_name(**args))


if __name__ == '__main__':
    main()
