def test_basic_2d_affine():
    linear_component = np.array([[1, -6],
                                 [-3, 2]])
    translation_component = np.array([7, -8])
    h_matrix = np.eye(3, 3)
    h_matrix[:-1, :-1] = linear_component
    h_matrix[:-1, -1] = translation_component
    affine = Affine(h_matrix)
    x = np.array([[0, 1],
                  [1, 1],
                  [-1, -5],
                  [3, -5]])
    # transform x explicitly
    solution = np.dot(x, linear_component.T) + translation_component
    # transform x using the affine transform
    result = affine.apply(x)
    # check that both answers are equivalent
    assert_allclose(solution, result)
    # create several copies of x
    x_copies = np.array([x, x, x, x, x, x, x, x])
    # transform all of copies at once using the affine transform
    results = affine.apply(x_copies)
    # check that all copies have been transformed correctly
    for r in results:
        assert_allclose(solution, r)

def test_align_2d_affine():
    linear_component = np.array([[1, -6], [-3, 2]])
    translation_component = np.array([7, -8])
    h_matrix = np.eye(3, 3)
    h_matrix[:-1, :-1] = linear_component
    h_matrix[:-1, -1] = translation_component
    affine = Affine(h_matrix)
    source = PointCloud(np.array([[0, 1], [1, 1], [-1, -5], [3, -5]]))
    target = affine.apply(source)
    # estimate the transform from source and target
    estimate = AlignmentAffine(source, target)
    # check the estimates is correct
    assert_allclose(affine.h_matrix, estimate.h_matrix)

def extract_targets(frame, target_centre, target_shape, n_perturbations=10,
                    noise_std=0.04):
    # initialize targets
    w, h = target_shape
    targets = np.empty((n_perturbations + 1, frame.n_channels, w, h))

    # extract original target
    targets[0] = frame.extract_patches(
        target_centre, patch_size=target_shape,
        as_single_array=True)

    for j in range(n_perturbations):
        # perturb identity affine transform
        params = noise_std * np.random.randn(6)
        transform = Affine.init_identity(2).from_vector(params)
        # warp frame using previous affine transform
        perturbed_frame = frame.warp_to_mask(frame.as_masked().mask,
                                             transform)
        # apply inverse of affine transform to target centre
        perturbed_centre = transform.pseudoinverse().apply(target_centre)
        # extract perturbed target + context region from frame
        perturbed_target = perturbed_frame.extract_patches(
            perturbed_centre, patch_size=target_shape,
            as_single_array=True)
        # store target
        targets[j+1] = perturbed_target

    return targets

def test_c_warp_gray():
    target_transform = Affine.identity(2).from_vector(initial_params)
    warped_im = gray_image.warp_to(template_mask, target_transform,
                                   interpolator='c')

    assert(warped_im.shape == gray_template.shape)
    assert_allclose(warped_im.pixels, gray_template.pixels)

def test_warp_multi():
    rgb_image = mio.import_builtin_asset('takeo.ppm')
    target_transform = Affine.init_identity(2).from_vector(initial_params)
    warped_im = rgb_image.warp_to_mask(template_mask, target_transform)

    assert(warped_im.shape == rgb_template.shape)
    assert_allclose(warped_im.pixels, rgb_template.pixels)

    def fit_from_bb(self, image, bounding_box, gt_shape=None, **kwargs):
        algo_result = self.algorithm.run(image, bounding_box, gt_shape=gt_shape)

        # TODO: This should be a basic result instead.
        return MultiFitterResult(image, self, [algo_result],
                                 Affine.init_identity(2),
                                 gt_shape=gt_shape)

def test_warp_to_mask_masked_image_all_true():
    img = MaskedImage.init_blank((10, 10), fill=2.5)

    template_mask = BooleanImage.init_blank((10, 10), fill=False)
    template_mask.pixels[:, :5, :5] = True
    t = Affine.init_identity(2)
    warped_img = img.warp_to_mask(template_mask, t)
    assert(type(warped_img) == MaskedImage)

def test_warp_gray_batch():
    rgb_image = mio.import_builtin_asset('takeo.ppm')
    gray_image = rgb_image.as_greyscale()
    target_transform = Affine.init_identity(2).from_vector(initial_params)
    warped_im = gray_image.warp_to_mask(template_mask, target_transform,
                                        batch_size=100)

    assert(warped_im.shape == gray_template.shape)
    assert_allclose(warped_im.pixels, gray_template.pixels)

def residual_wrapper(residual, algorithm, interpolator, expected_error):
    image, template, initial_params = setup_conditions(interpolator)
    align_algorithm = algorithm(
        template, residual, Affine.identity(2).from_vector(
            initial_params))
    fitting = align_algorithm.fit(image, initial_params)
    transform = fitting.final_transform
    rms_error = compute_fixed_error(transform)
    assert_approx_equal(rms_error, expected_error)

def setup_error():
    target_transform = Affine.identity(2).from_vector(target_params)
    original_box = np.array([[0,               0],
                             [target_shape[0], 0],
                             [target_shape[0], target_shape[1]],
                             [0,               target_shape[1]]]).T
    target_pts = target_transform.apply(original_box.T)

    return target_pts, original_box

def test_warp_to_mask_boolean():
    b = BooleanImage.init_blank((10, 10))
    b.pixels[:, :5] = False
    template_mask = BooleanImage.init_blank((10, 10))
    template_mask.pixels[:5, :] = False
    t = Affine.init_identity(2)
    warped_mask = b.warp_to_mask(template_mask, t)
    assert(type(warped_mask) == BooleanImage)
    result = template_mask.pixels.copy()
    result[:, :5] = False
    assert(np.all(result == warped_mask.pixels))

def test_warp_to_mask_image():
    img = Image.init_blank((10, 10), n_channels=2)
    img.pixels[:, :, :5] = 0.5
    template_mask = BooleanImage.init_blank((10, 10))
    template_mask.pixels[:, 5:, :] = False
    t = Affine.init_identity(2)
    warped_img = img.warp_to_mask(template_mask, t)
    assert(type(warped_img) == MaskedImage)
    result = Image.init_blank((10, 10), n_channels=2).pixels
    result[:, :5, :5] = 0.5
    assert(np.all(result == warped_img.pixels))

def test_affine_jacobian_2d_with_positions():
    params = np.array([0, 0.1, 0.2, 0, 30, 70])
    t = Affine.identity(2).from_vector(params)
    explicit_pixel_locations = np.array(
        [[0, 0],
        [0, 1],
        [0, 2],
        [1, 0],
        [1, 1],
        [1, 2]])
    dW_dp = t.d_dp(explicit_pixel_locations)
    assert_equal(dW_dp, jac_solution2d)

def test_affine_jacobian_3d_with_positions():
    params = np.ones(12)
    t = Affine.identity(3).from_vector(params)
    explicit_pixel_locations = np.array(
        [[0, 0, 0],
        [0, 0, 1],
        [0, 1, 0],
        [0, 1, 1],
        [0, 2, 0],
        [0, 2, 1],
        [1, 0, 0],
        [1, 0, 1],
        [1, 1, 0],
        [1, 1, 1],
        [1, 2, 0],
        [1, 2, 1]])
    dW_dp = t.d_dp(explicit_pixel_locations)
    assert_equal(dW_dp, jac_solution3d)

def test_warp_to_mask_masked_image():
    mask = BooleanImage.blank((10, 10))
    # make a funny mask on the original image
    mask.pixels[2:, :] = False
    img = MaskedImage.blank((10, 10), n_channels=2, mask=mask)
    img.pixels[...] = 2.5
    template_mask = BooleanImage.blank((10, 10), fill=False)
    template_mask.pixels[:5, :5] = True
    t = Affine.identity(2)
    warped_img = img.warp_to_mask(template_mask, t)
    assert(type(warped_img) == MaskedImage)
    result = Image.blank((10, 10), n_channels=2).pixels
    result[:5, :5, :] = 2.5
    result_mask = BooleanImage.blank((10, 10), fill=False).pixels
    result_mask[:2, :5] = True
    assert(warped_img.n_true_pixels() == 10)
    assert(np.all(result == warped_img.pixels))
    assert(np.all(result_mask == warped_img.mask.pixels))

def test_warp_to_mask_masked_image():
    mask = BooleanImage.init_blank((15, 15))
    # make a truncated mask on the original image
    mask.pixels[0, -1, -1] = False
    img = MaskedImage.init_blank((15, 15), n_channels=2, mask=mask,
                                 fill=2.5)
    template_mask = BooleanImage.init_blank((10, 10), fill=False)
    template_mask.pixels[:, :5, :5] = True
    t = Affine.init_identity(2)
    warped_img = img.warp_to_mask(template_mask, t)
    assert(type(warped_img) == MaskedImage)

    result = Image.init_blank((10, 10), n_channels=2).pixels
    result[:, :5, :5] = 2.5
    result_mask = BooleanImage.init_blank((10, 10), fill=False).pixels
    result_mask[:, :5, :5] = True
    assert(warped_img.n_true_pixels() == 25)
    assert_allclose(result, warped_img.pixels)
    assert_allclose(result_mask, warped_img.mask.pixels)

def test_affine_identity_3d():
    assert_allclose(Affine.identity(3).h_matrix, np.eye(4))

def test_affine_compose_inplace_affine():
    a = Affine.init_identity(2)
    b = Affine.init_identity(2)
    a.compose_before_inplace(b)
    assert(np.all(a.h_matrix == b.h_matrix))

def test_affine_pseudoinverse():
    s = NonUniformScale([4, 3])
    inv_man = NonUniformScale([1./4, 1./3])
    b = Affine(s.h_matrix)
    i = b.pseudoinverse()
    assert_allclose(i.h_matrix, inv_man.h_matrix)

def test_affine_identity_2d():
    assert_allclose(Affine.init_identity(2).h_matrix, np.eye(3))