/*

 * Copyright (c) 2016, Alliance for Open Media. All rights reserved

 *

 * This source code is subject to the terms of the BSD 2 Clause License and

 * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License

 * was not distributed with this source code in the LICENSE file, you can

 * obtain it at https://www.aomedia.org/license/software-license. If the Alliance for Open

 * Media Patent License 1.0 was not distributed with this source code in the

 * PATENTS file, you can obtain it at https://www.aomedia.org/license/patent-license.

 */

#include <memory.h>

#include <math.h>

#include <stdlib.h>

#include <assert.h>

#include "ransac.h"

#include "mathutils.h"

#include "random.h"

#include "common_dsp_rtcd.h"

#include "utility.h"

#define MAX_MINPTS 4

#define MAX_DEGENERATE_ITER 10

#define MINPTS_MULTIPLIER 5

#define INLIER_THRESHOLD_SQUARED 1.5625 /*(1.25 * 1.25)*/

// Number of initial models to generate

#define NUM_TRIALS 20

// Number of times to refine the best model found

#define NUM_REFINES 5

#define MIN_TRIALS 20

////////////////////////////////////////////////////////////////////////////////

// ransac

// Return -1 if 'a' is a better motion, 1 if 'b' is better, 0 otherwise.

static int compare_motions(const void* arg_a, const void* arg_b) {

    const RANSAC_MOTION* motion_a = (RANSAC_MOTION*)arg_a;

    const RANSAC_MOTION* motion_b = (RANSAC_MOTION*)arg_b;

    if (motion_a->num_inliers > motion_b->num_inliers) {

        return -1;

    }

    if (motion_a->num_inliers < motion_b->num_inliers) {

        return 1;

    }

    if (motion_a->sse < motion_b->sse) {

        return -1;

    }

    if (motion_a->sse > motion_b->sse) {

        return 1;

    }

    return 0;

}

static int is_better_motion(const RANSAC_MOTION* motion_a, const RANSAC_MOTION* motion_b) {

    return compare_motions(motion_a, motion_b) < 0;

}

static void score_translation(const double* mat, const Correspondence* points, int num_points, RANSAC_MOTION* model) {

    model->num_inliers = 0;

    model->sse         = 0.0;

    for (int i = 0; i < num_points; ++i) {

        const double x1 = points[i].x;

        const double y1 = points[i].y;

        const double x2 = points[i].rx;

        const double y2 = points[i].ry;

        const double proj_x = x1 + mat[0];

        const double proj_y = y1 + mat[1];

        const double dx  = proj_x - x2;

        const double dy  = proj_y - y2;

        const double sse = dx * dx + dy * dy;

        if (sse < INLIER_THRESHOLD_SQUARED) {

            model->inlier_indices[model->num_inliers++] = i;

            model->sse += sse;

        }

    }

}

static void score_affine(const double* mat, const Correspondence* points, int num_points, RANSAC_MOTION* model) {

    model->num_inliers = 0;

    model->sse         = 0.0;

    for (int i = 0; i < num_points; ++i) {

        const double x1 = points[i].x;

        const double y1 = points[i].y;

        const double x2 = points[i].rx;

        const double y2 = points[i].ry;

        const double proj_x = mat[2] * x1 + mat[3] * y1 + mat[0];

        const double proj_y = mat[4] * x1 + mat[5] * y1 + mat[1];

        const double dx  = proj_x - x2;

        const double dy  = proj_y - y2;

        const double sse = dx * dx + dy * dy;

        if (sse < INLIER_THRESHOLD_SQUARED) {

            model->inlier_indices[model->num_inliers++] = i;

            model->sse += sse;

        }

    }

}

static bool find_translation(const Correspondence* points, const int* indices, int num_indices, double* params) {

    double sumx = 0;

    double sumy = 0;

    for (int i = 0; i < num_indices; ++i) {

        int          index = indices[i];

        const double sx    = points[index].x;

        const double sy    = points[index].y;

        const double dx    = points[index].rx;

        const double dy    = points[index].ry;

        sumx += dx - sx;

        sumy += dy - sy;

    }

    params[0] = sumx / num_indices;

    params[1] = sumy / num_indices;

    params[2] = 1;

    params[3] = 0;

    params[4] = 0;

    params[5] = 1;

    return true;

}

static bool find_rotzoom(const Correspondence* points, const int* indices, int num_indices, double* params) {

    const int n = 4; // Size of least-squares problem

    double    mat[4 * 4]; // Accumulator for A'A

    double    y[4]; // Accumulator for A'b

    double    a[4]; // Single row of A

    least_squares_init(mat, y, n);

    for (int i = 0; i < num_indices; ++i) {

        int          index = indices[i];

        const double sx    = points[index].x;

        const double sy    = points[index].y;

        const double dx    = points[index].rx;

        const double dy    = points[index].ry;

        a[0]     = 1;

        a[1]     = 0;

        a[2]     = sx;

        a[3]     = sy;

        double b = dx; // Single element of b

        least_squares_accumulate(mat, y, a, b, n);

        a[0] = 0;

        a[1] = 1;

        a[2] = sy;

        a[3] = -sx;

        b    = dy;

        least_squares_accumulate(mat, y, a, b, n);

    }

    // Fill in params[0] .. params[3] with output model

    if (!least_squares_solve(mat, y, params, n)) {

        return false;

    }

    // Fill in remaining parameters

    params[4] = -params[3];

    params[5] = params[2];

    return true;

}

static bool find_affine(const Correspondence* points, const int* indices, int num_indices, double* params) {

    // Note: The least squares problem for affine models is 6-dimensional,

    // but it splits into two independent 3-dimensional subproblems.

    // Solving these two subproblems separately and recombining at the end

    // results in less total computation than solving the 6-dimensional

    // problem directly.

    //

    // The two subproblems correspond to all the parameters which contribute

    // to the x output of the model, and all the parameters which contribute

    // to the y output, respectively.

    const int n = 3; // Size of each least-squares problem

    double    mat[2][3 * 3]; // Accumulator for A'A

    double    y[2][3]; // Accumulator for A'b

    double    x[2][3]; // Output vector

    double    a[2][3]; // Single row of A

    double    b[2]; // Single element of b

    least_squares_init(mat[0], y[0], n);

    least_squares_init(mat[1], y[1], n);

    for (int i = 0; i < num_indices; ++i) {

        int          index = indices[i];

        const double sx    = points[index].x;

        const double sy    = points[index].y;

        const double dx    = points[index].rx;

        const double dy    = points[index].ry;

        a[0][0] = 1;

        a[0][1] = sx;

        a[0][2] = sy;

        b[0]    = dx;

        least_squares_accumulate(mat[0], y[0], a[0], b[0], n);

        a[1][0] = 1;

        a[1][1] = sx;

        a[1][2] = sy;

        b[1]    = dy;

        least_squares_accumulate(mat[1], y[1], a[1], b[1], n);

    }

    if (!least_squares_solve(mat[0], y[0], x[0], n)) {

        return false;

    }

    if (!least_squares_solve(mat[1], y[1], x[1], n)) {

        return false;

    }

    // Rearrange least squares result to form output model

    params[0] = x[0][0];

    params[1] = x[1][0];

    params[2] = x[0][1];

    params[3] = x[0][2];

    params[4] = x[1][1];

    params[5] = x[1][2];

    return true;

}

// Returns true on success, false on error

static bool ransac_internal(const Correspondence* matched_points, int npoints, MotionModel* motion_models,

                            int num_desired_motions, const RansacModelInfo* model_info, bool* mem_alloc_failed) {

    assert(npoints >= 0);

    int  i       = 0;

    int  minpts  = model_info->minpts;

    bool ret_val = true;

    unsigned int seed = (unsigned int)npoints;

    int indices[MAX_MINPTS] = {0};

    // Store information for the num_desired_motions best transformations found

    // and the worst motion among them, as well as the motion currently under

    // consideration.

    RANSAC_MOTION *motions, *worst_kept_motion = NULL;

    RANSAC_MOTION  current_motion;

    // Store the parameters and the indices of the inlier points for the motion

    // currently under consideration.

    double params_this_motion[MAX_PARAMDIM];

    // Initialize output models, as a fallback in case we can't find a model

    for (i = 0; i < num_desired_motions; i++) {

        memcpy(motion_models[i].params, kIdentityParams, MAX_PARAMDIM * sizeof(*(motion_models[i].params)));

        motion_models[i].num_inliers = 0;

    }

    if (npoints < minpts * MINPTS_MULTIPLIER || npoints == 0) {

        return false;

    }

    int min_inliers = AOMMAX((int)(MIN_INLIER_PROB * npoints), minpts);

    EB_CALLOC_ARRAY_NO_CHECK(motions, num_desired_motions);

    // Allocate one large buffer which will be carved up to store the inlier

    // indices for the current motion plus the num_desired_motions many

    // output models

    // This allows us to keep the allocation/deallocation logic simple, without

    // having to (for example) check that `motions` is non-null before allocating

    // the inlier arrays

    int* inlier_buffer;

    EB_MALLOC_ARRAY_NO_CHECK(inlier_buffer, npoints * (num_desired_motions + 1));

    if (!(motions && inlier_buffer)) {

        ret_val           = false;

        *mem_alloc_failed = true;

        goto finish_ransac;

    }

    // Once all our allocations are known-good, we can fill in our structures

    worst_kept_motion = motions;

    for (i = 0; i < num_desired_motions; ++i) {

        motions[i].inlier_indices = inlier_buffer + i * npoints;

    }

    memset(&current_motion, 0, sizeof(current_motion));

    current_motion.inlier_indices = inlier_buffer + num_desired_motions * npoints;

    for (int trial_count = 0; trial_count < NUM_TRIALS; trial_count++) {

        lcg_pick(npoints, minpts, indices, &seed);

        if (!model_info->find_transformation(matched_points, indices, minpts, params_this_motion)) {

            continue;

        }

        model_info->score_model(params_this_motion, matched_points, npoints, &current_motion);

        if (current_motion.num_inliers < min_inliers) {

            // Reject models with too few inliers

            continue;

        }

        if (is_better_motion(&current_motion, worst_kept_motion)) {

            // This motion is better than the worst currently kept motion. Remember

            // the inlier points and sse. The parameters for each kept motion

            // will be recomputed later using only the inliers.

            worst_kept_motion->num_inliers = current_motion.num_inliers;

            worst_kept_motion->sse         = current_motion.sse;

            // Rather than copying the (potentially many) inlier indices from

            // current_motion.inlier_indices to worst_kept_motion->inlier_indices,

            // we can swap the underlying pointers.

            //

            // This is okay because the next time current_motion.inlier_indices

            // is used will be in the next trial, where we ignore its previous

            // contents anyway. And both arrays will be deallocated together at the

            // end of this function, so there are no lifetime issues.

            int* tmp                          = worst_kept_motion->inlier_indices;

            worst_kept_motion->inlier_indices = current_motion.inlier_indices;

            current_motion.inlier_indices     = tmp;

            // Determine the new worst kept motion and its num_inliers and sse.

            for (i = 0; i < num_desired_motions; ++i) {

                if (is_better_motion(worst_kept_motion, &motions[i])) {

                    worst_kept_motion = &motions[i];

                }

            }

        }

    }

    // Sort the motions, best first.

    qsort(motions, num_desired_motions, sizeof(RANSAC_MOTION), compare_motions);

    // Refine each of the best N models using iterative estimation.

    //

    // The idea here is loosely based on the iterative method from

    // "Locally Optimized RANSAC" by O. Chum, J. Matas and Josef Kittler:

    // https://cmp.felk.cvut.cz/ftp/articles/matas/chum-dagm03.pdf

    //

    // However, we implement a simpler version than their proposal, and simply

    // refit the model repeatedly until the number of inliers stops increasing,

    // with a cap on the number of iterations to defend against edge cases which

    // only improve very slowly.

    for (i = 0; i < num_desired_motions; ++i) {

        if (motions[i].num_inliers <= 0) {

            // Output model has already been initialized to the identity model,

            // so just skip setup

            continue;

        }

        bool bad_model = false;

        for (int refine_count = 0; refine_count < NUM_REFINES; refine_count++) {

            int num_inliers = motions[i].num_inliers;

            assert(num_inliers >= min_inliers);

            if (!model_info->find_transformation(

                    matched_points, motions[i].inlier_indices, num_inliers, params_this_motion)) {

                // In the unlikely event that this model fitting fails, we don't have a

                // good fallback. So leave this model set to the identity model

                bad_model = true;

                break;

            }

            // Score the newly generated model

            model_info->score_model(params_this_motion, matched_points, npoints, &current_motion);

            // At this point, there are three possibilities:

            // 1) If we found more inliers, keep refining.

            // 2) If we found the same number of inliers but a lower SSE, we want to

            //    keep the new model, but further refinement is unlikely to gain much.

            //    So commit to this new model

            // 3) It is possible, but very unlikely, that the new model will have

            //    fewer inliers. If it does happen, we probably just lost a few

            //    borderline inliers. So treat the same as case (2).

            if (current_motion.num_inliers > motions[i].num_inliers) {

                motions[i].num_inliers        = current_motion.num_inliers;

                motions[i].sse                = current_motion.sse;

                int* tmp                      = motions[i].inlier_indices;

                motions[i].inlier_indices     = current_motion.inlier_indices;

                current_motion.inlier_indices = tmp;

            } else {

                // Refined model is no better, so stop

                // This shouldn't be significantly worse than the previous model,

                // so it's fine to use the parameters in params_this_motion.

                // This saves us from having to cache the previous iteration's params.

                break;

            }

        }

        if (bad_model) {

            continue;

        }

        // Fill in output struct

        memcpy(motion_models[i].params, params_this_motion, MAX_PARAMDIM * sizeof(*motion_models[i].params));

        for (int j = 0; j < motions[i].num_inliers; j++) {

            int                   index         = motions[i].inlier_indices[j];

            const Correspondence* corr          = &matched_points[index];

            motion_models[i].inliers[2 * j + 0] = (int)rint(corr->x);

            motion_models[i].inliers[2 * j + 1] = (int)rint(corr->y);

        }

        motion_models[i].num_inliers = motions[i].num_inliers;

    }

finish_ransac:

    EB_FREE_ARRAY(inlier_buffer);

    EB_FREE_ARRAY(motions);

    return ret_val;

}

static const RansacModelInfo ransac_model_info[TRANS_TYPES] = {

    // IDENTITY

    {NULL, NULL, 0},

    // TRANSLATION

    {find_translation, score_translation, 1},

    // ROTZOOM

    {find_rotzoom, score_affine, 2},

    // AFFINE

    {find_affine, score_affine, 3},

};

// Returns true on success, false on error

bool svt_aom_ransac(const Correspondence* matched_points, int npoints, TransformationType type,

                    MotionModel* motion_models, int num_desired_motions, bool* mem_alloc_failed) {

    assert(type > IDENTITY && type < TRANS_TYPES);

    return ransac_internal(

        matched_points, npoints, motion_models, num_desired_motions, &ransac_model_info[type], mem_alloc_failed);

}