r"""
Module of tools for working with Tide Gauge, DART Buoy, and related data sets.

:Functions:

 - get_coops_gauge
 - timestr2num
 - check_missing
 - read_tide_gauge
 - read_harmonic_constituents
 - dominant_constituents
 - make_tide
 - fit_tide_harmonic
 - get_periods
 - fit_tide_poly

"""

import numpy 
from time import mktime
import dateutil.parser
import datetime

from clawpack.geoclaw.util import get_remote_file

def get_coops_gauge(stnid, bdate, edate, btime="0000", etime="2359",
            output_dir='.', file_name=None, verbose=False):
    r"""

    Attempt to download file from CO-OPS website.
    Note: Only works for 5 days (?) of data or less.

    See http://tidesandcurrents.noaa.gov/1mindata.html
    for a list of stations

    :Example:

        stnid = 9419750      # Crescent City
        bdate = "20110311"
        edate = "20110316"

    """

    if file_name is None:
        file_name = "%s_%s_%s_to_%s_%s.csv" \
                % (stnid, bdate, btime, edate, etime)

    url = "http://tidesandcurrents.noaa.gov/cgi-bin/tsunamicsv.cgi?&tmpname=&stnid=%s&bdate=%s&edate=%s&bdatetime=%s&edatetime=%s" \
            % (stnid, bdate, edate, btime, etime)

    
    get_remote_file(url, output_dir=output_dir, file_name=file_name,
                    verbose=verbose)

    print "Attempted to download to "
    print "    %s/%s" % (output_dir, file_name)



def timestr2num(s):
    hours = int(s[:2])
    minutes_past_hour = int(s[3:5])
    minutes = hours*60 + minutes_past_hour
    return minutes

def check_missing(s):
    if s in ['-','',' ']:
        return numpy.nan
        #return 0.
    else:
        try:
            return float(s)
        except:
            print "cannot convert s = ",s


def read_tide_gauge(fname, t0_datetime=None):

    r""" 
    Read tide gauge file of the form:
        DATE,TIME,1MIN,6MIN,ALTERNATE,PREDICTED,RESIDUAL
        03/11/2011,02:29,0.38,-,-,-,-
        03/11/2011,02:30,0.40,0.39,-,0.40,0.00
        03/11/2011,02:31,0.38,-,-,0.40,-0.02
        etc.

    :Inputs:

      - *fname* (str) -  e.g. '1612340_Honolulu_2011-03-11_to_2011-03-11.csv'
      - *t0_datetime* - a *datetime.datetime* object specifying start time
        for resulting time series (output array *tsec*). 
        
        Or a string that can be converted to a *datetime.datetime* using
        *dateutil.parser.parse*,

        For example, to output seconds since the Tohoku event:

           t0_datetime = "05:46:24 UTC on March 11, 2011"

    :Outputs:

      - *t* - numpy array of *datetime.datetime* objects
      - *tsec* - numpy array of times in seconds.  
        If *t0_datetime* is specified, the time is relative to that start time
        Otherwise, it is seconds since "the epoch" so you might want to
        shift by *tsec[0]* before further processing.
      - *obs_1min* - Observations at 1 minute intervals
      - *predicted* - Tidal constituent predictions (may only be at 6-minutes)
      - *residual* - Observations minus predicted
        

    """

    from pylab import datestr2num, num2date

    converters = {}
    converters[0] = datestr2num
    converters[1] = timestr2num
    converters[2] = check_missing
    converters[3] = check_missing
    converters[4] = check_missing
    converters[5] = check_missing
    converters[6] = check_missing

    d = numpy.loadtxt(fname, converters=converters, delimiter=',', skiprows=1)
    days = d[:,0] + d[:,1]/(24*60.)
    t_datetime = num2date(days)
    obs_1min = d[:,2]
    alternate = d[:,4]
    predicted = d[:,5]
    residual = d[:,6]

    tsec = numpy.array([mktime(tj.timetuple()) for tj in t_datetime])

    if type(t0_datetime) == str:
        try:
            t0_datetime = dateutil.parser.parse(t0_datetime)
        except:
            print "*** Could not convert t0_datetime: ignoring"

    if type(t0_datetime) == datetime.datetime:
        t0sec = mktime(t0_datetime.timetuple())
        tsec = tsec - t0sec

    return t_datetime,tsec,obs_1min,predicted,residual


def read_harmonic_constituents(path, skiprows=None, verbose=True):
    r"""
    Read harmonic constituents from a file that possibly contains a header
    and then has lines of the form

        1   M2  0.716   339.1   28.9841042  Principal lunar semidiurnal
        2   S2  0.183   351.7   30.0    Principal solar semidiurnal
        etc.

    The first *skiprow* lines are ignored.

    Such a file can be generated by cut and paste from a website such as
    http://tidesandcurrents.noaa.gov/harcon.html?id=9419750.

    For a list of stations, see
    http://tidesandcurrents.noaa.gov/stations.html?type=Harmonic+Constituents

    :Input:

     - *path* (str) - path to file
     - *skiprows* (int) - number of rows to ignore.
       If *skiprows* is None, will attempt to figure out assuming M2 is first
     - *verbose* (bool) - If True, prints out a table of constituents

    :Output:

     - *amplitude* - Dictionary of amplitudes, keys are 'M2', 'S2', etc.
     - *phase* - Dictionary of phases (in degrees), keys are 'M2', 'S2', etc.
     """

    periods = get_periods()
    
    lines = open(path).readlines()
    
    if skiprows is None:
        # attempt to figure it out, assuming first line of data is for 'M2'
        skiprows = 0
        while 'M2' not in lines[skiprows]:
            skiprows = skiprows+1
        print "Ignoring first %s lines of file" % skiprows

    amplitude = {}
    phase = {}
    if verbose:
        print "%s  %6s  %7s  %13s  %13s" \
            % (' name','Amplitude','Phase','speed','period')

    for line in lines[skiprows:]:
        tokens = line.split()
        name = tokens[1]
        amplitude[name] = float(tokens[2])
        phase[name] = float(tokens[3])
        if verbose:
            print "%s  %6.3f  %10.2f  %13.7f  %13.7f" \
               % (name.ljust(5),amplitude[name],phase[name],
                 360./periods[name],periods[name])
    return amplitude,phase

def dominant_constituents(amplitude,tol=0.):

    r"""
    :Input:
    
     - *amplitude* (dict): dictionary of amplitudes with harmonic
       constituent names as keys.
     - *tol* (float): tolerance

    :Output:
     - *dominant_periods* (dict): dictionary of periods with
       harmonic constituent names as keys that only contains those
       for which *amplitude[name] > tol*.
    """

    periods = get_periods()
    dominant_periods = {}
    for name in amplitude.keys():
        if amplitude[name] > tol:
            dominant_periods[name] = periods[name]
    return dominant_periods
            

def make_tide(t,amplitude,phase,offset=0,t0=0.):
    r"""
    Make a time series of surface elevation given tidal constituents.
    
    :Inputs:
    
     - *t* (ndarray) - times (in hours)
     - *amplitude* (dict) - dictionary of amplitudes, keys should be names of 
       constituents, e.g. 'M2', 'S2', etc.
     - *phase* (dict) - dictionary of phases (in degrees), same keys as *amplitude*.
     - *offset* (float) - constant term added to cosine series
     - *t0* (float) - Initial time for phase shift.
     
    :Output:
    
     - *eta* (ndarray) - surface elevation relative to MSL at times *t*.
    
    """
    
    periods = get_periods()
    eta = numpy.zeros(t.shape) + offset

    for k in amplitude.keys():
        eta = eta + amplitude[k] * numpy.cos(2.*numpy.pi*((t-t0)/periods[k] \
                + phase[k]/360.))
        
    return eta


def get_periods():
    """
    Returns dictionary of tidal harmonic constituent periods (in hours).
    """
    
    periods = { \
        'K1': 23.9344697,
        'O1': 25.8193417,
        'M2': 12.4206012,
        'S2': 12.0000000,
        'M3': 08.2804008,
        'M4': 06.2103006,
        '2MK5': 04.9308802,
        'M6': 04.1402004,
        '3MK7': 03.10515030,
        'M8': 03.1051503,
        'N2': 12.6583482,
        'Q1': 26.8683567,
        'MK3': 08.1771399,
        'S4': 06.0000000,
        'MN4': 06.2691739,
        'NU2': 12.6260044,
        'S6': 04.0000000,
        'MU2': 12.8717576,
        '2N2': 12.9053745,
        'OO1': 22.3060742,
        'LAM2': 12.2217742,
        'S1': 24.0000000,
        'M1': 24.8332484,
        'J1': 23.0984768,
        'MM': 661.3092049,
        'SSA': 4382.9052087,
        'SA': 8765.8210896,
        'MSF': 354.3670522,
        'MF': 327.8589689,
        'RHO': 26.7230533,
        'T2': 12.0164492,
        'R2': 11.9835958,
        '2Q1': 28.0062225,
        'P1': 24.0658902,
        '2SM2': 11.6069516,
        'L2': 12.1916202,
        '2MK3': 08.3863030,
        'K2': 11.9672348,
        'MS4': 06.1033393,
        }
    return periods


periods = get_periods()

# For harbors in Hawaii these appear to be dominant periods 
# (Amplitude less than 1% of these in Honolulu, Kahului, Hilo)
constituents_hawaii = ['J1','K1','K2','M2','N2','O1','P1','Q1','S2','SA']
periods_hawaii = {k:periods[k] for k in constituents_hawaii}


def fit_tide_poly(t,eta,degree):
    """
    Fit a polynomial of the specified degree to data 
    Returns the coefficents c of c[0] + c[1]*t + ...
    and the polynomial fit eta_fit.
    """
    from numpy.linalg import lstsq, svd
    from pylab import find
    

    if numpy.any(numpy.isnan(eta)):
        eta2 = numpy.where(numpy.isnan(eta), 1e50, eta)
        j_nonan = find(eta2 < 1e40)
        t_nonan = t[j_nonan]
        eta_nonan = eta[j_nonan]
        print "Ignoring %i NaN values" % (len(eta)-len(eta_nonan))
    else:
        t_nonan = t
        eta_nonan = eta

    
    # Scale data so matrix better conditioned:

    scale_factor = abs(t_nonan).max()
    t_nonan = t_nonan/scale_factor
    t = t/scale_factor

    
    # Use Newton polynomial basis using these points:
    tpts = numpy.linspace(t_nonan.min(),t_nonan.max(),degree+1)
    
    # Form A matrix for least squares fit 
    A = numpy.ones((len(t_nonan),degree+1))
    for j in range(1,degree+1):
        A[:,j] = A[:,j-1] * (t_nonan - tpts[j])
    ncols = A.shape[1]
        
    # Perform least squares fit:
    c = lstsq(A,eta_nonan)[0]
    
    if numpy.any(numpy.isnan(eta)):
        # Reconstruct A using all times for calculating predicted:
        A = numpy.ones((len(t),degree+1))
        for j in range(1,degree+1):
            A[:,j] = A[:,j-1] * (t - tpts[j])

    eta_fit = numpy.dot(A,c)
    
    return eta_fit
    


def fit_tide_harmonic(t, eta, periods, t0=0, svd_tol=0.01):
    from numpy.linalg import lstsq, svd, norm
    from pylab import find
    
    if numpy.any(numpy.isnan(eta)):
        eta2 = numpy.where(numpy.isnan(eta), 1e50, eta)
        j_nonan = find(eta2 < 1e40)
        t_nonan = t[j_nonan]
        eta_nonan = eta[j_nonan]
    else:
        t_nonan = t
        eta_nonan = eta

    A = numpy.ones(t_nonan.shape)
    names = periods.keys()
    for k in names:
        s1 = numpy.sin(2*numpy.pi*t_nonan/periods[k])
        c1 = numpy.cos(2*numpy.pi*t_nonan/periods[k])
        A = numpy.vstack([A,s1,c1])
    A = A.T
    ncols = A.shape[1]
    
    # c,res,rank,s = lstsq(A,eta)  ## Does not work well!
    # Using full least squares solution gives very large coefficients
    # Instead use SVD based pseudo-inverse throwing away singular
    # values below the specified tolerance:

    U,S,V = svd(A, full_matrices=False)
    #print "Singular values: ",S
    
    c = numpy.zeros(ncols)
    num_sv = 0
    for k in range(ncols):
        if S[k]/S[0] > svd_tol:
            c = c + (numpy.dot(U[:,k],eta_nonan) / S[k]) * V[k,:]
            num_sv += 1
    print "Inverting using %s singular values out of %s" % (num_sv, ncols)
    
    c_sin = c[1::2]
    c_cos = c[2::2]
    c_cos = numpy.where(abs(c_cos)<1e-10, 1e-10, c_cos)
    phi = -numpy.arctan(c_sin / c_cos) * 180./numpy.pi
    phi = numpy.where(c_cos < 0., phi+180, phi)
    

    # Determine offset, amplitude, phase so that fit has the form
    #  eta_fit = offset + sum_k amplitude[k] * cos(2*pi*(t - t0) + phase[k])
    # where the sum is over all harmonic constituents in periods.keys()

    offset = c[0]  # constant term in fit
    phase = {}
    amplitude = {}
    for i,name in enumerate(names):
        amplitude[name] = numpy.sqrt(c_sin[i]**2 + c_cos[i]**2)
        phase[name] = phi[i] + 360.*t0/periods[name]

    #print "c: ",c
    #print "c_cos: ",c_cos
    #print "c_sin: ",c_sin
    #print "+++ offset, amplitude, phase: ",offset, amplitude, phase
        
    if numpy.any(numpy.isnan(eta)):
        # Reconstruct A using all times for calculating predicted:
        A = numpy.ones(t.shape)
        for k in names:
            s1 = numpy.sin(2*numpy.pi*t/periods[k])
            c1 = numpy.cos(2*numpy.pi*t/periods[k])
            A = numpy.vstack([A,s1,c1])
        A = A.T

    eta_fit = numpy.dot(A,c)

    # residual = eta - eta_fit
    # print "Norm of residual: ",norm(residual,2)  # might return NaN
    # print "Norm of amplitudes: ",norm(c[1:],2)
    
    return eta_fit, amplitude, phase, offset
    
