#include <math.h>
#include "mex.h"
void mexFunction( int nlhs, mxArray *plhs[],
          int nrhs, const mxArray *prhs[] )

/* function [z,c] = mandelbrot_step(z,c,z0,k);
 * Take one step of the Mandelbrot iteration.
 * Complex arithmetic:
 *    z = z.^2 + z0
 *    c(abs(z) < 2) == k
 * Real arithmetic:
 *    x <-> real(z);
 *    y <-> imag(z);
 *    u <-> real(z0);
 *    v <-> imag(z0);
 *    [x,y] = [x.^2-y.^2+u, 2*x.*y+v];
 *    c(x.^2+y.^2 < 4) = k;
 *
 *    prhs[0] = z       The input matrix of complex numbers
 *    prhs[1] = c       The input matrix of iterations / count
 *    prhs[2] = z0      The input complex number
 *    prhs[3] = k       The iteration count
 *
 *    plhs[0] = z       The resultant matrix of complex numbers
 *    plhs[1] = c       The resultant matrix of iterations / count
 */

{
    double *x,*y,*u,*v,t;
    unsigned short *c,k;
    int j,n;

    x = mxGetPr(prhs[0]);   // Get array of real components of 'z'
    y = mxGetPi(prhs[0]);   // Get array of imaginary components of 'z'
    c = (unsigned short *) mxGetData(prhs[1]); // Get 'c' array from input
    u = mxGetPr(prhs[2]);   // Get array of real components of 'z0'
    v = mxGetPi(prhs[2]);   // Get array of imaginary components of 'z0'
    k = (unsigned short) mxGetScalar(prhs[3]);  // Get 'k' interate
    /* In the following, the output parameter is a pointer to the input
       data which is about to be modified. A clever trick! */
    plhs[0] = prhs[0];      // &(z') = &(z)
    plhs[1] = prhs[1];      // &(c') = &(c) Output points to 'c' which is about to be modified

    n = mxGetN(prhs[0]);    // Number of cols. in 'z' (but a vector in C)
    for (j=0; j<n*n; j++) { // For every element in 'z'...
        if (c[j] == k-1) {  // If element made it to k last time around...
            t = x[j];       // Save the real component
            /* Apply real arithmetic to X and Y components of Z[] */
            x[j] = x[j]*x[j] - y[j]*y[j] + u[j];    // x^2 - y^2 + u
            y[j] = 2*t*y[j] + v[j];                 // 2*x(orig)*y + v
            if (x[j]*x[j] + y[j]*y[j] < 4) {    // |If new value|^2 < 4...
                c[j] = k;                       // the element made it to k
            }
        }
    }
    return;
}