#!/usr/bin/env lua
cv=require('luacv')

DETECTION = 0
CAPTURING = 1
CALIBRATED = 2

board_size = cv.Size(8,6)
square_size = 1
aspect_ratio = 1
image_count = 4
write_extrinsics = 0
write_points = 0
flags = 0

flip_vertical = 0
delay = 1000
prev_timestamp = 0
font = cv.Font( 1, 1 )
avg_reproj_err = 0
mode = DETECTION
undistord_image=0
image_points_buf ={}
for i=1,board_size.width*board_size.height do
  image_points_buf[i]=cv.Point2D32f(0,0)
end


capture = cv.CaptureFromFile('cam.avi') or error("The input file could not be opened")

function printMat(mat)
  for i=0,mat.height-1 do 
    local str='\t'
    for j=0,mat.width-1 do
      local scal=cv.Get2D(mat,i,j)
      str=str..'\t'..string.format('%04.2f',scal.val[1])
    end
    print(str)
  end
end


camera=cv.CreateMat(3,3,cv.CV_64F)
dist_coeffs=cv.CreateMat(1,4,cv.CV_64F);

function save_camera_params(out_filename,image_count,img_size,board_size,square_size,aspect_ratio,flags,camera_matrix,dist_coeffs,extr_params,image_points_seq,reproj_errs,avg_reproj_err)
  print('calibration_time',os.date())
  print('image_count',image_count)
  print('image_width',img_size.width)
  print('image_height',img_size.height)
  print('board_width',board_size.width)
  print('board_height',board_size.height)
  print('square_size',square_size)

  if bit32.band(flags,cv.CV_CALIB_FIX_ASPECT_RATIO)>0 then
    print('aspect_ratio',aspect_ratio)
  end
  if flags~=0 then
    print('flags:',
    bit32.band(flags,cv.CV_CALIB_USE_INTRINSIC_GUESS) and "+use_intrinsic_guess" or "",
    bit32.band(flags,cv.CV_CALIB_FIX_ASPECT_RATIO) and "+fix_aspect_ratio" or "",
    bit32.band(flags,cv.CV_CALIB_FIX_PRINCIPAL_POINT) and "+fix_principal_point" or "",
    bit32.band(flags,cv.CV_CALIB_ZERO_TANGENT_DIST) and "+zero_tangent_dist" or "" )
  end
  print('camera_matrix')
  printMat(camera_matrix)
  print('distortion_coefficients')
  printMat(dist_coeffs)
  print('avg_reprojection_error',avg_reproj_err)
end

function run_calibration(image_points_seq,img_size,board_size,square_size,aspect_ratio,flags,camera,dist_coeffs,extr_params,reproj_errs,avg_reproj_err)
  local image_count=image_points_seq.total/#image_points_buf
  local point_count=board_size.width*board_size.height

  local image_points=cv.CreateMat(1,image_count*point_count,cv.CV_32FC2)
  local object_points=cv.CreateMat(1,image_count*point_count,cv.CV_32FC3)
  local point_counts=cv.CreateMat(1,image_count,cv.CV_32SC1)
  local rot_vects=cv.CreateMatHeader(image_count,3,cv.CV_32FC1)
  local trans_vects=cv.CreateMatHeader(image_count,3,cv.CV_32FC1)
  for i=0,image_count-1 do
    index=0
    for j=0,board_size.height-1 do
      for k=0,board_size.width-1 do
        for a=0,image_count-1 do
          src_img_pt=cv.GetSeqElem(image_points_seq,index,cv.CvPoint2D32f["name"])
          cv.Set2D(object_points,0,index*image_count+a,cv.Scalar(j*square_size,k*square_size,0))
          cv.Set2D(image_points,0,index*image_count+a,cv.Scalar(src_img_pt.x,src_img_pt.y))
        end
        index=index+1
      end
    end
    collectgarbage('collect')
  end

  cv.Set(point_counts,cv.Scalar(point_count))
  extr_params=cv.CreateMat(image_count,6,cv.CV_32FC1)

  local rot_vects=cv.GetCols(extr_params,rot_vects,0,3)
  local trans_vects=cv.GetCols(extr_params,trans_vects,3,6)

  cv.Zero(camera)
  cv.Zero(dist_coeffs)
  
  if bit32.band(flags,cv.CV_CALIB_FIX_ASPECT_RATIO)>0 then
    cv.Set1D(camera,0,cv.Scalar(aspect_ratio))
    cv.Set1D(camera,4,cv.Scalar(1))
  end
  cv.CalibrateCamera2(object_points,image_points,point_counts,img_size,camera,dist_coeffs,rot_vects,trans_vects,flags)

  local code=(cv.CheckArr(camera)>0 and cv.CheckArr(dist_coeffs)>0 and cv.CheckArr(extr_params)>0 )

  cv.ReleaseMat(image_points)
  cv.ReleaseMat(object_points)
  cv.ReleaseMat(point_counts)
  return code
end


elem_size = board_size.width*board_size.height*cv.CvPoint2D32f["size"]
storage = cv.CreateMemStorage( elem_size*4> bit32.lshift(1, 16)and elem_size*4 or bit32.lshift(1, 16))

image_points_seq = cv.CreateSeq( 0, cv.CvSeq["size"], elem_size, storage )

cv.NamedWindow( "Image View", 1 )

mode=CAPTURING

while true do
  text_origin=cv.Point(0,0)
  text_size = cv.Size(0,0)
  base_line = 0
  blink=0

  view0 = cv.QueryFrame( capture ) 

  if not view0 then
    break
  end


  if view0  then
    view = cv.CreateImage( cv.GetSize(view0), cv.IPL_DEPTH_8U, view0.nChannels )
    if view0.origin == cv.IPL_ORIGIN_BL then
      cv.Flip( view0, view )
    else
      cv.Copy( view0, view )
    end
  end

  if flip_vertical >0 then
    cv.Flip( view, view)
  end


  img_size=cv.GetSize(view)

  found,count = cv.FindChessboardCorners( view, board_size,image_points_buf, cv.CV_CALIB_CB_ADAPTIVE_THRESH)

  view_gray = cv.CreateImage( cv.GetSize(view), cv.IPL_DEPTH_8U, 1 )
  cv.CvtColor( view, view_gray, cv.CV_BGR2GRAY )
  cv.FindCornerSubPix( view_gray, image_points_buf, cv.Size(11,11),cv.Size(-1,-1), cv.TermCriteria( cv.CV_TERMCRIT_EPS+cv.CV_TERMCRIT_ITER, 30, 0.1 ))

  cv.ReleaseImage( view_gray )
  if ( mode==CAPTURING and found>0 and os.clock()-prev_timestamp > delay*0.001) then
    for i=1,#image_points_buf do
      cv.SeqPush(image_points_seq,image_points_buf[i])
    end
    prev_timestamp=os.clock()
    blink=1
  end
    
  cv.DrawChessboardCorners( view, board_size, image_points_buf, count,found)

  cv.GetTextSize( '100/100', font,text_size, base_line )
  text_origin.x = view.width - text_size.width - 10
  text_origin.y = view.height - base_line - 10
  
  if mode==CAPTURING then
    if image_count > 0 then
      s=image_points_seq.total/#image_points_buf..'/'..image_count
    else
      s=image_points_seq.total/#image_points_buf..'/?'
    end
  elseif mode==CALIBRATED then
    s='Calibrated'
  end

  cv.PutText( view, s, text_origin, font, mode ~= CALIBRATED and cv.CV_RGB(255,0,0) or cv.CV_RGB(0,255,0))
  
  if (blink>0) then
    cv.Not(view,view)
  end
  
  if mode==CALIBRATED and undistord_image>0 then
    t=cv.CloneImage(view)
    cv.Undistort2(t,view,camera,dist_coeffs)
    cv.ReleaseImage(t)
  end

  cv.ShowImage( "Image View", view )
  cv.ReleaseImage( view )
  key = cv.WaitKey(50)

  if key == 27 then
    break
  elseif key==string.byte('u') then
    undistord_image=bit32.bnot(undistord_image)
  end

  if mode==CAPTURING and image_points_seq.total/#image_points_buf >= image_count then
    local code=run_calibration(image_points_seq,img_size,board_size,square_size,aspect_ratio,flags,camera,dist_coeffs,extr_params,reproj_errs,avg_reproj_err)
    save_camera_params('cam.out',image_count,img_size,board_size,square_size,aspect_ratio,flags,camera,dist_coeffs,0,0,reproj_errs,avg_reproj_err)
    if code then
      mode=CALIBRATED
    else
      mode=DETECTION
    end
  end

end
cv.ReleaseCapture( capture )
cv.ReleaseMemStorage( storage )
