#!/usr/bin/env python
import numpy as np
import cv2
import json

# local modules
from common import splitfn

# built-in modules
import os


USAGE = '''
USAGE: calib.py [--save <filename>] [--debug <output path>] [--square_size] [<left image mask>] [<right image mask>]
'''



if __name__ == '__main__':
    import sys
    import getopt
    from glob import glob

    args, img_mask = getopt.getopt(sys.argv[1:], '', ['save=', 'debug=', 'square_size='])
    args = dict(args)
    try:
        img_mask_left = img_mask[0]    
        img_mask_right = img_mask[1]    
    except:
        img_mask_left = '../data/left*.jpg'
        img_mask_right = '../data/right*.jpg'

    img_names = [glob(img_mask_left), glob(img_mask_right)]
    img_names[0].sort()
    img_names[1].sort()
    debug_dir = args.get('--debug')
    save_name = args.get('--save')
    square_size = float(args.get('--square_size', 1.0))
    print "Square Size:", square_size, "cm"

    pattern_size = (9, 6)
    pattern_points = np.zeros( (np.prod(pattern_size), 3), np.float32 )
    pattern_points[:,:2] = np.indices(pattern_size).T.reshape(-1, 2)
    pattern_points *= square_size

    obj_points = [[], []]
    img_points = [[], []]
    h = [0, 0]
    w = [0, 0]
    
    for i in range(2):
        for fn in img_names[i]:
            print 'processing %s...' % fn,
            img = cv2.imread(fn, 0)
            #gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            if img is None:
              print "Failed to load", fn
              continue

            h[i], w[i] = img.shape[:2]
            found, corners = cv2.findChessboardCorners(img, pattern_size)
            if found:
                term = ( cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_COUNT, 30, 0.1 )
                cv2.cornerSubPix(img, corners, (5, 5), (-1, -1), term)
            if debug_dir:
                vis = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
                cv2.drawChessboardCorners(vis, pattern_size, corners, found)
                path, name, ext = splitfn(fn)
                cv2.imwrite('%s/%s_chess.bmp' % (debug_dir, name), vis)
            if not found:
                print 'chessboard not found'
                continue
            img_points[i].append(corners.reshape(-1, 2))
            obj_points[i].append(pattern_points)

            print 'ok'
            
    assert (len(img_points[0]) == len(img_points[1]))
    assert (len(obj_points[0]) == len(obj_points[1]))
    obj_points = obj_points[0]
    
    # run monocular calibration on each camera to get intrinsic parameters
    rms_l, camera_matrix_l, dist_coeffs_l, _, _ = cv2.calibrateCamera(obj_points, img_points[0], (w[0], h[0]), None, None)
    rms_r, camera_matrix_r, dist_coeffs_r, _, _ = cv2.calibrateCamera(obj_points, img_points[1], (w[1], h[1]), None, None)
    
    print "---------- Camera Left ------------"
    print "RMS:", rms_l
    print "camera matrix:\n", camera_matrix_l
    print "distortion coefficients: ", dist_coeffs_l.ravel()
    print "\n---------- Camera Right -----------"
    print "RMS:", rms_r
    print "camera matrix:\n", camera_matrix_r
    print "distortion coefficients: ", dist_coeffs_r.ravel()
    
    # set stereo flags
    stereo_flags = 0
    stereo_flags |= cv2.CALIB_USE_INTRINSIC_GUESS
    stereo_flags |= cv2.CALIB_FIX_INTRINSIC
    
    # More availabe flags...
    # stereo_flags |= cv2.CALIB_USE_INTRINSIC_GUESS    \
        # Refine intrinsic parameters
    # stereo_flags |= cv2.CALIB_FIX_PRINCIPAL_POINT    \
        # Fix the principal points during the optimization.
    # stereo_flags |= cv2.CALIB_FIX_FOCAL_LENGTH       \
        # Fix focal length
    # stereo_flags |= cv2.CALIB_FIX_ASPECT_RATIO       \
        # fix aspect ratio
    # stereo_flags |= cv2.CALIB_SAME_FOCAL_LENGTH      \
        # Use same focal length
    # stereo_flags |= cv2.CALIB_ZERO_TANGENT_DIST      \
        # Set tangential distortion to zero
    # stereo_flags |= cv2.CALIB_RATIONAL_MODEL         \
        # Use 8 param

    
    # run stereo calibration
    rms_stereo, camera_matrix_l, dist_coeffs_l, camera_matrix_r, dist_coeffs_r, R, T, E, F  = cv2.stereoCalibrate(
        obj_points, img_points[0], img_points[1], camera_matrix_l, dist_coeffs_l, camera_matrix_r, dist_coeffs_r, (w[0], h[0]), flags=stereo_flags)

    # run stereo rectification
    rectification_matrix_l, rectification_matrix_r, projection_matrix_l, projection_matrix_r, _, _, _ = cv2.stereoRectify(camera_matrix_l, dist_coeffs_l, camera_matrix_r, dist_coeffs_r, (w[0], h[0]), R, T)
    
    print "\n---------- Camera Stereo ----------"
    print "RMS:", rms_stereo
    print "camera matrix left:\n", camera_matrix_l
    print "distortion coefficients left: ", dist_coeffs_l.ravel()
    print "camera matrix right:\n", camera_matrix_r
    print "distortion coefficients right: ", dist_coeffs_r.ravel()
    print "R:\n", R
    #print "T:\n", T
    #print "E:\n", E
    #print "F:\n", F
    print "Rectification matrix Left :\n", rectification_matrix_l
    print "Rectification matrix Right :\n", rectification_matrix_r
    print "Projection Matrix Left :\n", projection_matrix_l
    print "Projection Matrix Right :\n", projection_matrix_r

    if save_name:
        data = {"camera_matrix_l": camera_matrix_l.tolist(), "camera_matrix_r": camera_matrix_r.tolist(),
				"dist_coeffs_l": dist_coeffs_l.tolist(), "dist_coeffs_r": dist_coeffs_r.tolist(),
				"rectification_matrix_l": rectification_matrix_l.tolist(), "rectification_matrix_r": rectification_matrix_r.tolist(),
				"projection_matrix_l": projection_matrix_l.tolist(), "projection_matrix_r": projection_matrix_r.tolist(),
				"w": w, "h": h}
        with open(save_name, "w") as f:
            json.dump(data, f)
    
    
    cv2.destroyAllWindows()