# Path: lfd/tpsopt/registration.py
# def unit_boxify(x_na):    
#     ranges = x_na.ptp(axis=0)
#     dlarge = ranges.argmax()
#     unscaled_translation = - (x_na.min(axis=0) + x_na.max(axis=0))/2
#     scaling = 1./ranges[dlarge]
#     scaled_translation = unscaled_translation * scaling
#     return x_na*scaling + scaled_translation, (scaling, scaled_translation)    
# 
# def loglinspace(a,b,n):
#     "n numbers between a to b (inclusive) with constant ratio between consecutive numbers"
#     return np.exp(np.linspace(np.log(a),np.log(b),n))    
# 
# Path: lfd/rapprentice/clouds.py
# DEFAULT_F = 535.
#     X = (x - cx)*(Z/f)
#     Y = (y - cy)*(Z/f)
#     XYZ = np.empty((480,640,3))
#     Z = XYZ[:,:,2] = depth / 1000. # convert mm -> meters
#     XYZ[:,:,0] = (x - cx)*(Z/f)
#     XYZ[:,:,1] = (y - cy)*(Z/f)
# def xyZ_to_XY(x,y,Z,f=DEFAULT_F):
# def XYZ_to_xy(X,Y,Z,f=DEFAULT_F):
# def depth_to_xyz(depth,f=DEFAULT_F):
# def downsample(xyz, v):


import h5py
import argparse
import sys
import numpy as np
import scipy.linalg
import scikits.cuda.linalg
from pycuda import gpuarray
from scikits.cuda.linalg import pinv as cu_pinv
from settings import GRIPPER_OPEN_CLOSE_THRESH
from tps import tps_kernel_matrix, tps_kernel_matrix2
from lfd.tpsopt.registration import unit_boxify, loglinspace
from lfd.rapprentice import clouds
from culinalg_exts import get_gpu_ptrs, dot_batch_nocheck, m_dot_batch
from settings import N_ITER_CHEAP, DEFAULT_LAMBDA, DS_SIZE, BEND_COEF_DIGITS, EXACT_LAMBDA, N_ITER_EXACT
    N_ptr_gpu = get_gpu_ptrs(N_arr_gpu)
    
    dot_batch_nocheck(Q_arr_gpu, Q_arr_gpu, H_arr_gpu,
                      Q_ptr_gpu, Q_ptr_gpu, H_ptr_gpu,
                      transa = 'T')
    # N'HN
    NHN_arr_gpu, NHN_ptr_gpu = m_dot_batch((N_arr_gpu, N_ptr_gpu, 'T'),
                                           (H_arr_gpu, H_ptr_gpu, 'N'),
                                           (N_arr_gpu, N_ptr_gpu, 'N'))
    iH_arr = []
    for NHN in NHN_arr_gpu:
        iH_arr.append(scipy.linalg.inv(NHN.get()).copy())
    iH_arr_gpu = [gpuarray.to_gpu_async(iH) for iH in iH_arr]
    iH_ptr_gpu = get_gpu_ptrs(iH_arr_gpu)

    proj_mats   = m_dot_batch((N_arr_gpu,  N_ptr_gpu,   'N'),
                              (iH_arr_gpu, iH_ptr_gpu, 'N'),
                              (N_arr_gpu,  N_ptr_gpu,   'T'),
                              (Q_arr_gpu,  Q_ptr_gpu,   'T'))

    offset_mats = m_dot_batch((N_arr_gpu,  N_ptr_gpu,   'N'),
                              (iH_arr_gpu, iH_ptr_gpu, 'N'),
                              (N_arr_gpu,  N_ptr_gpu,   'T'),
                              (F_arr_gpu,  F_ptr_gpu,   'N'))

    return proj_mats, offset_mats

def test_batch_get_sol_params(f, bend_coefs, rot_coef, atol=1e-7, index=0):
    seg_info = f.items()[index][1]
    inv_group =  seg_info['inv']

# Path: lfd/tpsopt/registration.py
# def unit_boxify(x_na):    
#     ranges = x_na.ptp(axis=0)
#     dlarge = ranges.argmax()
#     unscaled_translation = - (x_na.min(axis=0) + x_na.max(axis=0))/2
#     scaling = 1./ranges[dlarge]
#     scaled_translation = unscaled_translation * scaling
#     return x_na*scaling + scaled_translation, (scaling, scaled_translation)    
# 
# def loglinspace(a,b,n):
#     "n numbers between a to b (inclusive) with constant ratio between consecutive numbers"
#     return np.exp(np.linspace(np.log(a),np.log(b),n))    
# 
# Path: lfd/rapprentice/clouds.py
# DEFAULT_F = 535.
#     X = (x - cx)*(Z/f)
#     Y = (y - cy)*(Z/f)
#     XYZ = np.empty((480,640,3))
#     Z = XYZ[:,:,2] = depth / 1000. # convert mm -> meters
#     XYZ[:,:,0] = (x - cx)*(Z/f)
#     XYZ[:,:,1] = (y - cy)*(Z/f)
# def xyZ_to_XY(x,y,Z,f=DEFAULT_F):
# def XYZ_to_xy(X,Y,Z,f=DEFAULT_F):
# def depth_to_xyz(depth,f=DEFAULT_F):
# def downsample(xyz, v):


# Path: lfd/tpsopt/precompute.py
import h5py
import argparse
import sys
import numpy as np
import scipy.linalg
import scikits.cuda.linalg
from pycuda import gpuarray
from scikits.cuda.linalg import pinv as cu_pinv
from settings import GRIPPER_OPEN_CLOSE_THRESH
from tps import tps_kernel_matrix, tps_kernel_matrix2
from lfd.tpsopt.registration import unit_boxify, loglinspace
from lfd.rapprentice import clouds
from culinalg_exts import get_gpu_ptrs, dot_batch_nocheck, m_dot_batch
from settings import N_ITER_CHEAP, DEFAULT_LAMBDA, DS_SIZE, BEND_COEF_DIGITS, EXACT_LAMBDA, N_ITER_EXACT

    N_ptr_gpu = get_gpu_ptrs(N_arr_gpu)
    
    dot_batch_nocheck(Q_arr_gpu, Q_arr_gpu, H_arr_gpu,
                      Q_ptr_gpu, Q_ptr_gpu, H_ptr_gpu,
                      transa = 'T')
    # N'HN
    NHN_arr_gpu, NHN_ptr_gpu = m_dot_batch((N_arr_gpu, N_ptr_gpu, 'T'),
                                           (H_arr_gpu, H_ptr_gpu, 'N'),
                                           (N_arr_gpu, N_ptr_gpu, 'N'))
    iH_arr = []
    for NHN in NHN_arr_gpu:
        iH_arr.append(scipy.linalg.inv(NHN.get()).copy())
    iH_arr_gpu = [gpuarray.to_gpu_async(iH) for iH in iH_arr]
    iH_ptr_gpu = get_gpu_ptrs(iH_arr_gpu)

    proj_mats   = m_dot_batch((N_arr_gpu,  N_ptr_gpu,   'N'),
                              (iH_arr_gpu, iH_ptr_gpu, 'N'),
                              (N_arr_gpu,  N_ptr_gpu,   'T'),
                              (Q_arr_gpu,  Q_ptr_gpu,   'T'))

    offset_mats = m_dot_batch((N_arr_gpu,  N_ptr_gpu,   'N'),
                              (iH_arr_gpu, iH_ptr_gpu, 'N'),
                              (N_arr_gpu,  N_ptr_gpu,   'T'),
                              (F_arr_gpu,  F_ptr_gpu,   'N'))

    return proj_mats, offset_mats

def test_batch_get_sol_params(f, bend_coefs, rot_coef, atol=1e-7, index=0):
    seg_info = f.items()[index][1]
    inv_group =  seg_info['inv']

    ds_key = 'DS_SIZE_{}'.format(DS_SIZE)