# -*- coding: utf-8 -*-

from numpy import *

from scipy.misc import imread
import matplotlib.pyplot as plot
import pickle
import io

#from urllib.request import urlopen
from urllib2 import urlopen

R_min = 0.3
R_max = 0.5
z_min = -0.1
z_max = 0.1

phi_window = 8
lam = 11

picklePath = "pickle_py2/"

# Geometry calculation ############################################################################

def getPixelImage_top(R_min,R_max,z_min,z_max,R_camera,z_camera):
	
	recalc_z = linspace(0.157,0.495,17)
	recalc_pix_m = array([137,124,113,104,97,90,84,79,75,71,67,64,61,58,56,54,52])/0.01083/4.2

	recalc_coef = polyfit(recalc_z,1/recalc_pix_m,1)
	
	pointDistance = 0.001
	pointDistance_z = 0.001
	pointRadius = 0.001
	image = zeros((336,floor(96/phi_window)))
	[xgrid,zgrid] = meshgrid(arange(image.shape[1]),arange(image.shape[0]))
#	ipdb.set_trace()
	xgrid = (xgrid-xgrid.max()/2)*phi_window
	zgrid = zgrid-zgrid.max()/2
	for R in arange(R_min,R_max,pointDistance):
		angleDistance = pointDistance/R
		for fi in arange(-0.1,0.1,angleDistance):
			for z in arange(z_min,z_max,pointDistance_z):
				d = z_camera - z
				x = R*sin(fi)
				recalc = recalc_coef[0]*d+recalc_coef[1]
				xg = x/recalc
				zg = (R*cos(fi)-R_camera)/recalc
				rg = pointRadius/recalc
				image = image + exp(-((xgrid-xg)**2+(zgrid-zg)**2)/rg**2)
	return image

def getPixelImage_side(R_min,R_max,z_min,z_max,R_camera,z_camera):
	
	recalc_z = linspace(0.157,0.495,17)
	recalc_pix_m = array([137,124,113,104,97,90,84,79,75,71,67,64,61,58,56,54,52])/0.01083/4.2

	recalc_coef = polyfit(recalc_z,1/recalc_pix_m,1)
	
	pointDistance = 0.001
	pointRadius = 0.001
	image = zeros((336,floor(96/phi_window)))
	[xgrid,zgrid] = meshgrid(arange(image.shape[1]),arange(image.shape[0]))
	xgrid = (xgrid-xgrid.max()/2)*phi_window
	zgrid = zgrid-zgrid.max()/2
	for R in arange(R_min,R_max,pointDistance):
		angleDistance = pointDistance/R
		for fi in arange(-0.1,0.1,angleDistance):
			for z in arange(z_min,z_max,pointDistance):
				d = R_camera - R*cos(fi)
				x = R*sin(fi)
				recalc = recalc_coef[0]*d+recalc_coef[1]
				xg = x/recalc
				zg = (z-z_camera)/recalc
				rg = pointRadius/recalc
				image = image + exp(-((xgrid-xg)**2+(zgrid-zg)**2)/rg**2)
	return image

def createGeomMatrices():
	
	R_topCamera = 0.4
	z_topCamera = 0.4
	
	R_sideCamera = 0.75
	z_sideCamera = 0.0
	
	nR = 30
	
	R_grid = linspace(R_min,R_max,nR+1)
	z_grid = linspace(z_min,z_max,nR+1)
	
	T = zeros((2*336,nR*nR,floor(96/phi_window)))
	pix = 0
	for i_R in range(nR):
		for i_z in range(nR):
			image_top = getPixelImage_top(R_grid[i_R],R_grid[i_R+1],z_grid[i_z],z_grid[i_z+1],R_topCamera,z_topCamera)
			image_side = getPixelImage_side(R_grid[i_R],R_grid[i_R+1],z_grid[i_z],z_grid[i_z+1],R_sideCamera,z_sideCamera)
			
			for i_fi in range(T.shape[2]):
				T[0:336,pix,i_fi] = image_top[:,i_fi]
				T[336:,pix,i_fi] = image_side[:,i_fi]
			
			print(pix)
			pix = pix+1
	
	for i_fi in range(T.shape[2]):
		file = open(picklePath+"T_{0}_fi_{1}".format(nR,i_fi),"wb")
		pickle.dump(T[:,:,i_fi],file)
		file.close()

def getD_edgeFisher(xgrid,ygrid,pixels,boundary): # boundary: 0 .. free, 1 .. fixed
	
	nx = xgrid.size
	ny = ygrid.size

	pixels_logical = zeros(((nx+1)*ny))
	pixels_logical[pixels] = 1
	
	x_coords = repeat(arange(nx),ny)
	y_coords = repeat(arange(ny),nx).reshape((nx,ny)).T.reshape((nx*ny))
	
	pixelsWithRightNeighbor = pixels[logical_and(pixels_logical[pixels[pixels+ny<nx*ny]+ny]>0,x_coords[pixels[pixels+ny<nx*ny]]<x_coords.max())]
	pixelsWithTopNeighbor = pixels[logical_and(pixels_logical[pixels[pixels+1<nx*ny]+1]>0,y_coords[pixels[pixels+1<nx*ny]]<y_coords.max())]

	pixelsWithoutRightNeighbor = concatenate((pixels[pixels_logical[pixels[pixels+ny<nx*ny]+ny]==0], logical_and(pixels_logical[:-ny], x_coords==x_coords.max()).nonzero()[0]))
	pixelsWithoutTopNeighbor = concatenate((pixels[pixels_logical[pixels[pixels+1<nx*ny]+1]==0], logical_and(pixels_logical[:-ny], y_coords==y_coords.max()).nonzero()[0]))
	pixels_logical = pixels_logical[:-ny]
	pixelsWithoutBottomNeighbor = logical_and(pixels_logical==1,concatenate((pixels_logical[ny:],zeros((ny))))==0).nonzero()[0]
	pixelsWithoutLeftNeighbor = logical_and(pixels_logical==1,concatenate((zeros((ny)),pixels_logical[:-ny]))==0).nonzero()[0]
	
	innerEdgesCount = pixelsWithRightNeighbor.size+pixelsWithTopNeighbor.size
	outerEdgesCount = pixelsWithoutRightNeighbor.size+pixelsWithoutLeftNeighbor.size+pixelsWithoutTopNeighbor.size+pixelsWithoutBottomNeighbor.size
	
	weightIndices = zeros((innerEdgesCount,2))
	D = zeros((nx*ny,innerEdgesCount))
	
	for i in range(pixelsWithRightNeighbor.size):
		weightIndices[i,:] = [pixelsWithRightNeighbor[i],pixelsWithRightNeighbor[i]+ny]
		D[pixelsWithRightNeighbor[i],   i] = -1
		D[pixelsWithRightNeighbor[i]+ny,i] =  1
	ind0 = pixelsWithRightNeighbor.size
	
	for i in range(pixelsWithTopNeighbor.size):
		weightIndices[ind0+i,:] = [pixelsWithTopNeighbor[i],pixelsWithTopNeighbor[i]+1]
		D[pixelsWithTopNeighbor[i],  ind0+i] = -1
		D[pixelsWithTopNeighbor[i]+1,ind0+i] =  1
	
	if boundary != 0:
		weightIndices = concatenate((weightIndices,zeros((outerEdgesCount,2))),axis=0)
		D = concatenate((D,zeros((nx*ny,outerEdgesCount))),axis=1)
		ind0 = ind0+pixelsWithTopNeighbor.size
		
		for i in range(pixelsWithoutRightNeighbor.size):
			weightIndices[ind0+i,:] = [pixelsWithoutRightNeighbor[i],nx*ny+1]
			D[pixelsWithoutRightNeighbor[i],ind0+i] = boundary
		ind0 = ind0+pixelsWithoutRightNeighbor.size

		for i in range(pixelsWithoutLeftNeighbor.size):
			weightIndices[ind0+i,:] = [pixelsWithoutLeftNeighbor[i],nx*ny+1]
			D[pixelsWithoutLeftNeighbor[i],ind0+i] = boundary
		ind0 = ind0+pixelsWithoutLeftNeighbor.size

		for i in range(pixelsWithoutTopNeighbor.size):
			weightIndices[ind0+i,:] = [pixelsWithoutTopNeighbor[i],nx*ny+1]
			D[pixelsWithoutTopNeighbor[i],ind0+i] = boundary
		ind0 = ind0+pixelsWithoutTopNeighbor.size

		for i in range(pixelsWithoutBottomNeighbor.size):
			weightIndices[ind0+i,:] = [pixelsWithoutBottomNeighbor[i],nx*ny+1]
			D[pixelsWithoutBottomNeighbor[i],ind0+i] = boundary
	
	H = dot(D,D.T)
	H = H[:,pixels]
	H = H[pixels,:]
	return H

def getPixels(nR):
	Rgrid,zgrid = meshgrid(arange(nR),arange(nR))
	Rgrid = Rgrid-nR/2
	zgrid = zgrid-nR/2
	distances = sqrt(Rgrid**2 + zgrid**2)/nR*(R_max-R_min)
	pixels = (distances.reshape(nR*nR)<0.11).nonzero()[0]
	return pixels

def drawPixels(data,xgrid,ygrid,pixels):
	data_all = zeros((xgrid.size*ygrid.size))
	data_all[pixels] = data
	data_all = rot90(data_all.reshape((xgrid.size,ygrid.size)))
#	plot.imshow(data_all,interpolation='nearest')
	plot.imshow(data_all)
	plot.show()

# Image handling ##################################################################################

def loadImages(shot):
	path_top = "http://golem.fjfi.cvut.cz/shots/{0}/diagnostics/Radiation/0211FastCamera.ON/1/plasma_all.png".format(shot)
	file = io.BytesIO(urlopen(path_top).read())
	img_top = imread(file)
	path_side = "http://golem.fjfi.cvut.cz/shots/{0}/diagnostics/Radiation/0211FastCamera.ON/2/plasma_all.png".format(shot)
	file = io.BytesIO(urlopen(path_side).read())
	img_side = imread(file)
	return img_top, img_side

def getSynchronizedImages(img1, img2):
	
	img1 = array(img1,dtype=float_)
	img2 = array(img2,dtype=float_)
	
	if img1.shape[1] > img2.shape[1]:
		img2 = concatenate((img2,zeros((img2.shape[0],img1.shape[1]-img2.shape[1],img2.shape[2]))),axis=1)
	elif img2.shape[1] > img1.shape[1]:
		img1 = concatenate((img1,zeros((img1.shape[0],img2.shape[1]-img1.shape[1],img1.shape[2]))),axis=1)
	
	brightness1 = img1.sum(axis=2).mean(axis=0)
	brightness2 = img2.sum(axis=2).mean(axis=0)
	
	ok1 = brightness1 > brightness1.mean()*0.01
	ok2 = brightness2 > brightness2.mean()*0.01
	
	phase1 = zeros(ok1.shape)
	center1 = ok1.nonzero()[0].mean()
	for i in range(ok1.size-1):
		indices_fromLeft = logical_not(ok1[:i]).nonzero()[0]
		phase_fromLeft = i-indices_fromLeft.max() if indices_fromLeft.size>0 else 0
		indices_fromRight = logical_not(ok1[i+1:]).nonzero()[0]
		phase_fromRight = 97-indices_fromRight.min() if indices_fromRight.size>0 else ok1.size
		if (phase_fromLeft>phase_fromRight) ^ (i>center1):
			phase1[i] = phase_fromLeft
		else:
			phase1[i] = phase_fromRight
	phase1[logical_not(ok1)] = -1
	
	phase2 = zeros(ok2.shape)
	center2 = ok2.nonzero()[0].mean()
	for i in range(ok2.size-1):
		indices_fromLeft = logical_not(ok2[:i]).nonzero()[0]
		phase_fromLeft = i-indices_fromLeft.max() if indices_fromLeft.size>0 else 0
		indices_fromRight = logical_not(ok2[i+1:]).nonzero()[0]
		phase_fromRight = 97-indices_fromRight.min() if indices_fromRight.size>0 else ok2.size
		if (phase_fromLeft>phase_fromRight) ^ (i>center2):
			phase2[i] = phase_fromLeft
		else:
			phase2[i] = phase_fromRight
	phase2[logical_not(ok2)] = -1
	
	brightness1 = interp(arange(brightness1.size),ok1.nonzero()[0],brightness1[ok1])
	brightness1[isnan(brightness1)] = 0
	brightness1 = convolve(brightness1,ones(100)/100,mode='same')
	brightness2 = interp(arange(brightness2.size),ok2.nonzero()[0],brightness2[ok2])
	brightness2[isnan(brightness2)] = 0
	brightness2 = convolve(brightness2,ones(100)/100,mode='same')
	
	bestShift = -((brightness1>brightness1.max()*0.1).nonzero()[0].min()+(brightness1>brightness1.max()*0.1).nonzero()[0].max()
	                 -(brightness2>brightness2.max()*0.1).nonzero()[0].min()-(brightness2>brightness2.max()*0.1).nonzero()[0].max())/2
	
	if bestShift<0:
		img1_s = concatenate((img1,zeros((336,abs(bestShift),3))),axis=1)
		phase1 = concatenate((phase1,-ones(abs(bestShift))))
		img2_s = concatenate((zeros((336,abs(bestShift),3)),img2),axis=1)
		phase2 = concatenate((-ones(abs(bestShift)),phase2))
	else:
		img1_s = concatenate((zeros((336,abs(bestShift),3)),img1),axis=1)
		phase1 = concatenate((-ones(abs(bestShift)),phase1))
		img2_s = concatenate((img2,zeros((336,abs(bestShift),3))),axis=1)
		phase2 = concatenate((phase2,-ones(abs(bestShift))))

	return img1_s, phase1, img2_s, phase2, bestShift

def selectFrames(img_top, img_side):
	frameStep = 20
	img_top, phase_top, img_side, phase_side, shift = getSynchronizedImages(img_top, img_side)
	bothPhases = logical_and(logical_and(phase_top>0, phase_side>0),logical_and(phase_top<92, phase_side<92)).nonzero()[0]
	frames = bothPhases[::frameStep]
	return img_top[:,frames,:], phase_top[frames], img_side[:,frames,:], phase_side[frames], frames, shift

# Reconstruction ##################################################################################

def getReconstruction_allFrames(img_top, phase_top, img_side, phase_side):
	print('1/{0}..'.format(img_top.shape[1]))
	reconstruction = getReconstruction_singleFrame(img_top[:,0,:],phase_top[0],img_side[:,0,:],phase_side[0])
	allRecons = zeros((reconstruction.shape[0],reconstruction.shape[1],img_top.shape[1]))
	allRecons[:,:,0] = reconstruction
	for s in range(1,img_top.shape[1]):
		print('{0}/{1}..'.format(s+1,img_top.shape[1]))
		reconstruction = getReconstruction_singleFrame(img_top[:,s,:],phase_top[s],img_side[:,s,:],phase_side[s])
		allRecons[:,:,s] = reconstruction
	return allRecons

def getReconstruction_singleFrame(img_top,phase_top,img_side,phase_side,draw=False):

	nR = 30
	
	a = 43000
	b = 140
	
	data_nl = concatenate((array(img_top,dtype=float_),array(img_side,dtype=float_)),axis=0)
	data = data_nl/(a-b*data_nl)/2
	data = a*data/(1+b*data)
	
	phase_top = (phase_top/8).round().clip(min=1,max=round(96/phi_window))-1
	file = open(picklePath+"T_{0}_fi_{1:.0f}".format(nR,phase_top),"rb")
	T_fi = pickle.load(file)
	file.close()
	geomMat = zeros((2*336,T_fi.shape[1]))
	geomMat[:336,:] = fliplr(T_fi[:336,:])
	
	phase_side = (phase_side/8).round().clip(min=1,max=round(96/phi_window))-1
	file = open(picklePath+"T_{0}_fi_{1:.0f}".format(nR,phase_side),"rb")
	T_fi = pickle.load(file)
	file.close()
	geomMat[336:,:] = rot90(T_fi[336:,:],2)*0.75
	
	pixels = getPixels(nR)
	Rgrid,zgrid = meshgrid(arange(nR),arange(nR))
	Rgrid = Rgrid-nR/2
	zgrid = zgrid-nR/2
	Rgrid = Rgrid.reshape((nR*nR))
	zgrid = zgrid.reshape((nR*nR))
	distances = sqrt(Rgrid[pixels]**2 + zgrid[pixels]**2)/nR*(R_max-R_min)
	
	geomMat = geomMat[:,pixels]
	
	H = getD_edgeFisher(arange(nR),arange(nR),pixels,0)
	H = dot(dot(H,H),H)
#	H = dot(H,H)
	reconstruction = ones((pixels.size,3))
	
	mask = arctan((distances-0.1)*5E3)+3.1415926535897932/2
	H = H + diag(mask*0.5)
	
	for c in range(3):
		reconstruction[:,c] = linalg.solve(dot(geomMat.T,geomMat)+exp(lam)*H,dot(geomMat.T,data[:,c]))
	
	if draw:
		drawPixels(reconstruction[:,2],arange(nR),arange(nR),pixels)

	return reconstruction