from __future__ import print_function
import numpy as np
import matplotlib as plt
import cv2
import os
import time
import urllib2
from zipfile import ZipFile
from PIL import Image
import io

np.random.seed(7)  # for reproducibility

from keras import backend as K
from keras.models import Model, model_from_json, model_from_yaml
from keras.layers import Input, Convolution2D, MaxPooling2D, LSTM, Reshape, Merge, TimeDistributed, Flatten, Activation, Dense, Dropout, merge
from attention import SpatialTransformer

def numparams(model): # Compute number of params in a model (the actual number of floats)
    return sum([np.prod(K.get_value(w).shape) for w in model.trainable_weights])

def dense_optical_flow(prev,current):
    prvs = np.uint16(prev)
    next = np.uint16(current)
    prvs = cv2.cvtColor(prev,cv2.COLOR_BGR2GRAY)
    next = cv2.cvtColor(current,cv2.COLOR_BGR2GRAY)
    hsv = np.zeros_like(prev)
    hsv[...,1] = 255
 
    flow = cv2.calcOpticalFlowFarneback(prvs,next, None, 0.5, 3, 15, 3, 5, 1.2, 0) 
    mag, ang = cv2.cartToPolar(flow[...,0], flow[...,1])
    hsv[...,0] = ang*180/np.pi/2
    hsv[...,2] = cv2.normalize(mag,None,0,255,cv2.NORM_MINMAX)
    bgr = cv2.cvtColor(hsv,cv2.COLOR_HSV2BGR)

    return bgr
    
def lucas_kanade_optical_flow(old_frame, frame):
    feature_params = dict( maxCorners = 100,
                    qualityLevel = 0.3,
                    minDistance = 7,
                    blockSize = 7 )
 
     # Parameters for lucas kanade optical flow
    lk_params = dict( winSize  = (15,15),
                   maxLevel = 2,
                   criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03))
     
     # Create some random colors
    color = np.random.randint(0,255,(100,3))
     
     # Take first frame and find corners in it
    old_gray = cv2.cvtColor(old_frame, cv2.COLOR_BGR2GRAY)
    p0 = cv2.goodFeaturesToTrack(old_gray, mask = None, **feature_params)
     
    # Create a mask image for drawing purposes
    mask = np.zeros_like(old_frame)
     
    frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
     
     # calculate optical flow
    p1, st, err = cv2.calcOpticalFlowPyrLK(old_gray, frame_gray, p0, None, **lk_params)
     
     # Select good points
    good_new = p1[st==1]
    good_old = p0[st==1]
     
     # draw the tracks
    for i,(new,old) in enumerate(zip(good_new,good_old)):
         a,b = new.ravel()
         c,d = old.ravel()
         mask = cv2.line(mask, (a,b),(c,d), color[i].tolist(), 2)
         frame = cv2.circle(frame,(a,b),5,color[i].tolist(),-1)
    img = cv2.add(frame,mask)
     
    cv2.imshow('frame',img)
    cv2.waitKey(10000)
    cv2.imshow('mask',mask)
    cv2.waitKey(10000)
    return img
    
def generate_arrays_from_file(folder):
    images = []
    data_folders = [f for f in os.listdir(folder + '/') if os.path.isdir(os.path.join(folder + '/', f))]
    for data_folder in data_folders:
	    files = [f for f in os.listdir(folder + '/' + data_folder) if os.path.isfile(os.path.join(folder + '/' + data_folder + '/', f))]
	    for f in files:
		x = cv2.imread(folder + '/' + data_folder + '/' + f, 1)
		images.append(x)
    return images

def get_optical_flow(images, img_rows, img_cols):
    optical_flow = []
    for (prev,current) in zip(images[:-1],images[1:]):
        optical_flow.append(dense_optical_flow(prev, current))
    print(optical_flow[0].shape)
    print(len(optical_flow))
    black = np.zeros((1,img_cols, img_rows,3))
    print(black.shape)
    optical_flow = np.concatenate((black, optical_flow), axis = 0)
    return optical_flow
    
def save_model(model, model_file, weights_file):
    model_yaml = model.to_yaml()
    timestamp = str(time.strftime("%H-%M-%S_%d-%m-%Y"))
    with open(model_file + timestamp + '.yaml', "w") as yaml_file:
        yaml_file.write(model_yaml)
    # serialize weights to HDF5
    model.save_weights(weights_file + timestamp + '.h5')
    print("Saved model to disk")
    
def load_model(model_file, weights_file):
    yaml_file = open(model_file, 'r')
    loaded_model_yaml = yaml_file.read()
    yaml_file.close()
    loaded_model = model_from_yaml(loaded_model_yaml)
    # load weights into new model
    loaded_model.load_weights(weights_file)
    print("Loaded model from disk")
    return loaded_model