#include "mex.h"
#include "matrix.h"
#include "math.h"
#include "stdlib.h"

/*  Copyright (C) 2009 Cesare Magri                                        */

/*
 * -------
 * LICENSE
 * -------
 * This software is distributed free under the condition that:
 *
 * 1. it shall not be incorporated in software that is subsequently sold;
 *
 * 2. the authorship of the software shall be acknowledged and the following
 *    article shall be properly cited in any publication that uses results
 *    generated by the software:

 *      Magri C, Whittingstall K, Singh V, Logothetis NK, Panzeri S: A
 *      toolbox for the fast information analysis of multiple-site LFP, EEG
 *      and spike train recordings. BMC Neuroscience 2009 10(1):81;
 * 3.  this notice shall remain in place in each source file.
 */

void mexFunction(int nlhs, mxArray *plhs[],
                 int nrhs, const mxArray *prhs[])
{
    /* Input:                                                            */
    double *R;
    double *nt;
    double Npart;
    double partIndx;
    
    /* Dimensions of R:                                                  */
    mwSize Nc;
    mwSize maxNt;
    mwSize Ns;
    
    /* Indexes:                                                          */
    mwIndex c, t, s;
    
    mwIndex tmin;
    mwIndex tmax;

    mwSize maxNtOut;
    mwSize *outDims;
    
    /* Output:                                                           */
    double *Rout;
    double *totNtOut;
    
    /* Can be either an input or output (computed if not provided):      */
    double *ntOut;

    
    /* ----------------------------------------------------------------- */

    
    /* Read input:                                                       */
    R        = mxGetPr(prhs[0]);
    nt       = mxGetPr(prhs[1]);
    Npart    = *mxGetPr(prhs[2]);
    partIndx = *mxGetPr(prhs[3]);
    if(nrhs==5)
        ntOut = mxGetPr(prhs[4]);
    
    /* Get dimensions of R:                                              */
    Nc    = mxGetDimensions(prhs[0])[0];
    maxNt = mxGetDimensions(prhs[0])[1];
    Ns    = mxGetDimensions(prhs[0])[2];
    
    /* Set dimensions of output:                                         */
    maxNtOut = maxNt / Npart; /* floored int ratio                       */
    
    outDims = mxCalloc(3, sizeof(mwSize));
    outDims[0] = Nc;
    outDims[1] = maxNtOut;
    outDims[2] = Ns;

    /* Allocate output                                                   */
    plhs[0] = mxCreateNumericArray(3, outDims, mxDOUBLE_CLASS, 0);
    Rout = mxGetPr(plhs[0]);
    
    /*
     * If ntOut is not provided as an input we need to build it together
     * with the total number of output trials:
     */
    if(nrhs<5) {
        plhs[1] = mxCreateDoubleMatrix(Ns, 1, mxREAL);
        ntOut = mxGetPr(plhs[1]);
        
        plhs[2] = mxCreateDoubleScalar(0);
        totNtOut = mxGetPr(plhs[2]);
    }

    /* Copy data into partition:                                         */
    for(s=0; s<Ns; s++)
    {
        /* if ntOut is not provided as an input we compute it:           */
        if(nrhs<5)
        {
            ntOut[s] = floor(nt[s] / Npart);
            *totNtOut += ntOut[s];
        }

        tmax = ntOut[s] * partIndx;
        tmin = tmax - ntOut[s];

        for(t=tmin; t<tmax; t++)
        {    
            for(c=0; c<Nc; c++)
            {
                Rout[c + (t-tmin)*Nc + s*Nc*maxNtOut] = R[c + t*Nc + s*Nc*maxNt];
            }
        }
    }
    
    mxFree(outDims);
}