# 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
    lin_in = np.eye(d)*p
    trans_in = q
    aff_in = Affine(lin_in, trans_in)
    
    lin_out = np.eye(d)/r
    trans_out = -s/r
    aff_out = Affine(lin_out, trans_out)

    return Composition([aff_in, f, aff_out])
        
# @profile
def tps_rpm_bij(x_nd, y_md, fsolve, gsolve,n_iter = 20, reg_init = .1, reg_final = .001, rad_init = .1,
                rad_final = .005, rot_reg = 1e-3, outlierprior=1e-1, outlierfrac=2e-1, vis_cost_xy=None,
                return_corr=False, check_solver=False):
    """
    tps-rpm algorithm mostly as described by chui and rangaran
    reg_init/reg_final: regularization on curvature
    rad_init/rad_final: radius for correspondence calculation (meters)
    plotting: 0 means don't plot. integer n means plot every n iterations
    """
    
    _,d=x_nd.shape
    regs = np.around(loglinspace(reg_init, reg_final, n_iter), BEND_COEF_DIGITS)
    rads = loglinspace(rad_init, rad_final, n_iter)

    f = ThinPlateSpline(d)
    scale = (np.max(y_md,axis=0) - np.min(y_md,axis=0)) / (np.max(x_nd,axis=0) - np.min(x_nd,axis=0))
    f.lin_ag = np.diag(scale) # align the mins and max
    f.trans_g = np.median(y_md,axis=0) - np.median(x_nd,axis=0) * scale  # align the medians
    

# 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

    lin_in = np.eye(d)*p
    trans_in = q
    aff_in = Affine(lin_in, trans_in)
    
    lin_out = np.eye(d)/r
    trans_out = -s/r
    aff_out = Affine(lin_out, trans_out)

    return Composition([aff_in, f, aff_out])
        
# @profile
def tps_rpm_bij(x_nd, y_md, fsolve, gsolve,n_iter = 20, reg_init = .1, reg_final = .001, rad_init = .1,
                rad_final = .005, rot_reg = 1e-3, outlierprior=1e-1, outlierfrac=2e-1, vis_cost_xy=None,
                return_corr=False, check_solver=False):
    """
    tps-rpm algorithm mostly as described by chui and rangaran
    reg_init/reg_final: regularization on curvature
    rad_init/rad_final: radius for correspondence calculation (meters)
    plotting: 0 means don't plot. integer n means plot every n iterations
    """
    
    _,d=x_nd.shape
    regs = np.around(loglinspace(reg_init, reg_final, n_iter), BEND_COEF_DIGITS)
    rads = loglinspace(rad_init, rad_final, n_iter)

    f = ThinPlateSpline(d)
    scale = (np.max(y_md,axis=0) - np.min(y_md,axis=0)) / (np.max(x_nd,axis=0) - np.min(x_nd,axis=0))
    f.lin_ag = np.diag(scale) # align the mins and max
    f.trans_g = np.median(y_md,axis=0) - np.median(x_nd,axis=0) * scale  # align the medians
    

    g = ThinPlateSpline(d)