# Path: lfd/tpsopt/transformations.py
# class ThinPlateSpline(Transformation):
#     """
#     members:
#         x_na: centers of basis functions
#         w_ng: 
#         lin_ag: transpose of linear part, so you take x_na.dot(lin_ag)
#         trans_g: translation part
#     
#     """
#     def __init__(self, d=3):
#         "initialize as identity"
#         self.x_na = np.zeros((0,d))
#         self.lin_ag = np.eye(d)
#         self.trans_g = np.zeros(d)
#         self.w_ng = np.zeros((0,d))
# 
#     def transform_points(self, x_ma):
#         y_ng = tps.tps_eval(x_ma, self.lin_ag, self.trans_g, self.w_ng, self.x_na)
#         return y_ng
#     def compute_jacobian(self, x_ma):
#         grad_mga = tps.tps_grad(x_ma, self.lin_ag, self.trans_g, self.w_ng, self.x_na)
#         return grad_mga
# 
# def fit_ThinPlateSpline(x_na, y_ng, bend_coef=.1, rot_coef = 1e-5, wt_n=None):
#     """
#     x_na: source cloud
#     y_nd: target cloud
#     smoothing: penalize non-affine part
#     angular_spring: penalize rotation
#     wt_n: weight the points        
#     """
#     f = ThinPlateSpline()
#     f.lin_ag, f.trans_g, f.w_ng = tps.tps_fit3(x_na, y_ng, bend_coef, rot_coef, wt_n)
#     f.x_na = x_na
#     return f


import numpy as np
import scipy.spatial.distance as ssd
import tps
from lfd.tpsopt.transformations import ThinPlateSpline, fit_ThinPlateSpline
from settings import BEND_COEF_DIGITS
    
    g = ThinPlateSpline(d)
    g.lin_ag = np.diag(1./scale)
    g.trans_g = -np.diag(1./scale).dot(f.trans_g)


    # r_N = None
    
    for i in xrange(n_iter):
        xwarped_nd = f.transform_points(x_nd)
        ywarped_md = g.transform_points(y_md)
        
        fwddist_nm = ssd.cdist(xwarped_nd, y_md,'euclidean')
        invdist_nm = ssd.cdist(x_nd, ywarped_md,'euclidean')
        
        r = rads[i]
        prob_nm = np.exp( -(fwddist_nm + invdist_nm) / (2*r) )
        corr_nm, r_N, _ =  balance_matrix(prob_nm, 10, outlierprior, outlierfrac)
        corr_nm += 1e-9
        
        wt_n = corr_nm.sum(axis=1)
        wt_m = corr_nm.sum(axis=0)

        xtarg_nd = (corr_nm/wt_n[:,None]).dot(y_md)
        ytarg_md = (corr_nm/wt_m[None,:]).T.dot(x_nd)                

        fsolve.solve(wt_n, xtarg_nd, regs[i], rot_reg, f)        
        gsolve.solve(wt_m, ytarg_md, regs[i], rot_reg, g)
        if check_solver:
            f_test = fit_ThinPlateSpline(x_nd, xtarg_nd, bend_coef = regs[i], wt_n=wt_n, rot_coef = rot_reg)

# Path: lfd/tpsopt/transformations.py
# class ThinPlateSpline(Transformation):
#     """
#     members:
#         x_na: centers of basis functions
#         w_ng: 
#         lin_ag: transpose of linear part, so you take x_na.dot(lin_ag)
#         trans_g: translation part
#     
#     """
#     def __init__(self, d=3):
#         "initialize as identity"
#         self.x_na = np.zeros((0,d))
#         self.lin_ag = np.eye(d)
#         self.trans_g = np.zeros(d)
#         self.w_ng = np.zeros((0,d))
# 
#     def transform_points(self, x_ma):
#         y_ng = tps.tps_eval(x_ma, self.lin_ag, self.trans_g, self.w_ng, self.x_na)
#         return y_ng
#     def compute_jacobian(self, x_ma):
#         grad_mga = tps.tps_grad(x_ma, self.lin_ag, self.trans_g, self.w_ng, self.x_na)
#         return grad_mga
# 
# def fit_ThinPlateSpline(x_na, y_ng, bend_coef=.1, rot_coef = 1e-5, wt_n=None):
#     """
#     x_na: source cloud
#     y_nd: target cloud
#     smoothing: penalize non-affine part
#     angular_spring: penalize rotation
#     wt_n: weight the points        
#     """
#     f = ThinPlateSpline()
#     f.lin_ag, f.trans_g, f.w_ng = tps.tps_fit3(x_na, y_ng, bend_coef, rot_coef, wt_n)
#     f.x_na = x_na
#     return f


# Path: lfd/tpsopt/registration.py
import numpy as np
import scipy.spatial.distance as ssd
import tps
from lfd.tpsopt.transformations import ThinPlateSpline, fit_ThinPlateSpline
from settings import BEND_COEF_DIGITS

    
    g = ThinPlateSpline(d)
    g.lin_ag = np.diag(1./scale)
    g.trans_g = -np.diag(1./scale).dot(f.trans_g)


    # r_N = None
    
    for i in xrange(n_iter):
        xwarped_nd = f.transform_points(x_nd)
        ywarped_md = g.transform_points(y_md)
        
        fwddist_nm = ssd.cdist(xwarped_nd, y_md,'euclidean')
        invdist_nm = ssd.cdist(x_nd, ywarped_md,'euclidean')
        
        r = rads[i]
        prob_nm = np.exp( -(fwddist_nm + invdist_nm) / (2*r) )
        corr_nm, r_N, _ =  balance_matrix(prob_nm, 10, outlierprior, outlierfrac)
        corr_nm += 1e-9
        
        wt_n = corr_nm.sum(axis=1)
        wt_m = corr_nm.sum(axis=0)

        xtarg_nd = (corr_nm/wt_n[:,None]).dot(y_md)
        ytarg_md = (corr_nm/wt_m[None,:]).T.dot(x_nd)                

        fsolve.solve(wt_n, xtarg_nd, regs[i], rot_reg, f)        
        gsolve.solve(wt_m, ytarg_md, regs[i], rot_reg, g)
        if check_solver:
            f_test = fit_ThinPlateSpline(x_nd, xtarg_nd, bend_coef = regs[i], wt_n=wt_n, rot_coef = rot_reg)

            g_test = fit_ThinPlateSpline(y_md, ytarg_md, bend_coef = regs[i], wt_n=wt_m, rot_coef = rot_reg)