# Calculates comparison output values for DistortionMapperTest.cpp:CompareToOpenCV
#
# Assumes a python that has numpy and cv2 (OpenCV) available

import numpy as np
import cv2

Fx = 1000
Fy = 1000
Cx = 500
Cy = 500
# s = 0 - not supported by OpenCV

K = np.array([[Fx, 0, Cx],[0, Fy, Cy],[0, 0, 1]])

# Order is k1, k2, t1, t2, k3
dist = np.array([0.1, -0.003, 0.02, 0.01, 0.004])

np.random.seed(1234)

activeArray = np.array([[1000, 750]])

rawCoords = np.floor(np.random.rand(1000,2) * activeArray)

# OpenCV needs either row count or col count = 1 for some reason
rawCoords2 = rawCoords.reshape(-1, 1, 2)

# P is the output camera matrix, K is the input; use the same for both
expCoords = cv2.undistortPoints(rawCoords2, K, dist, P = K)

with open('DistortionMapperTest_OpenCvData.h','w') as f:
  f.write('// Generated by DistortionMapperComp.py\n');
  f.write('// for use by DistortionMapperTest.cpp\n\n');

  f.write('namespace openCvData {\n')
  f.write('std::array<int32_t, %d> rawCoords = {\n' % (rawCoords.shape[0] * rawCoords.shape[1]))
  for i in range(rawCoords.shape[0]):
    f.write('  %d, %d,\n' % (rawCoords[i][0], rawCoords[i][1]))
  f.write('};\n')

  f.write('std::array<int32_t, %d> expCoords = {\n' % (expCoords.shape[0] * expCoords.shape[2]))
  for i in range(expCoords.shape[0]):
    f.write('  %d, %d,\n' % (expCoords[i][0][0], expCoords[i][0][1]))
  f.write('};\n')
  f.write('} // namespace openCvData\n')

print "DistortionMapperTest_OpenCvData.h generated"