import cv2
import math
import numpy
from math import sin,cos

cap = cv2.VideoCapture("../vid.mp4")
writer = cv2.VideoWriter("../outvid.avi",cv2.cv.CV_FOURCC(*'MP42'),25,(1600,600),1)

###### CONFIG ######

feature_params = dict( maxCorners = 100,
                       qualityLevel = 0.3,
                       minDistance = 20,
                       blockSize = 7 )

scale_factor=0.5
scr_width=1600
scr_height=600


# 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))




###### INITALISATION ######


def get_maps(xres,yres):
    xlist=[]
    ylist=[]

    constant = 30./41.5*3

    # calculated from official spec: maximum angle (i.e., over the diagonals)
    # is 92 deg. -> angle over half a diagonal's length is 92/2 deg
    constant2= 1./(math.sqrt((1280/2)**2+(720/2)**2))*(92/2)/180*math.pi

    for y in xrange(0,yres):
        xtmp=[]
        ytmp=[]

        yy=(y-yres/2.)/xres # yes, xres.
        for x in xrange(0,xres):
            xx=math.tan((x-xres/2.)/xres)

            dist = math.sqrt(xx**2 + yy**2)
            
            if (dist != 0):
                xtmp=xtmp+[ 1280/2+  xx/dist/constant2*math.atan(constant*dist)  ]
                ytmp=ytmp+[ 720/2 +  yy/dist/constant2*math.atan(constant*dist)  ]
            else:
                xtmp=xtmp+[0+1280/2]
                ytmp=ytmp+[0+720/2]
            
        xlist=xlist+[xtmp]
        ylist=ylist+[ytmp]
        
        if y % 10 == 0:
            print y

    xmap=numpy.array(xlist).astype('float32')
    ymap=numpy.array(ylist).astype('float32')

    return xmap,ymap

xmap,ymap = get_maps(1280,720)



ret,oldframe_=cap.read()
rawheight, rawwidth, bpp = oldframe_.shape

oldframe = cv2.remap(oldframe_, xmap, ymap, cv2.INTER_CUBIC)
oldgray=cv2.cvtColor(oldframe,cv2.COLOR_BGR2GRAY)

height, width, bpp = oldframe.shape

mask_ = numpy.ones((rawheight,rawwidth, 1), numpy.uint8) * 255
mask = cv2.remap(mask_, xmap, ymap, cv2.INTER_CUBIC)

screencontent = numpy.zeros((scr_height, scr_width,3), numpy.uint8)

total_angle=0.
total_x=1500
total_y=-300

while(cap.isOpened()):
    ret, frame_ = cap.read()
    frame = cv2.remap(frame_, xmap, ymap, cv2.INTER_CUBIC)
    gray = cv2.cvtColor(frame,cv2.COLOR_BGR2GRAY)

    
    # calculate movement, rotation and stretch. (we will ignore the stretch factor.)
    mat = cv2.estimateRigidTransform(gray,oldgray,False)
    angle = math.atan2(mat[0,1],mat[0,0])
    stretch = int((math.sqrt(mat[0,1]**2+mat[0,0]**2)-1)*100)

    # calculate shift_x and _y is if one would rotate-and-stretch around the center of the image, not the topleft corner
    shift_x = mat[0,2] - width/2 + ( mat[0,0]*width/2 + mat[0,1]*height/2 )
    shift_y = mat[1,2] - height/2 + ( mat[1,0]*width/2 + mat[1,1]*height/2 )


    # accumulate values
    total_x = total_x + shift_x
    total_y = total_y + shift_y
    total_angle=total_angle+angle
    
    print angle/3.141592654*180,'deg\t',stretch,"%\t", shift_x,'\t',shift_y
    
    # rotate and move current frame into a global context
    mat2=cv2.getRotationMatrix2D((width/2,height/2), total_angle/3.141593654*180, scale_factor) 
    mat2[0,2] = mat2[0,2]+total_x*scale_factor
    mat2[1,2] = mat2[1,2]+total_y*scale_factor


    frame2= cv2.warpAffine(frame, mat2, (scr_width,scr_height) )
    mask2 = cv2.warpAffine(mask,  mat2, (scr_width,scr_height) )
    ret, mask2 = cv2.threshold(mask2, 254, 255, cv2.THRESH_BINARY)
    mask3 = cv2.erode(mask2, numpy.ones((20,200),numpy.uint8)) # strip off the potentially-badlooking edges. left/right borders are darkened, strip them off
    mask4 = cv2.erode(mask2, numpy.ones((1,1),numpy.uint8))

    screencontent = cv2.bitwise_and(screencontent,screencontent, mask=cv2.bitwise_not(mask3)) # blank out
    screencontent = cv2.add(screencontent, cv2.bitwise_and(frame2,frame2,mask=mask3))         # and redraw

    screencontent2=screencontent.copy()
    screencontent2=cv2.bitwise_and(screencontent2,screencontent2, mask=cv2.bitwise_not(mask4))
    screencontent2=cv2.add(screencontent2, cv2.bitwise_and(frame2,frame2,mask=mask4))

    cv2.imshow('frame', frame)
    cv2.imshow('screencontent', screencontent)
    cv2.imshow('screencontent2', screencontent2)

    writer.write(screencontent2)
    
    oldframe=frame
    oldgray=gray

    key = cv2.waitKey(1) & 0xFF
    if key == ord("q"):
        break

    if key == ord("r"):
        total_angle=0.
        total_x=700
        total_y=-200
        screencontent = numpy.zeros((scr_height, scr_width,3), numpy.uint8)