import numpy as np
import scipy.optimize as sco
import scipy.odr
from scipy.stats import chi2
import scipy.fftpack

def linreg_xy(x,y,ex,ey):
    ex = np.array(ex)
    ey = np.array(ey)
    x = np.array(x)
    y = np.array(y)
    if ex.shape[0] == 1:
        temp = np.zeros(x.shape)
        temp[:] = ex
        ex = temp
    if ey.shape[0] == 1:
        temp = np.zeros(y.shape)
        temp[:] = ey
        ey = temp

    a_ini,ea_ini,b_ini,eb_ini,chiq_ini,corr_ini = linreg(x,y,ey)

    model  = scipy.odr.Model(lin)
    data   = scipy.odr.RealData(x, y, sx=ex, sy=ey)
    odr    = scipy.odr.ODR(data, model, beta0=[a_ini, b_ini])
    output = odr.run()
    ndof = x.shape[0]-2
    #chiq = output.res_var*ndof
    #corr = output.cov_beta[0,1]/np.sqrt(output.cov_beta[0,0]*output.cov_beta[1,1])
    chiq, p = chiSq(x, y, np.sqrt(ey**2 + (output.beta[0]*ex)**2), [output.beta[0], output.beta[1]])

    # Es scheint, dass die sd_beta von ODR auf ein chi2/dof=1 reskaliert werden. Deshalb nehmen
    # wir den Fehler direkt aus der Kovarianzmatrix.
    return output.beta[0], np.sqrt(output.cov_beta[0,0]), output.beta[1], np.sqrt(output.cov_beta[1,1]), chiq, p

def linreg_xy_cov(x,y,ex,ey):
    ex = np.array(ex)
    ey = np.array(ey)
    x = np.array(x)
    y = np.array(y)
    if ex.shape[0] == 1:
        temp = np.zeros(x.shape)
        temp[:] = ex
        ex = temp
    if ey.shape[0] == 1:
        temp = np.zeros(y.shape)
        temp[:] = ey
        ey = temp

    a_ini,ea_ini,b_ini,eb_ini,chiq_ini,corr_ini = linreg(x,y,ey)

    model  = scipy.odr.Model(lin)
    data   = scipy.odr.RealData(x, y, sx=ex, sy=ey)
    odr    = scipy.odr.ODR(data, model, beta0=[a_ini, b_ini])
    output = odr.run()
    ndof = x.shape[0]-2
    #chiq = output.res_var*ndof
    #corr = output.cov_beta[0,1]/np.sqrt(output.cov_beta[0,0]*output.cov_beta[1,1])
    chiq, p = chiSq(x, y, np.sqrt(ey**2 + (output.beta[0]*ex)**2), [output.beta[0], output.beta[1]])

    # Es scheint, dass die sd_beta von ODR auf ein chi2/dof=1 reskaliert werden. Deshalb nehmen
    # wir den Fehler direkt aus der Kovarianzmatrix.
    return output.beta[0], output.beta[1], output.cov_beta, chiq, p

def linreg(x,y,ey):
    s   = np.sum(1./ey**2)
    sx  = np.sum(x/ey**2)
    sy  = np.sum(y/ey**2)
    sxx = np.sum(x**2/ey**2)
    sxy = np.sum(x*y/ey**2)
    delta = s*sxx-sx*sx
    b   = (sxx*sy-sx*sxy)/delta
    a   = (s*sxy-sx*sy)/delta
    eb  = np.sqrt(sxx/delta)
    ea  = np.sqrt(s/delta)
    cov = -sx/delta
    #corr = cov/(ea*eb)
    #chiq  = np.sum(((y-(a*x+b))/ey)**2)
    chiq, p = chiSq(x, y, ey, [a, b])

    return(a,ea,b,eb, chiq, p)

def lin(params, x):
    #y=ax+b
    return params[0]*x+params[1]

def chiSq(x, y, err, params):
    temp = 0
    for i in range(x.shape[0]):
        temp += (y[i]- (x[i]*params[0] + params[1]))**2 / err[i]**2
    chiq = temp
    #chiqdof = chiq / (x.shape[0] - 2)
    p = 1 - chi2.cdf(chiq, x.shape[0] - 2)
    return chiq, p

def fourier(t,y):
    dt = (t[-1]-t[0])/(len(t)-1)
    fmax = 0.5/dt
    step = fmax/len(t)
    freq=np.arange(0.,fmax,2.*step)
    amp = np.zeros(len(freq))
    i=0
    for f in freq:
        omega=2.*np.pi*f
        sc = sum(y*np.cos(omega*t))/len(t)
        ss = sum(y*np.sin(omega*t))/len(t)
        amp[i] = np.sqrt(sc**2+ss**2)
        i+=1
    return (freq,amp)


def fourier_fft(t,y):
    dt = (t[-1]-t[0])/(len(t)-1)
    amp = abs(scipy.fftpack.fft(y))
    freq = scipy.fftpack.fftfreq(t.size,dt)
    return (freq,amp)