def fit_manifold(data_to_fit, fit_params):
    '''fit_params takes nKnots : number of knots, dalpha : resolution for
    sampled curve, knot_order : method to initially order knots, penalty_type : 
    penalty'''
    # fit_params is a superset of the initial params that PiecewiseLinearFit needs
    fitter = PiecewiseLinearFit(data_to_fit, fit_params)
    unord_knots = fitter.get_new_initial_knots()
    init_knots = fitter.order_knots(unord_knots,
                                    method=fit_params['knot_order'])
    curr_fit_params = {
        'init_knots': init_knots,
        'penalty_type': fit_params['penalty_type']
    }
    fitter.fit_data(curr_fit_params)
    fit_results = dict(fit_params)
    fit_results['init_knots'] = init_knots
    # The fit class appends the results of each fit to a list called saved_knots
    # Here we're just using the class once, hence saved_knots[0]
    fit_results['final_knots'] = fitter.saved_knots[0]['knots']
    fit_results['fit_err'] = fitter.saved_knots[0]['err']
    fit_results['loop_final_knots'] = fhf.loop_knots(
        fitter.saved_knots[0]['knots'])
    fit_results['tt'], fit_results['curve'] = fhf.get_curve_from_knots(
        fit_results['loop_final_knots'], 'eq_vel')
    return fit_results
    def get_curve_from_knots_internal(self, inp_knots):
        '''Turn a list of knots into a curve, sampled at the pre-specified
        resolution.'''

        # Repeat the first knot at the end so we get a loop.
        loop_knots = fhf.loop_knots(inp_knots)
        return loop_knots[self.t_int_idx] + (
            loop_knots[self.t_int_idx + 1] -
            loop_knots[self.t_int_idx]) * self.t_rsc[:, np.newaxis]
        def cost_fn(flat_knots, penalty_params=fit_params):
            knots = np.reshape(flat_knots.copy(), (self.nKnots, self.nDims))
            loop_knots = fhf.loop_knots(
                np.reshape(flat_knots.copy(), (self.nKnots, self.nDims)))
            fit_curve = loop_knots[self.t_int_idx] + (
                loop_knots[self.t_int_idx + 1] -
                loop_knots[self.t_int_idx]) * self.t_rsc[:, np.newaxis]
            neighbgraph = NearestNeighbors(n_neighbors=1).fit(fit_curve)
            dists, inds = neighbgraph.kneighbors(self.data_to_fit)

            if penalty_params['penalty_type'] == 'none':
                cost = np.sum(dists)
            elif penalty_params['penalty_type'] == 'mult_len':
                cost = np.sum(dists) * self.tot_len(loop_knots)
            elif penalty_params['penalty_type'] == 'add_len':
                cost = np.mean(dists) + penalty_params[
                    'len_coeff'] * self.tot_len(loop_knots)
            return cost
data_to_fit = curr_mani[train_idx].copy()
data_to_decode = curr_mani[test_idx].copy()

# Automated fit does:
# fit_result = mff.fit_manifold(data_to_fit, fit_params)
# And then matches it to a reference angle using mff.decode_from_passed_fit.
# Here let's look at the knot placement and pick a good one.
# Set up the fit class and initial knots
fitter = mff.PiecewiseLinearFit(data_to_fit, fit_params)
unord_knots = fitter.get_new_initial_knots()

# Generate knots and look at them to make sure we're happy with them. Mess around
# with number of knots etc. if fits look bad.
# t0 = time.time()
init_knots = fitter.order_knots(unord_knots, method=fit_params['knot_order'])
loop_init = fhf.loop_knots(init_knots)
init_tt, init_curve = fhf.get_curve_from_knots(loop_init, 'eq_vel')
# print 'Time ', time.time()-t0

fig = plt.figure(figsize=(9, 9))
ax = fig.add_subplot(111, projection='3d')
ax.scatter(data_to_fit[::2, 0],
           data_to_fit[::2, 1],
           data_to_fit[::2, 2],
           c='r',
           s=5)
ax.plot(loop_init[:, 0], loop_init[:, 1], loop_init[:, 2], c='k', lw=2)
plt.show()

# Now optimize to get the best fit
t0 = time.time()
    k = ('Mouse28-140313', 3, 15, 'wt_per_len', 'mult_len', 0.8)
    # Or do k = fit_data['fit_results'].keys()[0] or just pick one
    sel_knots = fit_data['fit_results'][k][0][-1]
elif fit_source == 'single_interactive':
    # This version from single interactive fit
    dd = gen_params['results_dir'] + '2019_06_03_interactive_curve_fits/'
    file_pattern = '%s_%s_dim%d_trainfrac%.2f_interactive_fits_sd*.p'%(
        session, state, embed_dim, train_frac)
    fit_data, fit_fname = gff.load_file_from_pattern(dd + file_pattern)
    sel_knots = fit_data['fit_results']['final_knots']

assert fit_data['embed_file'] == embed_fname, 'Fit seems to be to different data'
meas_angles = embed['meas_angles']

data_to_decode = curr_mani
loop_sel_knots = fhf.loop_knots(sel_knots)
tt, curve = fhf.get_curve_from_knots(loop_sel_knots, 'eq_vel')
dec_angle, mse = mfd.decode_from_passed_fit(data_to_decode, tt[:-1], 
    curve[:-1], meas_angles)

si = 2000; ei = 6000
fig = plt.figure(figsize=(12,4))
ax = fig.add_subplot(111)
ax.plot(meas_angles[si:ei], color='k', lw=2, label='Measured')
ax.plot(dec_angle[si:ei], color=cols['Wake'], lw=2, label='Decoded')
ax.set_xlabel('Time (s)')
ax.set_ylabel('Angle (rad)')
ax.legend()
plt.show()

# Plot diffusion curve for decoded angle