import numpy as np
import numpy.polynomial.chebyshev as cb
import h5py as h5

def read_trad (filename) :

    # Read the data

    f = h5.File(filename, 'r')

    cb_pos = read_cb(f, 'cb_pos')
    cb_neg = read_cb(f, 'cb_neg')
    cb_ctr = read_cb(f, 'cb_ctr')

    m = f.attrs['m']
    k = f.attrs['k']

    f.close()

    # Return it

    return {'m'      : m,
            'k'      : k,
            'cb_pos' : cb_pos,
            'cb_neg' : cb_neg,
            'cb_ctr' : cb_ctr}

#

def read_cb (f, name) :

    # Read the Chebyshev data from the named group

    x_a = f[name].attrs['x_a']
    x_b = f[name].attrs['x_b']

    c = f[name]['c'][...]

    c[0] *= 0.5

    return cb.Chebyshev(c, domain=[x_a,x_b])

#

def lambda_tr (tr, nu) :

    # Evaluate lambda at nu

    if nu < -1. :
        lambda_tr = tr['cb_neg'](1./nu)*lambda_norm(nu, tr['m'], tr['k'])
    elif nu > 1. :
        lambda_tr = tr['cb_pos'](1./nu)*lambda_norm(nu, tr['m'], tr['k'])
    else :
        lambda_tr = tr['cb_ctr'](nu)*lambda_norm(nu, tr['m'], tr['k'])

    return lambda_tr


#

def lambda_norm (nu, m, k) :

    # Evaluate the lambda scale

    if nu < -1 :
        lambda_norm = lambda_asymp(nu, m, k)
    elif nu > 1. :
        lambda_norm = lambda_asymp(nu, m, k)
    else :
        l = np.abs(m) + 1
        lambda_norm = l*(l+1)

    return lambda_norm

#

def lambda_asymp (nu, m, k) :

    # Evaluate the lambda using asymptotic expressions

    if m*nu >= 0 :

        if k > 0 :
            s = k - 1
            lambda_asymp = m*nu + m**2 + 0.5*nu**2*(2*s + 1)**2*(1. + np.sqrt(1. + 4.*(m*nu + m**2)/(nu**2*(2*s + 1)**2)))
        else :
            lambda_asymp = m**2*(2.*m*nu)/(2.*m*nu - 1.)
           
    else :

        if k >= -1 :
            s = k + 1
            lambda_asymp = m*nu + m**2 + 0.5*nu**2*(2*s + 1)**2*(1. + np.sqrt(1. + 4.*(m*nu + m**2)/(nu**2*(2*s + 1)**2)))
        else :
            s = -k -1
            lambda_asymp = (m*nu - m**2)**2/(nu**2*(2*s+1)**2)

    return lambda_asymp

#

    

    

    
