#include "matrix.h"
#include "mex.h"
#include <complex.h>
#include <math.h>

using namespace std;

#define EPS 1e-8

void mexQuadsolve(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[]) {

    enum {
        MAX_NLHS = 3,
    };

    if (nlhs < 2 || nlhs > 3) {
        mexErrMsgIdAndTxt("MATLAB:cubesolve:output_arguments",
                "number of output arguments must be 2 or 3");
    }

    size_t mrows = mxGetM(prhs[0]);
    size_t ncols = mxGetN(prhs[0]);
    for (int i = 0; i < nlhs; ++i) {
        plhs[i] = mxCreateDoubleMatrix((mwSize) mrows, (mwSize) ncols, mxCOMPLEX);
    }

    mxComplexDouble *res[MAX_NLHS];
    for (int i = 0; i < nlhs; ++i) {
        res[i] = (mxComplexDouble *) mxGetComplexDoubles(plhs[i]);
    }

    mxComplexDouble *A = (mxComplexDouble *) mxGetComplexDoubles(prhs[0]);
    mxComplexDouble *B = (mxComplexDouble *) mxGetComplexDoubles(prhs[1]);
    mxComplexDouble *C = (mxComplexDouble *) mxGetComplexDoubles(prhs[2]);
    __complex__ roots[] = {0, 0, 0};

    for (size_t i = 0; i < mrows * ncols; ++i) {
        if (A[i].real * A[i].real + A[i].imag * A[i].imag < EPS) {
            mexErrMsgIdAndTxt("MATLAB:cubesolve:null_coeff",
                    "A must be not null");
        }

        __complex__ p = (B[i].real + I * B[i].imag) / (A[i].real + I * A[i].imag);
        __complex__ q = (C[i].real + I * C[i].imag) / (A[i].real + I * A[i].imag);
        __complex__ w = csqrt(cpow(p / 3., 3.0) + cpow(q / 2., 2.0));
        __complex__ alpha = cpow(-q / 2. + w, 1.0 / 3.0);
        // and here we need to rotate our beta becauose cpow calculates wrong branch
        // of logarithm inside computing the complex power
        __complex__ beta = cpow(-q / 2. - w, 1.0 / 3.0) * (-0.5 - I * csqrt(3) / 2);

        roots[0] = -0.5 * (alpha + beta) + I * sqrt(3) / 2.0 * (alpha - beta);
        roots[1] = -0.5 * (alpha + beta) - I * sqrt(3) / 2.0 * (alpha - beta);

        if (nlhs == 3) {
            roots[2] = alpha + beta;
        }

        for (int k = 0; k < nlhs; ++k) {
            res[k][i].real = creal(roots[k]);
            res[k][i].imag = cimag(roots[k]);
        }
    }
}