/***************************** LICENSE START ***********************************

 Copyright 2012 ECMWF and INPE. This software is distributed under the terms
 of the Apache License version 2.0. In applying this license, ECMWF does not
 waive the privileges and immunities granted to it by virtue of its status as
 an Intergovernmental Organization or submit itself to any jurisdiction.

 ***************************** LICENSE END *************************************/

#include "macro.h"

static double Det(CMatrix& m)
{
    // Assumes that m is square
    double d = 0, p1, p2, p3;
    int n = m.Col();

    for (int i = 0; i < n; i++) {
        p1 = p2 = 1.0;
        p3 = m(0, i);

        int k = i;

        for (int j = 1; j < n; j++) {
            k = (k + 1) % n;
            p1 *= m(j, k);
            p2 *= m(n - j, k);
        }
        p3 *= (p1 - p2);
        d += p3;
    }
    return d;
}

static int Solve(CMatrix& m)
{
    int i, j, k;
    int rows = m.Row();
    int cols = m.Col();

    for (i = 0; i < rows; i++) {
        for (j = i; j < rows && m(j, j) == 0.0; j++)
            ;

        if (j == rows)
            return 1;  // Cannot solve

        if (j != i) {
            for (k = 0; k < cols; k++) {
                double d = m(i, k);
                m(i, k) = m(j, k);
                m(j, k) = d;
            }
        }
        for (j = cols - 1; j >= 0; j--)
            m(i, j) /= m(i, i);

        for (j = i + 1; j < rows; j++) {
            for (k = cols - 1; k >= i; k--)
                m(j, k) -= m(j, i) * m(i, k);
        }
    }
    for (i = rows - 2; i >= 0; i--) {
        for (j = cols - 2; j > i; j--) {
            m(i, cols - 1) -= m(i, j) * m(j, cols - 1);
            m(i, j) = 0;
        }
    }
    return 0;
}

static int Invert(CMatrix& sm, CMatrix& dm)
{
    int i, j, k;
    int n = sm.Row();

    if (Det(sm) == 0.0)
        return 1;

    CMatrix d(n, n + 1);

    for (i = 0; i < n; i++) {
        int nrow = i - 1;
        int ncol = i - 1;
        for (j = 0; j < n; j++) {
            nrow = (nrow + 1) % n;

            d(j, n) = (j == 0.0);

            for (k = 0; k < n; k++) {
                ncol = (ncol + 1) % n;
                d(j, k) = sm(nrow, ncol);
            }
        }

        if (Solve(d))
            return 1;
        else {
            nrow = i - 1;
            for (j = 0; j < n; j++) {
                nrow = (nrow + 1) % n;
                dm(nrow, i) = d(j, n);
            }
        }
    }
    return 0;
}

static void Multiply(CMatrix& a, CMatrix& b, CMatrix& c)
{
    for (int i = 0; i < a.Row(); i++)
        for (int j = 0; j < b.Col(); j++) {
            double d = 0;
            for (int k = 0; k < a.Col(); k++)
                d += a(i, k) * b(k, j);
            c(i, j) = d;
        }
}


CMatrix::CMatrix(int r, int c) :
    Content(tmatrix)
{
    row = r > 0 ? r : 1;
    col = c > 0 ? c : 1;

    int size = row * col;

    values = new double[size];
    for (int i = 0; i < size; i++)
        values[i] = 0.0;
}

void CMatrix::Copy(const CMatrix& v)
{
    row = v.row;
    col = v.col;

    int size = col * row;

    values = new double[size];

    for (int i = 0; i < size; i++)
        values[i] = v.values[i];
}

void CMatrix::Dump(int)
{
    std::cout << "matrix(" << row << ',' << col << ')';
}

void CMatrix::Print(void)
{
    for (int i = 0; i < row; i++) {
        std::cout << '|';
        for (int j = 0; j < col; j++) {
            std::cout << values[i + row * j];
            if (j != col - 1)
                std::cout << ',';
        }
        std::cout << '|' << '\n';
    }
}
int CMatrix::Write(FILE* file)
{
    for (int i = 0; i < row; i++) {
        for (int j = 0; j < col; j++) {
            fprintf(file, "%g", values[i + row * j]);
            if (j != col - 1)
                fprintf(file, ",");
        }
        fprintf(file, "\n");
    }
    return ferror(file);
}

void CMatrix::SetSubValue(Value& v, int arity, Value* arg)
{
    if (!Check(1, v, arity, arg, tnumber, 2, tnumber, tnumber))
        return;

    int a, b;
    arg[0].GetValue(a);
    arg[1].GetValue(b);

    double d;
    v.GetValue(d);

    CMatrix& m = *this;

    if (a < 1 || a > m.Row() || b < 1 || b > m.Col()) {
        Error(
            "Matrix index [%d,%d] is out of range."
            " Matrix is %d x %d.",
            a, b, m.Row(), m.Col());
        return;
    }

    m(a - 1, b - 1) = d;
}

//==============================================================================

class Matrix1Function : public Function
{
public:
    Matrix1Function(const char* n) :
        Function(n) {}
    virtual Value Execute(int arity, Value* arg);
};

Value Matrix1Function::Execute(int arity, Value* arg)
{
    CMatrix* m;
    int n = 1;

    int i;
    for (i = 0; i < arity; i++) {
        CVector* v;
        arg[i].GetValue(v);
        if (v->Count() > n)
            n = v->Count();
    }

    m = new CMatrix(arity, n);

    for (i = 0; i < arity; i++) {
        CVector* v;
        arg[i].GetValue(v);

        for (int j = 0; j < v->Count(); j++)
            (*m)(i, j) = v->getIndexedValue(j);
    }

    return Value(m);
}

//==============================================================================

class Matrix2Function : public Function
{
public:
    Matrix2Function(const char* n) :
        Function(n, 2, tnumber, tnumber)
    {
        info = "Builds a new matrix given ist dimensions";
    }
    virtual Value Execute(int arity, Value* arg);
};

Value Matrix2Function::Execute(int, Value* arg)
{
    int a, b;
    arg[0].GetValue(a);
    arg[1].GetValue(b);
    return Value(new CMatrix(a, b));
}

//==============================================================================

class MatrixMulFunction : public Function
{
public:
    MatrixMulFunction(const char* n) :
        Function(n, 2, tmatrix, tmatrix)
    {
        info = "Matrix multiplication";
    }
    virtual Value Execute(int arity, Value* arg);
};

Value MatrixMulFunction::Execute(int, Value* arg)
{
    CMatrix *a, *b;
    arg[0].GetValue(a);
    arg[1].GetValue(b);

    if (a->Col() != b->Row())
        return Error("Cannot multiply a %d x %d matrix by a %d x %d matrix",
                     a->Row(), a->Col(), b->Row(), b->Col());

    auto* c = new CMatrix(a->Row(), b->Col());

    Multiply(*a, *b, *c);

    return Value(c);
}

//==============================================================================

class MatrixAddFunction : public Function
{
public:
    MatrixAddFunction(const char* n) :
        Function(n, 2, tmatrix, tmatrix)
    {
        info = "Matrix addition";
    }
    virtual Value Execute(int arity, Value* arg);
};

Value MatrixAddFunction::Execute(int, Value* arg)
{
    CMatrix *a, *b;
    arg[0].GetValue(a);
    arg[1].GetValue(b);

    if (a->Col() != b->Col() || a->Row() != b->Row())
        return Error("Cannot add a %d x %d matrix by a %d x %d matrix",
                     a->Row(), a->Col(), b->Row(), b->Col());

    auto* c = new CMatrix(a->Row(), a->Col());

    for (int i = 0; i < a->Row(); i++)
        for (int j = 0; j < a->Col(); j++)
            (*c)(i, j) = (*a)(i, j) + (*b)(i, j);

    return Value(c);
}

//==============================================================================

class MatrixSubFunction : public Function
{
public:
    MatrixSubFunction(const char* n) :
        Function(n, 2, tmatrix, tmatrix)
    {
        info = "Matrix addition";
    }
    virtual Value Execute(int arity, Value* arg);
};

Value MatrixSubFunction::Execute(int, Value* arg)
{
    CMatrix *a, *b;
    arg[0].GetValue(a);
    arg[1].GetValue(b);

    if (a->Col() != b->Col() || a->Row() != b->Row())
        return Error("Cannot substract a %d x %d matrix by a %d x %d matrix",
                     a->Row(), a->Col(), b->Row(), b->Col());

    auto* c = new CMatrix(a->Row(), a->Col());

    for (int i = 0; i < a->Row(); i++)
        for (int j = 0; j < a->Col(); j++)
            (*c)(i, j) = (*a)(i, j) - (*b)(i, j);

    return Value(c);
}

//==============================================================================

class MatrixNegFunction : public Function
{
public:
    MatrixNegFunction(const char* n) :
        Function(n, 1, tmatrix)
    {
        info = "Matrix negation";
    }
    virtual Value Execute(int arity, Value* arg);
};

Value MatrixNegFunction::Execute(int, Value* arg)
{
    CMatrix* a;
    arg[0].GetValue(a);

    auto* c = new CMatrix(a->Row(), a->Col());

    for (int i = 0; i < a->Row(); i++)
        for (int j = 0; j < a->Col(); j++)
            (*c)(i, j) = -(*a)(i, j);

    return Value(c);
}

//==============================================================================

class MatrixInvFunction : public Function
{
public:
    MatrixInvFunction(const char* n) :
        Function(n, 1, tmatrix)
    {
        info = "Matrix invertion";
    }
    virtual Value Execute(int arity, Value* arg);
};

Value MatrixInvFunction::Execute(int, Value* arg)
{
    CMatrix* a;
    arg[0].GetValue(a);


    if (a->Col() != a->Row())
        return Error("Cannot inverse %d x %d matrix. Matrix must by square",
                     a->Row(), a->Col());

    auto* c = new CMatrix(a->Row(), a->Col());

    if (Invert(*a, *c)) {
        delete c;
        return Error("Matrix is singular. It cannot be inverted");
    }

    return Value(c);
}

//==============================================================================

class MatrixDetFunction : public Function
{
public:
    MatrixDetFunction(const char* n) :
        Function(n, 1, tmatrix)
    {
        info = "Matrix determinant";
    }
    virtual Value Execute(int arity, Value* arg);
};

Value MatrixDetFunction::Execute(int, Value* arg)
{
    CMatrix* a;
    arg[0].GetValue(a);


    if (a->Col() != a->Row())
        return Error("Cannot get determinat of a %d x %d matrix. Matrix must by square",
                     a->Row(), a->Col());

    return Value(Det(*a));
}

//==============================================================================

class MatrixTransFunction : public Function
{
public:
    MatrixTransFunction(const char* n) :
        Function(n, 1, tmatrix)
    {
        info = "Transpose a matrix";
    }
    virtual Value Execute(int arity, Value* arg);
};

Value MatrixTransFunction::Execute(int, Value* arg)
{
    CMatrix* a;
    arg[0].GetValue(a);

    auto* b = new CMatrix(a->Col(), a->Row());

    for (int i = 0; i < a->Row(); i++)
        for (int j = 0; j < a->Col(); j++)
            (*b)(j, i) = (*a)(i, j);

    return Value(b);
}


//==============================================================================

class MatrixGetElemFunction : public Function
{
public:
    MatrixGetElemFunction(const char* n) :
        Function(n, 3, tmatrix, tnumber, tnumber) {}
    virtual Value Execute(int arity, Value* arg);
};


Value MatrixGetElemFunction::Execute(int, Value* arg)
{
    CMatrix* m;
    int a, b;

    arg[0].GetValue(m);
    arg[1].GetValue(a);
    arg[2].GetValue(b);

    if (a < 1 || a > m->Row() || b < 1 || b > m->Col())
        return Error(
            "Matrix index [%d,%d] is out of range."
            " Matrix is %d x %d.",
            a, b, m->Row(), m->Col());

    return Value((*m)(a - 1, b - 1));
}

static void install(Context* c)
{
    c->AddFunction(new Matrix1Function("_matrix"));
    c->AddFunction(new Matrix2Function("matrix"));
    c->AddFunction(new MatrixGetElemFunction("[]"));

    c->AddFunction(new MatrixMulFunction("*"));
    // c->AddFunction(new MatrixProdFunction("*"));
    c->AddFunction(new MatrixAddFunction("+"));
    c->AddFunction(new MatrixSubFunction("-"));
    c->AddFunction(new MatrixNegFunction("-"));
    c->AddFunction(new MatrixInvFunction("inverse"));
    c->AddFunction(new MatrixDetFunction("det"));
    c->AddFunction(new MatrixTransFunction("transpose"));
}

static Linkage linkage(install);
