from display import parse_and_show, displacement_field
import matplotlib.pyplot as plt
import subprocess
import time
import numpy as np
import cPickle as pickle
import os
import sys
from config import frame_cols, frame_rows, get_px_ndx, get_image_pair

image_dir = '../../data/vort_sim'
im_pair = get_image_pair(image_dir)

X, Y, U, V = displacement_field(pickle.load(open(os.path.join(image_dir, 'stdout.pck'), 'rb')), image_dir)

from scipy import mgrid, signal

def gauss_kern(size, sizey=None):
    """ Returns a normalized 2D gauss kernel array for convolutions """
    size = int(size)
    if not sizey:
        sizey = size
    else:
        sizey = int(sizey)
    x, y = mgrid[-size:size+1, -sizey:sizey+1]
    g = exp(-(x**2/float(size)+y**2/float(sizey)))
    return g / g.sum()
def blur_image(im, n, ny=None) :
    """ blurs the image by convolving with a gaussian kernel of typical
        size n. The optional keyword argument ny allows for a different
        size in the y direction.
    """
    g = gauss_kern(n, sizey=ny)
    improc = signal.convolve(im,g, mode='valid')
    return(improc)

Z = np.power(U, 2) + np.power(V, 2)

plt.figure()
plt.contourf(Z)
plt.colorbar()
plt.gca().invert_yaxis()
plt.savefig('sample_density.pdf')

plt.figure()
plt.contourf(blur_image(Z, 5))
plt.gca().invert_yaxis()
plt.colorbar()

P = blur_image(Z, 5)
cum_P = np.cumsum(np.reshape(P, None, 1))
total = cum_P[-1]

P = np.exp(Z)
P.size
P.shape

P = np.exp(Z / 6)

pixel_Z = np.kron(Z, np.ones((8, 8)))
pixel_Z = np.concatenate((pixel_Z, np.zeros((448, 32))), axis=1)
pixel_Z = np.concatenate((pixel_Z, np.zeros((32, 640))), axis=0)
print pixel_Z.shape
pixel_P = np.exp(pixel_Z/6)
cum_P = np.cumsum(np.reshape(pixel_P, None, 1))
total = cum_P[-1]
new = np.random.sample(1000)*total
ndx = np.searchsorted(cum_P, new)
y = ndx // im_pair.A.size[0]
x = ndx % im_pair.A.size[0]
ndx[x > 600] = ndx[x > 600] - 40
ndx[y > 440] = ndx[y > 440] - 40 * im_pair.A.size[0]
y = ndx // im_pair.A.size[0]
x = ndx % im_pair.A.size[0]
plt.figure()
plt.plot(x, y, 'r.')
plt.gca().invert_yaxis()
plt.savefig('sample_points.pdf')
pickle.dump(ndx, open('adaptive_coords.pck', 'wb'))
print max(ndx), 640*480

print max(x), max(y)

a = np.array([1,2,3])
a[a > 1]

np.searchsorted(cum_P, total)

3035 // frame_cols(im_pair.A)

frame_rows(im_pair.A)

frame_cols(im_pair.A)

X.size

Z.size

np.concatenate((np.array([[1],[2],[3]]), np.zeros((3,2))), axis=1)