#!/usr/bin/env python

import pickle
from numpy import *
import matplotlib
#matplotlib.use('AGG')
from matplotlib.pyplot import *

##################################################

# http://en.wikipedia.org/wiki/Propagation_of_uncertainty
# http://en.wikipedia.org/wiki/Linear_least_squares_%28mathematics%29#Weighted_linear_least_squares
# http://en.wikipedia.org/wiki/Least_squares#Weighted_least_squares

##################################################

def wfit(x,y,order=1,sigmay=array([])):
    # sigmay is the variance of the y values

    # form the vander matrix for x
    X = vander(x, order+1)
    # full weight matrix is inverse of the variance
    # little weight matrix is inverse of the sigma
    w = diag(ones(size(x)))
    if size(sigmay) != 0:
        w = w * 1.0/sigmay
    Xw = dot(w, X)
    yw = dot(w, y)
    (beta, residuals, rank, sing_vals) = linalg.lstsq(Xw, yw)
    fit = polyval(beta, x)
    # Mbeta is the covariance matrix off beta.
    # Even though the measurements might be uncorrelated, the fitted
    # parameters are.
    Mbeta = dot(Xw.T, Xw)
    Mbeta = linalg.inv(Mbeta)
    return beta, Mbeta, fit

def xint(b, Mb):
    # Calculate x-intercept from linear fit parameters
    xi = -b[1]/b[0]
    # Propagate the error
    saa = Mb[0,0]/b[0]
    sbb = -Mb[1,1]/b[1]
    rab = Mb[0,1]/Mb[0,0]/Mb[1,1]
    print rab
    sab = -2 * saa * sbb * rab
    rsxi = (saa)**2 + (sbb)**2 + sab
    print rsxi
    sxi = xi * sqrt(rsxi)
    print sxi
    return xi, sxi

##################################################

# load data, make fits, and plot
fmt = {'X':'b', 'Y':'g'}

xi = {}
sxi = {}
for t in ['X', 'Y']:
    f = open('schnupp_ETM' + t + '.pik', 'r')
    data = pickle.load(f)
    f.close()

    sphase = data['sphase']
    Imean = empty_like(sphase)
    Istd = empty_like(sphase)
    Qmean = empty_like(sphase)
    Qstd = empty_like(sphase)
    i = 0
    for p in sphase:
        Imean[i] = mean(data['I'][p])
        Istd[i] = std(data['I'][p])
        Qmean[i] = mean(data['Q'][p])
        Qstd[i] = std(data['Q'][p])
        i += 1

    Ib, IMb, Ifit = wfit(sphase, Imean, order=1, sigmay=Istd)

    xi[t], sxi[t] = xint(Ib, IMb)

    print "%s x-int: %f +/- %f" % (t, xi[t], sxi[t])
    errorbar(sphase, Imean, yerr=Istd,
             fmt='d'+fmt[t],
             barsabove=True,
             capsize=8,
             label=("%s" % t),
             )
    axvline(xi[t], color=fmt[t])
    axvspan(xi[t]-sxi[t], xi[t]+sxi[t],
            color=fmt[t],
            alpha=.2,
            )
    # errorbar(xi[t], 0, xerr=sxi[t],
    #          fmt='+'+fmt[t],
    #          capsize=60,
    #          )
    plot(sphase, Ifit, '--'+fmt[t], label=("%s fit" % (t)))

##################################################

# phase difference
pd = xi['X'] - xi['Y']
pa = (xi['X'] + xi['Y'])/2

# propagated phase difference error
spd = sqrt( sxi['X']**2 + sxi['Y']**2 )

# convert phase difference to length difference
# sideband frequency
f = 55e6
# speed of light
c = 3e8
# phase to length conversion
phase2len = (1./360. * c/f) * 1./2. * 100.
# 1/2 for round trip
# *100 for m -> cm

sa = phase2len * pd
ssa = phase2len * spd

text(pa, 4,
     r'''$\Delta \phi = %.2f \pm %.2f ^{\circ}$
$L_{sa} = %.2f \pm %.2f \ \mathrm{cm}$''' % (pd, spd, sa, ssa),
     size=30, va="center", ha="center",
     bbox=dict(boxstyle="round", fc="w", ec="0.5", alpha=0.9),
     backgroundcolor="white",
     )

text(63.7, -3,
     "55 MHz sideband\n103.313 Hz drive",
     size=12, va="center", ha="right",
     bbox=dict(boxstyle="round", fc="w", ec="0.5", alpha=0.9),
     backgroundcolor="white",
     )

grid('on')
ylim([-3.5, 6])
title('40m Schnupp asymmetry, relative to Y-arm')
ylabel('out-of-phase response amplitude')
xlabel('AS55 RD rotation phase (degrees)')
legend(numpoints=1,
       fancybox=True,
       loc='center right')

show()