"""
Test suite for Cython utilities.
"""
import numpy as np

def interp_conserve(x, xp, fp, left=0., right=0.):
    """
    Interpolate `xp`,`yp` array to the output x array, conserving flux.  
    `xp` can be irregularly spaced.
    """
    midpoint = (x[1:]-x[:-1])/2.+x[:-1]
    midpoint = np.append(midpoint, np.array([x[0],x[-1]]))
    midpoint = midpoint[np.argsort(midpoint)]
    int_midpoint = np.interp(midpoint, xp, fp, left=left, right=right)
    int_midpoint[midpoint > xp.max()] = 0.
    int_midpoint[midpoint < xp.min()] = 0.
    
    fullx = np.append(xp, midpoint)
    fully = np.append(fp, int_midpoint)
    
    so = np.argsort(fullx)
    fullx, fully = fullx[so], fully[so]
    
    outy = x*0.
    dx = midpoint[1:]-midpoint[:-1]
    for i in range(len(x)):
        bin = (fullx >= midpoint[i]) & (fullx <= midpoint[i+1])
        outy[i] = np.trapz(fully[bin], fullx[bin])/dx[i]
    
    return outy

def test_nmf():
    
    import interp_c
    import numpy as np
    import time
    import matplotlib.pyplot as plt
    
    x = np.arange(0,16,0.5)*np.pi*2
    NFILT = len(x)
    NTEMP = 7
    
    coeffs = np.random.random(NTEMP)
    
    y = x*0.
    templates = np.zeros((NTEMP, NFILT))
    norms = np.zeros(NTEMP)
    for i in range(0,NTEMP):
        templates[i,:] = np.sin(x/((i+1)*1.6))+1.1
        norms[i] = np.trapz(templates[i,:], x)
        
    y = np.dot(coeffs.reshape((1,NTEMP)), templates).flatten()
    
    err = y*0.+np.median(y)/10.
    
    yerr = y+np.random.normal(size=NFILT)*err
    
    t0 = time.time()
    
    amatrix = interp_c.prepare_nmf_amatrix(err**2, templates)
    t1= time.time()
    coeffs_fit = interp_c.run_nmf(yerr, err**2, templates, amatrix, toler=1.e-5)
    t2 = time.time()
    
    print 'Prepare: %.4f' %(t1-t0)
    print 'Fit    : %.4f' %(t2-t1)
    
    yfit = np.dot(coeffs_fit.reshape((1,-1)), templates).flatten()
    
    fig = plt.figure()
    
    ax = fig.add_subplot(311)
    ax.plot(x, y/np.median(y), color='blue')
    ax.errorbar(x, yerr/np.median(y), err/np.median(y), color='blue', marker='o', linestyle='None')
    ax.plot(x, yfit/np.median(y), color='red')
    ax.set_ylabel('"Flux"')
        
    ax = fig.add_subplot(312)
    ax.plot(x, y-yfit, color='blue')
    ax.errorbar(x, yerr-yfit, err, color='blue', marker='o', linestyle='None')
    ax.plot(x, yfit-yfit, color='red')
    ax.set_ylabel(r'$\Delta$(obs-fit)')
    chi2 = np.sum((yerr-yfit)**2/err**2)/(len(x)-1)
    ax.text(0.1,0.8,r'$\chi^2_\nu$=%.3f' %(chi2), transform=ax.transAxes)
    
    ax = fig.add_subplot(313)
    ax.plot(np.log10(coeffs/coeffs_fit), color='orange')
    ax.set_ylabel(r'$\log$($\Delta$coeff)')
    
    #### Is sum of normalizations conserved?
    norm_input = np.sum(norms*coeffs)
    norm_fit = np.sum(norms*coeffs_fit)
    int_fit = np.trapz(yfit, x)
    print 'Norm_in: %.2f, Norm_fit: %.2f, trapz_fit: %.2f' %(norm_input, norm_fit, int_fit)
    
    fig.savefig('/tmp/nmf.png')
    
def test():

    import interp_c
    import time
    import scipy
    import threedhst
    import numpy as np
    
    N = int(1.e6)
    
    xfull = np.arange(0,N+1,1)*1.