#!/usr/bin/python
# -*- coding: utf-8 -*-
# process mirror spec vs angle/pol data (SPECLAB 985 - 1099)
from netCDF4 import Dataset
from numpy import *
from matplotlib.pyplot import *
from binaryMatrixFile import binRead, binWrite
#from scipy.stats import nanmean
from lineColorSeries import lineCols
from scipy.stats import nanmean
from scipy.optimize import curve_fit
import re;

#function for fitting sine curves to polarisation trace
def sineFitFunc(x, amp, ctrast, period, xOffset) :
        return amp * ctrast * sin(2*pi*(x - xOffset) / period) + amp;


#wavelen
l0 = 651.0;
l1 = 655.0;

#p0=985;
#p0=1000;
#p0=1014;
#p0=1031;
#p0=1040;
#p0=1043;
#p0=1061;
#p0=1079; # (phase)

forceCollectData = False;
forceFits = False;

if "sets" not in dir() : 
	sets = [];

#pulseInfos = [[ mark index on table,
#		 pulse number, 
#		mirror to bolt distance, 
#		source arm to bolt distance, 
#		detector arm to bolt distance ], ...]

if p0 == 985 :
	setName = "mirror '3' (the dodgy one)";

	pulseInfos = [[1, 986, 303, 527, 541],
	[3, 987, 303, 456, 464],
	[5, 988, 353, 462, 471],
	[7, 989, 353, 395, 402],
	[9, 990, 353, 341, 333],
	[11, 991, 353, 279, 286],
	[13, 992, 353, 225,  229],
	[14, 993, 353, 189, 192],
	[12, 994, 353, 257, 262],
	[10, 995, 353, 304, 309],
	[8, 996, 353, 364, 372],
	[6, 997, 353, 440, 430],
	[4, 998, 353, 501, 511],
	[2, 999, 303, 491, 500]];

	#fit init estimates, depends on motor speed
	estPeriod = 700;
	iMax = 1760; 

elif p0 == 1000 : 
	setName = "mirror #3 main scan";
	
	pulseInfos = [[1, 1000, 303, 527, 541],
	[3, 1001, 303, 456, 464],
	[5, 1002, 353, 462, 471],
	[7, 1003, 353, 395, 402],
	[9, 1004, 353, 341, 333],
	[11, 1005, 353, 279, 286],
	[13, 1006, 353, 225,  229],
	[14, 1007, 353, 189, 192],
	[12, 1008, 353, 257, 262],
	[10, 1009, 353, 304, 309],
	[8, 1010, 353, 364, 372],
	[6, 1011, 353, 440, 430],
	[4, 1012, 353, 501, 511],
	[2, 1013, 303, 491, 500]];

	#fit init estimates, depends on motor speed
	estPeriod = 1300;
	iMax = 1750; 

elif p0 == 1014 :
	setName = "mirror #3 fine scan";
	
	pulseInfos = [[6.0, 1014, 353, 440, 430],
	[6.5, 1015, 353, 409, 416],
	[7.0, 1016, 353, 395, 402],
	[7.5, 1017, 353, 384, 379],
	[8.0, 1018, 353, 364, 372],
	[8.5, 1019, 353, 347, 353],
	[9.0, 1020, 353, 341, 333],
	[9.5, 1021, 353, 319, 322],
	[10., 1022, 353, 304, 309],
	[9.5, 1023, 353, 319, 322],
	[9.0, 1024, 353, 341, 333],
	[8.5, 1025, 353, 347, 353],
	[8.0, 1026, 353, 364, 372],
	[7.0, 1027, 353, 395, 402],
	[7.5, 1028, 353, 384, 379],
	[6.5, 1029, 353, 409, 416],
	[6.0, 1030, 353, 440, 430]];

	#fit init estimates, depends on motor speed
	estPeriod = 1200;
	iMax = 2500; 

elif p0 == 1031 :
	setName = "mirror #1 front side";
	pulseInfos = [[6.0, 1031, 353, 440, 430],
	[6.5, 1032, 353, 409, 416],
	[7.0, 1033, 353, 395, 402],
	[7.5, 1034, 353, 384, 379],
	[8.0, 1035, 353, 364, 372],
	[8.5, 1036, 353, 347, 353],
	[9.0, 1037, 353, 341, 333],
	[9.5, 1038, 353, 319, 322],
	[10., 1039, 353, 304, 309]];
	estPeriod = 830;
	iMax = 1750; 

elif p0 == 1040 :
	setName = "mirror #1 back side";
	pulseInfos = [[6.0, 1040, 353, 440, 430],
	[8.0, 1041, 353, 364, 372],
	[10., 1042, 353, 304, 309]];
	estPeriod = 830;
	iMax = 1750; 

elif p0 == 1043 :
	setName = "mirror 4 front side";
	pulseInfos = [[1.0, 1043, 303, 527, 541],
	[2.0, 1044, 303, 491, 500],
	[3.0, 1045, 303, 456, 464],
	[4.0, 1046, 353, 501, 511],
	[5.0, 1047, 353, 462, 471],
	[6.0, 1048, 353, 440, 430],
	[6.5, 1049, 353, 409, 416],
	[7.0, 1050, 353, 395, 402],
	[7.5, 1051, 353, 384, 379],
	[8.0, 1052, 353, 364, 372],
	[8.5, 1053, 353, 347, 353],
	[9.0, 1054, 353, 341, 333],
	[9.5, 1055, 353, 319, 322],
	[10., 1056, 353, 304, 309],
	[11., 1057, 353, 279, 286],
	[12., 1058, 353, 257, 262],
	[13., 1059, 353, 225, 229],
	[14., 1060, 353, 189, 192]];
	estPeriod = 830;
	iMax = 1740; 
	

elif p0 == 1061 :
	setName = "mirror 1 repeat";
	pulseInfos = [[1.0, 1061, 303, 527, 541],
	[2.0, 1062, 303, 491, 500],
	[3.0, 1063, 303, 456, 464],
	[4.0, 1064, 353, 501, 511],
	[5.0, 1065, 353, 462, 471],
	[6.0, 1066, 353, 440, 430],
	[6.5, 1067, 353, 409, 416],
	[7.0, 1068, 353, 395, 402],
	[7.5, 1069, 353, 384, 379],
	[8.0, 1070, 353, 364, 372],
	[8.5, 1071, 353, 347, 353],
	[9.0, 1072, 353, 341, 333],
	[9.5, 1073, 353, 319, 322],
	[10., 1074, 353, 304, 309],
	[11., 1075, 353, 279, 286],
	[12., 1076, 353, 257, 262],
	[13., 1077, 353, 225, 229],
	[14., 1078, 353, 189, 192]];
	estPeriod = 850;
	iMax = 1750; 

elif p0 == 1079 :
	setName = "mirror 1 phase";
	pulseInfos = [[2.0, 1079, 303, 491, 500],
	[1.0, 1080, 303, 527, 541],
	[3.0, 1081, 303, 456, 464],
	[4.0, 1082, 353, 501, 511],
	[5.0, 1083, 353, 462, 471],
	[6.0, 1084, 353, 440, 430],
	[6.5, 1085, 353, 409, 416],
	[7.0, 1086, 353, 395, 402],
	[7.5, 1087, 353, 384, 379],
	[8.0, 1088, 353, 364, 372],
	[8.5, 1089, 353, 347, 353],
	[9.0, 1090, 353, 341, 333],
	[9.5, 1091, 353, 319, 322],
	[10., 1092, 353, 304, 309],
	[11., 1093, 353, 279, 286],
	[12., 1094, 353, 257, 262],
	[13., 1095, 353, 225, 229],
	[14., 1096, 353, 189, 192]];
else :
	raise Exception("p0 is not a known set");

pulseInfos = array(pulseInfos);

np = pulseInfos.shape[0];
ns=2034;
na=4000;


pulses=zeros(np);
angs=zeros(np);

# Step 1: Load/Produce big data file from spectrometer raw text files
try :
	if forceCollectData : 
		raise Exception("forceCollectData is true");

	#you don't actually need this if you're working from the specFits
	d = Dataset("/work/data/gmds/SPECLAB/"+str(p0)+"/mirrorAngSpec.nc").variables['data'][:];

except Exception as err : 
	print("Collecting data for p0 = %i because %s" % (p0, err));

	d=zeros((np, na, ns));
	for iP in range(0, np) :
	
		num = pulseInfos[iP, 0];
		pulse = pulseInfos[iP, 1];
		ac = float(pulseInfos[iP, 2]);
		bc = float(pulseInfos[iP, 3]);
		dc = float(pulseInfos[iP, 4]);
		ang = arctan((bc+dc)/2.0 / ac);

		pulses[iP] = pulse;
		angs[iP] = ang;

		for ai in range(0, na) :
			try :
				a = loadtxt("/work/data/gmds/SPECLAB/"+str(pulse)+"/spec-"+str(pulse)+"-"+str(ai)+".txt");
				print(str(pulse) + "-" + str(ai) + " OK");
				wavelen = a[:,2];
				d[iP,ai,:] = a[:,3];

			except Exception as err:
				print(str(pulse) + "-" + str(ai) + ": X");
				continue;

		

	ds = Dataset("/work/data/gmds/SPECLAB/"+str(p0)+"/mirrorAngSpec.nc", mode="w");
	dim = ds.createDimension("pulse", np);
	dim = ds.createDimension("ang", na);
	dim = ds.createDimension("spec", ns);
	var = ds.createVariable("data", "f8", dimensions=('pulse', 'ang', 'spec' ));
	var[:] = d;
	ds.close();

	
a = loadtxt("/work/data/gmds/SPECLAB/"+("%i" % pulseInfos[0,1])+"/spec-"+("%i" % pulseInfos[0,1])+"-1.txt");
wavelen = a[:,2];
								
# calc angles from measurements 
ac = pulseInfos[:,2].astype('float');
bc = pulseInfos[:,3].astype('float');
dc = pulseInfos[:,4].astype('float');
angs = 2.0*arctan((bc+dc)/2.0 / ac);

#find relevant spectral region
il0 = argmin((wavelen - l0)**2);
il1 = argmin((wavelen - l1)**2);
ia = argsort(angs);

# Step 2: Load/Produce fits to polarisation variation for each angle/wavelength
try : 
	if forceFits : 
		raise Exception("forceFits is true");

	pFinal = Dataset("/work/data/gmds/SPECLAB/"+str(p0)+"/sineFits.nc").variables['data'][:];
	
except Exception as err : 
	print("Performing fits for p0 = %i because %s" % (p0, err));

	pFinal = zeros((np, ns, 4));
	pInit = zeros((np, ns, 4));
	for i in range(0, np) : 
		for j in range(530, 1670) :  # wavelen ~ 500-800nm
			x=arange(0, iMax);
			x=x[d[i,x,j] > 0];
	
			vMax = d[i, x, j].max();
			vMin = d[i, x, j].min();
			vMean = d[i, x, j].mean();
			estAmp = (vMax + vMin) / 2;
			estContrast = (vMax - vMin) / (vMax + vMin);
			
			idxMax = argmax(d[i, x, j]);
			estX0 = idxMax - estPeriod / 4;
			
			if estContrast < 0.00 : estContrast = 0.20;
			if estContrast > 0.99 : estContrast = 0.99;
			if estAmp < 100 : estAmp = 100;
			if estAmp > 30000 : estAmp = 30000;
	
			pInit[i, j, :] = array([estAmp, estContrast, estPeriod, estX0]);
			try :
				ret = curve_fit(sineFitFunc, x, d[i, x, j], pInit[i, j, :], 100);
				pFinal[i, j, :] = ret[0];
				if pFinal[i, j, 1] < 0 :  # correct negative amplitudes by moving offset by 180°
					pFinal[i, j, 1] *= -1;
					pFinal[i, j, 3] = pFinal[i, j, 3] + pFinal[i, j, 2] / 2; 
				#estX0 = pFinal[i, j, 3];
				#estPeriod = pFinal[i, j, 2];
			except RuntimeError as err : 
				#fitting failed, skip it
				print(err);
				pFinal[i, j, :] = array([nan, nan, nan, nan]);

			#some plotting every 50 wavelengths
			if pFinal[i, j, 0] < 0 or mod(j, 50) == 0 :
				figure(0);
				clf();
				plot(x, d[i, x, j], "b");
				plot(x, sineFitFunc(x, pInit[i,j,0], pInit[i,j,1], pInit[i,j,2], pInit[i,j,3]), "g");
				plot(x, sineFitFunc(x, pFinal[i,j,0], pFinal[i,j,1], pFinal[i,j,2], pFinal[i,j,3]), "r");
				ax=gca();
				ax.set_title("Fit i = %i, j = %i" % (i, j));
				draw();

				figure(1);
				clf();
				lineCols(np);
				amp = pFinal[ia,:,0].T; # (a-b)/2
				ctrast = pFinal[ia,:,1].T; # (a-b)/(a+b)			
				#offset=amp/ctrast = (a+b)/2;
				#offset - amp = b
				#offset + amp = a
				#AsAp = (offset - amp) / (offset + amp)
				#AsAp = (amp/ctrast - amp) / (amp/ctrast + amp)
 				AsAp = (1.0/ctrast - 1.0) / (1.0/ctrast + 1.0); # = b/a
				plot(wavelen, AsAp);
				legend((angs[ia]*180/pi).astype('int'));
				ax=gca();
				ax.set_title("Contrasts vs ");
				draw();
				
		print(i);

	ds = Dataset("/work/data/gmds/SPECLAB/"+str(p0)+"/sineFits.nc", mode="w");
	dim = ds.createDimension("pulse", np);
	dim = ds.createDimension("spec", ns);
	dim = ds.createDimension("param", 4);
	var = ds.createVariable("data", "f8", dimensions=('pulse', 'spec', 'param'));
	var[:] = pFinal;
	ds.close();


if p0 == 1079 : #phase shift between p, s
	# setup was polariser at 45°, mirror (delay plate of φ at 0°) then polariser at θ
	# I = amp * (1 + cos φ · sin 2θ) 
	# so our 'contrast' is just cos φ
	#
	# If there was an error in the input polariser angle: δ, then we have:
	# I = amp * (1 + cos(φ)·cos(2δ)·sin(2θ) + sin(2δ)·cos(2θ))
	# a = sin(2δ),   b = cos(φ)·cos(2δ)
	# I = amp * (1 + sqrt(a² + b²) cos(2θ + ?))
	# contrast = sqrt(a² + b²) 
	#          = sqrt(sin²(2δ) + cos²(φ)·cos²(2δ))
	# even for δ = 5° (we were definitely better than that)
	#   δ=5°,φ=0.0° --> ctrast=1 --> θ < 10^-7
	#   δ=5°,φ=10.0° --> ctrast=0.985 --> θ = 9.8°
	#   δ=5°,φ=19.6° --> ctrast=0.942 --> θ = 19.7°
	# so thi smeasurement is extremely robust!
	#   ... really?

	l = arange(650.0, 657.0, 2.0);
	l0 = l - 0.5;
	l1 = l + 0.5;
	nl = len(l);
	phi=zeros((nl, np));
	for i in range(0, nl) :
		il0 = argmin((wavelen - l0[i])**2);
		il1 = argmin((wavelen - l1[i])**2);
				
		y = nanmean(pFinal[:, il0:il1, 1], axis=1);
		phi[i, :] = arccos(y);

	clf();
	lineCols(nl);
	plot(angs[ia]*180/pi, phi[:, ia].T*180/pi);
	legend(l);
	ax=gca();
	ax.set_title("Mirror phase change for '" + setName + "'");
	ax.set_ylabel("Phase difference / $^o$");
	ax.set_xlabel("Incidence Angle / $^o$");
	draw();	

else :
	y = nanmean(pFinal[ia, 1095:1105, 1],axis=1);

	y = (1.0/y - 1.0) / (1.0/y + 1.0); # = a/b
	y = y * 100;

	plot(angs[ia]*180/pi, y, label=setName);
	ax=gca();
	legend();
	ax.set_title("Is/Ip reflection ratios");
	ax.set_xlabel("Incidence Angle / $^o$");
	ax.set_ylabel("Is / Ip in %");
	draw();



if False : #old code used in lab for fast plots


	#calc stdev and mean of valid regions
	dSpecRegion = d[ia, :, il0:il1].mean(axis=2);
	nAngs = dSpecRegion.shape[0]
	meanAll = zeros(nAngs);
	stdAll = zeros(nAngs);
	maxAll = zeros(nAngs);
	minAll = zeros(nAngs);
	for i in range(0, nAngs) : 
		isValid = dSpecRegion[i, :] > 0;
		meanAll[i] = dSpecRegion[i, isValid].mean(axis=0);
		stdAll[i] = dSpecRegion[i, isValid].std(axis=0);
		maxAll[i] = dSpecRegion[i, isValid].max(axis=0);
		minAll[i] = dSpecRegion[i, isValid].min(axis=0);

	#fill in set info
	thisSet = dict();
	thisSet['name'] = setName;
	thisSet['angs'] = angs[ia];
	thisSet['mean'] = meanAll;
	thisSet['stdev'] = stdAll;
	thisSet['fracChange'] = stdAll / meanAll;
	thisSet['fracChange2'] = (maxAll - minAll) / (maxAll + minAll);
	thisSet['wavelen'] = wavelen;
	sets.append(thisSet);

	plot(thisSet['angs']*180/pi, log10(thisSet['fracChange']), label=thisSet['name']);
	plot(thisSet['angs']*180/pi, log10(thisSet['fracChange2']), label=thisSet['name']+"minMax");

	legend();
	draw();