%matplotlib inline

import matplotlib.pylab as plt
import matplotlib.cm as cm
import numpy as np
import pymks
import cStringIO
import Image as Image
import urllib
import scipy.fftpack as spfft
import PIL
from IPython.display import Image as IPI
from IPython.display import display as display
from io import BytesIO

def get_data_url(url):
    file = cStringIO.StringIO(urllib.urlopen(url).read())
    img = Image.open(file)
    return np.asarray(img.convert('L'))

from pymks.bases import DiscreteIndicatorBasis
from pymks.stats import autocorrelate

def get_auto_correlations(data):
    basis = DiscreteIndicatorBasis(n_states=2, domain=[0, 1])
    data_ = basis.discretize(np.round(data / 255.).astype(int))
    return autocorrelate(data_[None])

def get_sym_data(scale_data):
    #scale_data = real_data[570:1521, 550:1501]
    L_small = scale_data.shape[0]
    L_large = 2 * scale_data.shape[0] - 1
    sym_data = np.zeros((L_large, L_large))
    sym_data[:L_small, (L_small - 1):] = scale_data
    sym_data[(L_small - 1):, (L_small - 1):] = scale_data[::-1]
    sym_data[(L_small - 1):, :L_small] = scale_data[::-1, ::-1]
    sym_data[:L_small, :L_small] = scale_data[:, ::-1]
    return sym_data

def display_array(arr):
    im = PIL.Image.fromarray(np.uint8(cm.gist_earth(arr) * 255))
    bio = BytesIO()
    im.save(bio, format='png')
    bio.getvalue()
    display(IPI(bio.getvalue(), format='png'))

url = 'https://farm4.staticflickr.com/3900/15362424645_b2c41c86b4_o_d.png'

#sim_data = get_data_url('https://farm4.staticflickr.com/3839/15173631407_859191a4db_o_d.jpg')
#sim_auto_correlation = get_auto_correlations(sim_data)
image_data = get_data_url(url)
image_auto_correlation = get_auto_correlations(image_data)

#display_array(image_data)

#display_array(np.fft.ifftshift(np.real_if_close(np.fft.ifftn(image_auto_correlation[0, ..., 1]))) * 1e8)
#lt.imshow(np.real_if_close(np.fft.ifftn(image_auto_correlation[0, ..., 0])))
#lt.colorbar()

#real_data = get_data_url('https://farm4.staticflickr.com/3918/15083260160_0e42df4bc9_o_d.jpg') # Log Scaled Image
real_data = np.float_(get_data_url('https://farm6.staticflickr.com/5591/15083317630_45562a1369_o_d.jpg')) # Original Image
print real_data.dtype
#fig = plt.gcf()
#fig.set_size_inches(24.5,18.5)
#plt.imshow(sym_data[710:730, 710:730], shape=(25, 25), interpolation='none')
fig = plt.gcf()
fig.set_size_inches(24.5,18.5)
plt.imshow(real_data, shape=(25, 25), interpolation='none')
plt.colorbar()

XX = 2

YY = -1
scale_data = (real_data[6: 727, 237 + XX: 958 + XX] / 39)
#scale_data = real_data[11:731, 235:956]
#plt.imshow(scale_data, cmap=cm.gray)
#display_array(scale_data)

sym_data = get_sym_data(scale_data)

fig = plt.gcf()
fig.set_size_inches(24.5,18.5)
#plt.imshow(sym_data[710:730, 710:730], shape=(25, 25), interpolation='none')
plt.imshow(sym_data, interpolation='none')
plt.colorbar()

beam_mask = np.where(sym_data[712:729, 712:729] < 0.35, 1.0, sym_data[712:729, 712:729])
sym_data[712:729, 712:729] = beam_mask
fig = plt.gcf()
fig.set_size_inches(24.5,18.5)
plt.imshow(sym_data[707:734, 707:734], shape=(25, 25), interpolation='none')
plt.colorbar()


fig = plt.gcf()
fig.set_size_inches(24.5,18.5)
plt.imshow(sym_data, shape=(25, 25), interpolation='none')
plt.colorbar()

assert(np.array_equal(sym_data, sym_data[::-1, ::-1]))
assert(np.array_equal(sym_data, sym_data[::-1]))
assert(np.array_equal(sym_data, sym_data[:, ::-1]))
assert(np.unique(sym_data.imag) == [0])


sym_data_shift = np.fft.ifftshift(sym_data)
real_auto = np.real_if_close(np.fft.ifftn(sym_data_shift)) # / np.product(np.array(sym_data.shape))

print np.min(real_auto)
print np.max(real_auto)
#plt.imshow(np.fft.fftshift(np.abs(real_auto)) * 2.8e3)
#plt.imshow(np.fft.fftshift(np.abs(real_auto)) * 0.7e2)
#plt.colorbar()
#print np.unique(np.real_if_close(real_auto).imag)
print np.min(real_auto.real), np.max(real_auto.real)
print np.min(real_auto.imag), np.max(real_auto.imag)
#display_array(np.fft.fftshift(np.abs(real_auto)) * 2.8e3) #linearly scaled from 0 to 1
#display_array(np.fft.fftshift(np.abs(real_auto)) * 1e3) #linearly scaled from 0 to 1 with beam mask
#display_array(np.fft.fftshift(np.abs(real_auto)) * 0.7e2) #log scaled data
fig = plt.gcf()
fig.set_size_inches(24.5,18.5)
#plt.imshow(np.fft.fftshift(np.abs(real_auto)) * 2.8e3, shape=(25, 25), interpolation='none')
plt.imshow(np.fft.fftshift(real_auto.real) * 2.8e3, shape=(25, 25), interpolation='none')
plt.colorbar()

print np.unique(np.abs(real_auto.imag))