"""This module contains a 'pure Python' implementation of the version 2.x 
HLSVDPRO package. 
The HLSVDPROPY package computes a 'sum of lorentzians' model for the complex 
'signals' data passed in via a 'black box' state space approach. We are using 
the scipy.linalg SVD libraries, rather than the PROPACK Fortran libraries used
in the HLSVDPRO package, but the algorithm is otherwise similarly based on the 
algorithm in:
Laudadio T, et.al. "Improved Lanczos algorithms for blackbox MRS data 
quantitation", Journal of Magnetic Resonance, Volume 157, p.292-297, 2002
Functions:
    hlsvd(data, nsv_sought, dwell_time, sparse=False) -> 6-tuple
    hlsvdpro(data, nsv_sought, m=None, sparse=True) -> 8-tuple
    convert_hlsvd_result(result, dwell)
    create_hlsvd_fids(result, npts, dwell, sum_results=False, convert=True)
    get_testdata()
    
Example:
    $ python hlvsd.py
    
    Running this module from the site-packages directory will run the internal 
    example. It requires matplotlib be already installed.
"""      

# Python modules
from __future__ import division
import math

# 3rd party modules
import numpy as np
import scipy.linalg
import scipy.sparse.linalg
import scipy.linalg.lapack as lapack


def hlsvd(data, nsv_sought, dwell_time):
    """
    This calls HLSVDPRO version 2.x code, but simulates the hlsvd.hlsvd()
    call from HLSVDPRO version 1.0.x to maintain the API. See doc string
    below for hlsvdpro() method
    Args:
        data (ndarray): an iterable of complex numbers.
        nsv_sought (int): The number of singular values sought. The function
            will return a maximum of this many singular values.
        dwell_time (float): Dwell time in milliseconds.
    Returns:
        tuple: a 6-tuple containing -
            (int), number of singular values found (nsv_found <= nsv_sought)
            (ndarray, floats) the singular values
            (ndarray, floats) the frequencies (in kilohertz)
            (ndarray, floats) the damping factors (in milliseconds?)
            (ndarray, floats) the amplitudes (in arbitrary units)
            (ndarray, floats) the phases (in degrees)
            Each list's length == nsv_found. The five lists are correlated
            (element N of one list is associated with element N in the other
            four lists) and are sorted by singular value with the largest
            (strongest signal) first. The frequencies and damping factors
            HAVE been adjusted by the dwell time, and phases converted to
            degrees for compatibility with HLSVDPRO version 1.x API.
    """

    m = len(data) // 2
    r = hlsvdpro(data, nsv_sought, m=m)
    r = convert_hlsvd_result(r, dwell_time)

    nsv_found, singular_values, frequencies, damping_factors, amplitudes, phases = r[0:6]

    return (nsv_found, singular_values, frequencies, damping_factors, amplitudes, phases)


    
def hlsvdpro(data, nsv_sought, m=None, sparse=False):
    """A pure Python implementation of the HLSVDPRO version 2.x package.
    
    Computes a 'sum of lorentzians' model for the complex 'data' passed in 
    using the scipy.linalg SVD libraries and based on the algorithm in:
      
    Laudadio T, et.al. "Improved Lanczos algorithms for blackbox MRS data 
    quantitation", Journal of Magnetic Resonance, Volume 157, p.292-297, 2002
    Args:
        data (list, ndarray): an iterable of complex numbers.
        nsv_sought (int): The number of singular values sought. The function 
            will return a maximum of this many singular values.
        m (int): (optional) default=len(data)/2, Use to set the size of
            the Hankel matrix used to compute the singular values. Hankel 
            matrix shape is (L+1,L) where L = len(data)-m-1
        sparse (bool): (optional) default False. If set to True, the  
            scipy.sparse.linalg.svds() is used to calculate singular values and 
            nsv_sought is passed in as a parameter. If False, scipy.linalg.svd()
            is used to calculate the singular values, and nsv_sought is used to 
            truncate the results returned.
    Returns:
        tuple: an 8-tuple containing -
            (int), number of singular values found (nsv_found <= nsv_sought)
            (ndarray, floats) the singular values
            (ndarray, floats) the frequencies (in kilohertz)
            (ndarray, floats) the damping factors 
            (ndarray, floats) the amplitudes (in arbitrary units)
            (ndarray, floats) the phases (in radians)
            (ndarray) The top nsv_found left singular vectors, shape=(L+1,3)
            (ndarray) The top nsv_found right singular vectors, shape=(3,L)
            For the first five ndarrays, length == nsv_found and element N of 
            one list is associated with element N in the other four lists. They
            are sorted by singular value with the largest (strongest signal) 
            first.  The frequencies and damping factors have NOT been adjusted 
            by the dwell time, nor the phases converted to degrees.
    """
    mode = 'f' # or 'b'
    xx = data
    k = nsv_sought
    n = len(xx)
    m = int(n/2) if m is None else m
    
    l = n - m - 1

    if mode == "f":
        x = scipy.linalg.hankel(xx[:l + 1], xx[l:])
    else:
        # for backward LP we need to make the hankel matrix:
        #    x_N-1 x_N-2 ... x_N-M-1
        #    x_N-2 x_N-3 ... x_N-M-2
        #      ...
        #    x_M   x_M-1 ... x_0

        x = scipy.linalg.hankel(xx[:m - 1:-1], xx[m::-1])
    
    if sparse:
        u, s, vh = scipy.sparse.linalg.svds(x, k=k)
    else:
        u, s, vh = scipy.linalg.svd(x, full_matrices=False)

    k = min(k,len(s))               # number of singular values found
        
    uk = np.mat(u[:, :k])           # trucated U matrix of rank K
    ub = uk[:-1]                    # Uk with bottom row removed
    ut = uk[1:]                     # Uk with top row removed

    zp, resid, rank, ss = scipy.linalg.lstsq(ub, ut)

    #----------------------------------------------------------------
    # Diagonalize Z' (=hx), yields 'signal'poles' aka 'roots'
    #   Eigenvalues are returned unordered. I sort them here to make
    #   it easier to compare them with output from the Fortran code.

    roots = scipy.linalg.eigvals(zp)
    roots = np.array(sorted(roots, reverse=True))

    #----------------------------------------------------------------
    #  Calculate dampings (damp) and frequencies (freq) from roots

    dampings = np.log(np.abs(roots))
    frequencies = np.arctan2(roots.imag, roots.real) / (math.pi * 2)

    #----------------------------------------------------------------
    #  Calculate complex-valued amplitudes, using the pseudoinverse
    #    of the Lrow*kfit Vandermonde matrix zeta.
    #
    # FIXME this code, like the Fortran, ignores possible errors
    #    reported in the "info" return value.

    zeta = np.vander(roots, N=len(data), increasing=True).T
    v1, x1, s1, rank, _, info = lapack.zgelss(zeta, data)

    #----------------------------------------------------------------
    # Discard uneeded values of x
    x1 = x1[:k]
    amplitudes = np.abs(x1)
    phases = np.arctan2(x1.imag, x1.real)

    return k, s[::-1], frequencies, dampings, amplitudes, phases, u, vh



def convert_hlsvd_result(result, dwell):
    """
    Use dwell time to convert output from hlsvdpro() method to standard units.
    Frequencies convert to [kHz] and damping factors to [ms]. Phases convert
    to [degrees]. Singular values and row and column matrices are maintained
    at their same values and output tuple locations.
    """
    nsv_found, sing_vals, frequencies, damping, ampls, phases = result[0:6]

    damping = np.array([1.0 / df for df in damping])
    damping = np.array([df * dwell for df in damping])
    frequencies = np.array([frequency / dwell for frequency in frequencies])
    phases = np.array([phase * (180.0/3.1415926) for phase in phases])
    return (nsv_found, sing_vals, frequencies, damping, ampls, phases)


def create_hlsvd_fids(result, npts, dwell, sum_results=False, convert=True):
    """ 
    This is a convenience function for creating time domain FID arrays from
    the results of the hlsvd() method. You can either send results in 
    directly from hlsvd() (the assumed condition) or run the conversion
    routine convert_hlsvd_results() manually yourself and send them in, but
    in the latter case set convert=False.
    
    See the self-test example at the bottom of this module for how to
    process the fids into spectra and compare in matplotlib.
    Args:
        result (tuple): output from hlsvd() method, an 8-tuple.
        npts (int): The number of points in the created fids.
        dwell (float): Dwell time in milliseconds for created fids.
        sum_results (bool): (optional) default False. If True, all fid lines
            will be summed up to one fid before being returned. If false, a
            (len(result[2]),npts) ndarray is returned
        convert (bool): (optional) default True. If True, the results parameter
            will be run through the convert_hlsvd_results() method before
            being used. If False, it will assume that the user has already
            converted the freq, damp and phase values.
    Returns:
        ndarray: either a (npts,) ndarray (if sum_results=True) or a
            (len(result[2]),npts) ndarray (if sum_results=False) with time
            domain fids.
    """
    if convert: result = convert_hlsvd_result(result, dwell)

    freqs, damps, areas, phase = result[2:6]

    fids = np.zeros((len(freqs), npts), dtype=complex)
    t = np.arange(npts) * dwell
    k = 1j * 2 * np.pi

    for i, damp in enumerate(damps):
        if damp:
            # hack for very small exp() values, temp disable numpy's error
            # report, it silently generate NaNs, then change those NaNs to 0.0
            old_settings = np.seterr(all='ignore')
            line = areas[i] * np.exp((t / damp) + k * (freqs[i] * t + phase[i] / 360.0))
            zeros = np.zeros_like(line)
            fids[i,:] = np.where(np.isnan(line), zeros, line)
            np.seterr(**old_settings)
        else:
            fids[i,:] = fids[i,:] * 0

    if sum_results: result = np.sum(fids, axis=0)

    return result


def _chop(data):
    return data * ((((np.arange(len(data)) + 1) % 2) * 2) - 1)


if __name__ == "__main__":
    print("In main, should not be here...")