#include "TransSolver.h"
#include "CVUtils.h"

#define M_LOW_TOLERANCE 0.000001

Matx33d affineTrans(const vector<Point2d>& src, const vector<Point2d>& dst)
{
    if (dst.empty() || src.empty()) {
        return Mat();
    }

    Point2d pc, qc;
    int smallerSize = src.size() < dst.size() ? src.size() : dst.size();

    for (int i = 0; i < smallerSize; i++) {
        pc += src[i];
        qc += dst[i];
    }
    pc.x /= smallerSize;
    pc.y /= smallerSize;
    qc.x /= smallerSize;
    qc.y /= smallerSize;

    Matx21d pit;
    Matx12d pi;
    Matx12d qi;
    Matx22d spitpi = Matx22d::zeros();
    Matx22d pitpi;
    Matx22d pitqi;
    Matx22d spitqi = Matx22d::zeros();

    for (int i = 0; i < src.size() && i < dst.size(); i++) {
        Point2d qpi = src[i] - pc;
        Point2d qqi = dst[i] - qc;
        
        pit(0) = qpi.x;
        pit(1) = qpi.y;
        pi(0) = qpi.x;
        pi(1) = qpi.y;
        qi(0) = qqi.x;
        qi(1) = qqi.y;
        pitpi = pit*pi;
        spitpi = pitpi + spitpi;
        pitqi = pit*qi;
        spitqi = pitqi + spitqi;
    }
    Matx22d ispitpi; 
    ispitpi = spitpi.inv();

    Matx22d M = ispitpi*spitqi;

    double m11 = M(0, 0);
    double m21 = M(0, 1);
    double m12 = M(1, 0);
    double m22 = M(1, 1);
    
    Matx33d qm(m11, m12, 0, m21, m22, 0, 0, 0, 1);
    Matx33d pcm(1.0, 0, -pc.x, 0, 1.0, -pc.y, 0, 0, 1);
    Matx33d qcm(1.0, 0, qc.x, 0, 1.0, qc.y, 0, 0, 1);

    Matx33d ret = qcm*qm*pcm;
    
    return ret;
}

Matx33d rigidTrans(const vector<Point2d>& src, const vector<Point2d>& dst)
{
    if (dst.empty() || src.empty()) {
        return Mat();
    }

    vector<double> weights(src.size(), 1.0 / src.size());
    Point2d pc, qc;
    double wsum = 0;

    for (int i = 0; i < src.size(); i++) {
        double w = 1.0 / src.size();
        weights[i] = w;
        pc += src[i] * w;
        qc += dst[i] * w;
        wsum += w;
    }
    pc.x /= wsum;
    pc.y /= wsum;
    qc.x /= wsum;
    qc.y /= wsum;

    double u = 0;
    double u1, u2;
    u1 = 0;
    u2 = 0;
    for (int i = 0; i < src.size() && i < dst.size(); i++) {
        Point2d qpi = src[i] - pc;
        Point2d qqi = dst[i] - qc;
        Point2d pi(qpi.x, qpi.y);
        Point2d qi(qqi.x, qqi.y);
        u1 += pi.dot(qi)*weights[i];
        Point2d pi_(pi.y, -pi.x);
        u2 += qi.dot(pi_)*weights[i];
    }
    u = sqrt(u1*u1 + u2*u2);
    if (u < M_LOW_TOLERANCE) {
        u = M_LOW_TOLERANCE;
    }

    Matx22d R = Matx22d::zeros();
    Matx22d r = Matx22d::zeros();

    for (int i = 0; i < src.size() && i < dst.size(); i++) {
        Point2d qpi = src[i] - pc;
        Point2d qqi = dst[i] - qc;
        Point2d pi(qpi.x, qpi.y);
        Point2d qi(qqi.x, qqi.y);
        Point2d pi_(pi.y, -pi.x);
        Point2d qi_(qi.y, -qi.x);

        r(0, 0) = pi.dot(qi);
        r(0, 1) = pi.dot(qi_);
        r(1, 0) = pi_.dot(qi);
        r(1, 1) = pi_.dot(qi_);

        R = R + r * (weights[i] / u);
    }

    double m11 = R(0, 0);
    double m21 = R(0, 1);
    double m12 = R(1, 0);
    double m22 = R(1, 1);
    Matx33d qm(m11, m12, 0, m21, m22, 0, 0, 0, 1);
    Matx33d pcm(1.0, 0, -pc.x, 0, 1.0, -pc.y, 0, 0, 1);
    Matx33d qcm(1.0, 0, qc.x, 0, 1.0, qc.y, 0, 0, 1);
    Matx33d ret = qcm*qm*pcm;

    return ret;
}

bool cmpPointVec(const vector<Point2d>& vec0, const vector<Point2d>& vec1, const Matx33d& mat, double tor)
{
    int smallerSize = vec0.size() < vec1.size() ? vec0.size() : vec1.size();
    for (int i = 0; i < smallerSize; ++i)
    {
        Point2d p0, p1;
        p0 = vec0[i];
        p1 = vec1[i];
        Point3d tp0 = mat*p0;
        tp0.x /= tp0.z;
        tp0.y /= tp0.z;
        if (abs(p1.x - tp0.x) > tor || abs(p1.y - tp0.y) > tor)
        {
            return false;
        }
    }
    return true;
}

void testTransSolver()
{
    {
        //rotation only

        vector<Point2d> vec0, vec1;
        vec0.push_back(Point2d());
        vec0.push_back(Point2d(1, 0));
        vec0.push_back(Point2d(1, 1));
        vec0.push_back(Point2d(0, 1));

        vec1.push_back(Point2d());
        vec1.push_back(Point2d(0, 1));
        vec1.push_back(Point2d(-1, 1));
        vec1.push_back(Point2d(-1, 0));

        Matx33d ret = affineTrans(vec0, vec1);   
        assert(cmpPointVec(vec0, vec1, ret, M_LOW_TOLERANCE));

        ret = rigidTrans(vec0, vec1);
        assert(cmpPointVec(vec0, vec1, ret, M_LOW_TOLERANCE));
    }
    {
        // rotation and scale

        vector<Point2d> vec0, vec1;
        vec0.push_back(Point2d());
        vec0.push_back(Point2d(1, 0));
        vec0.push_back(Point2d(1, 1));
        vec0.push_back(Point2d(0, 1));

        vec1.push_back(Point2d());
        vec1.push_back(Point2d(1, 1));
        vec1.push_back(Point2d(0, 2));
        vec1.push_back(Point2d(-1, 1));

        Matx33d ret = affineTrans(vec0, vec1);
        assert(cmpPointVec(vec0, vec1, ret, M_LOW_TOLERANCE));
    }
    {
        // scale only

        vector<Point2d> vec0, vec1;
        vec0.push_back(Point2d());
        vec0.push_back(Point2d(1, 0));
        vec0.push_back(Point2d(1, 1));
        vec0.push_back(Point2d(0, 1));

        vec1.push_back(Point2d());
        vec1.push_back(Point2d(2, 0));
        vec1.push_back(Point2d(2, 2));
        vec1.push_back(Point2d(0, 2));

        Matx33d ret = affineTrans(vec0, vec1);
        assert(cmpPointVec(vec0, vec1, ret, M_LOW_TOLERANCE));
    }
    {
        // translation only

        vector<Point2d> vec0, vec1;
        vec0.push_back(Point2d());
        vec0.push_back(Point2d(1, 0));
        vec0.push_back(Point2d(1, 1));
        vec0.push_back(Point2d(0, 1));

        vec1.push_back(Point2d(1, 1));
        vec1.push_back(Point2d(2, 1));
        vec1.push_back(Point2d(2, 2));
        vec1.push_back(Point2d(1, 2));

        Matx33d ret = affineTrans(vec0, vec1);
        assert(cmpPointVec(vec0, vec1, ret, M_LOW_TOLERANCE));

        ret = rigidTrans(vec0, vec1);
        assert(cmpPointVec(vec0, vec1, ret, M_LOW_TOLERANCE));
    }
    {
        for (int i = 0; i < 10000; ++i)
        {
            // random rotation and translation
            Mat mat23 = getRotationMatrix2D(Point2f(), rand() % 360, 1.0);
            mat23.at<double>(0, 2) = (rand() % 1000) / 1000.0;
            mat23.at<double>(1, 2) = (rand() % 1000) / 1000.0;
            Matx33d matx = Matx33d::eye();
            Mat mat(3, 3, CV_64FC1, matx.val);
            mat23.copyTo(mat.rowRange(0, 2));

            vector<Point2d> vec0, vec1;
            vec0.push_back(Point2d());
            vec0.push_back(Point2d(1, 0));
            vec0.push_back(Point2d(1, 1));
            vec0.push_back(Point2d(0, 1));

            vec1 = vec0;
            transPoints(vec1, matx);

            Matx33d ret = affineTrans(vec0, vec1);
            assert(cmpPointVec(vec0, vec1, ret, M_LOW_TOLERANCE));
            ret = rigidTrans(vec0, vec1);
            assert(cmpPointVec(vec0, vec1, ret, M_LOW_TOLERANCE));
        }
    }
    {
        // skew only

        vector<Point2d> vec0, vec1;
        vec0.push_back(Point2d());
        vec0.push_back(Point2d(1, 0));
        vec0.push_back(Point2d(1, 1));
        vec0.push_back(Point2d(0, 1));

        vec1.push_back(Point2d());
        vec1.push_back(Point2d(1, 0));
        vec1.push_back(Point2d(2, 1));
        vec1.push_back(Point2d(1, 1));

        Matx33d ret = affineTrans(vec0, vec1);
        assert(cmpPointVec(vec0, vec1, ret, M_LOW_TOLERANCE));
    }
    {
        // random affine

        for (int i = 0; i < 10000; ++i)
        {
            vector<Point2d> vec0, vec1;
            vec0.push_back(Point2d());
            vec0.push_back(Point2d(1, 0));
            vec0.push_back(Point2d(1, 1));
            vec0.push_back(Point2d(0, 1));

            Matx33d trans = Matx33d::eye();
            Mat mat(3, 3, CV_64FC1, trans.val);
            randu(mat.rowRange(0, 2), 0, 1.0);

            vec1 = vec0;
            transPoints(vec1, trans);

            Matx33d ret = affineTrans(vec0, vec1);
            assert(cmpPointVec(vec0, vec1, ret, M_LOW_TOLERANCE));
        }        
    }
}