/*
 * Implements omp in c++, only accepts real arguments,
 * stops for a fixed tol, must have 3 inputs and five 
 * outputs.
 * MATLAB CALL
 *  [mNewDictionary,iOldDictionary,mOrthogonalDictionary,...
 *      mBiorthogonal,vCoefficients]=omp(vSignal,...
 *      mDictionary,tolerance1);
 */

#include <cmath>
#include "mex.h"
#include "blas.h"

using namespace std;

int choose_atom_oomp(double residule[], double mUnselectedAtoms[], int m, int n);
void vector_matrix_BLAS(double *vector, double *matrix, double *returnVector, int m, int n);
double real_inner_product(double *v1, double *v2, int szVector);
void reorthognalize(double *mOrthogonalDictionary,int szSignal,int k, int nRepetitions);
void orthogonalize(double *mOrthogonalDictionary, int szSignal, int k, int nAtoms);
void calc_biorthogonal(double *mBiorthogonal, double *vNewAtom, double *vOrthogonalAtom, double norm, int szSignal, int k);
void calc_residule(double *residule, double *vSignal, double *mOrthogonalDictionary, int szSignal);
void swap_elements(double *swapA, double *swapB, int nElements);
void copy_elements(double *array1, double *array2, int nElements);
double normalize(double *atom, int szSignal, double normAtom = 0);
int omp(double *vSignal,
    double *mNewDictionary,
    double *tolerance1,
    int szSignal,
    int nAtoms,
    double *iOldDictionary,
    double *mOrthogonalDictionary,
    double *mBiorthogonal,
    double *vCoefficients);


void mexFunction(int nlhs, mxArray *plhs[],
    int nrhs, const mxArray *prhs[])
{
    
    if (nrhs != 3)
    {
        mexErrMsgTxt("Wrong number of input arguments");
    }
    
    if (nlhs != 6)
    {
        mexErrMsgTxt("Wrong number of output arguments");
    }

    if ( mxIsComplex(prhs[0]) || mxIsClass(prhs[0],"sparse") || mxIsChar(prhs[0]) )
    {
        mexErrMsgTxt("Input must be real, full,and nonstring");
    }

    int szSignal, n0, m1, nAtoms, m2, n2;
    // Inputs
    double *vSignal, *mDictionary, *tolerance1;
    // Outputs
    double *mNewDictionary, *iOldDictionary, *mOrthogonalDictionary, *mBiorthogonal, *vCoefficients, *chosen;
  

    /* Find the dimensions of input matrices */
    szSignal = mxGetM(prhs[0]);
    n0 = mxGetN(prhs[0]);

    m1 = mxGetM(prhs[1]);
    nAtoms = mxGetN(prhs[1]);
    
    m2 = mxGetM(prhs[2]);
    n2 = mxGetN(prhs[2]);

    /* Validate input dimensions */
    if (szSignal == 1 || n0 != 1)
    {
        mexErrMsgTxt("The signal must be a column vector");
    }

    if (szSignal != m1)
    {
        mexErrMsgTxt("Matrix dimensions must agree");
    }

    if (m2*n2 != 1)
    {
        mexErrMsgTxt("The tolerance cannot be a matrix");
    }
    

    /* Create mxArray's for the output data */
    // Include a comment to say what each of these is
    plhs[0] = mxCreateDoubleMatrix(szSignal, nAtoms, mxREAL);
    plhs[1] = mxCreateDoubleMatrix(nAtoms, 1, mxREAL);
    plhs[2] = mxCreateDoubleMatrix(szSignal, nAtoms, mxREAL);
    plhs[3] = mxCreateDoubleMatrix(szSignal, szSignal, mxREAL);
    plhs[4] = mxCreateDoubleMatrix(szSignal, 1, mxREAL);
    plhs[5] = mxCreateDoubleMatrix(1, 1, mxREAL);
    
    /* Get a pointer to the input data */
    vSignal = mxGetPr(prhs[0]);
    mDictionary = mxGetPr(prhs[1]);
    tolerance1 = mxGetPr(prhs[2]);
    
    /* Create pointer's to the output data */
    mNewDictionary = mxGetPr(plhs[0]);
    iOldDictionary = mxGetPr(plhs[1]);
    mOrthogonalDictionary = mxGetPr(plhs[2]);
    mBiorthogonal = mxGetPr(plhs[3]);
    vCoefficients = mxGetPr(plhs[4]);
    chosen = mxGetPr(plhs[5]);
    
    /* Perform deep copy of the original dictionary matrix */
    for ( int i = 0; i < nAtoms; i++ )
    {
        for ( int j = 0; j < szSignal; j++ )
        {
            mNewDictionary[(i*szSignal) + j] = mDictionary[(i*szSignal) + j];
        }
    }
    
    
    *chosen = static_cast<double>(omp(vSignal,mNewDictionary,tolerance1,szSignal, nAtoms,
        iOldDictionary,mOrthogonalDictionary,mBiorthogonal,vCoefficients));
    
}
    
    
    
int omp(   
            double *vSignal,
            double *mNewDictionary,
            double *tolerance1,
            int szSignal,
            int nAtoms,
            double *iOldDictionary,
            double *mOrthogonalDictionary,
            double *mBiorthogonal,
            double *vCoefficients           
        )
{
    /* Declare constants */
    const double tolerance2 = 1e-10;
    const double tolerance3 = 1e-7;
    
    /* Declare variables */
    int k, i, iChosenAtom, nIterations, iNewAtom;
    double normKthAtom, normResidule;
    
    
    /* Declare dynamic arrays remember to delete after use */
    double *residule;
    residule = new double[szSignal];

    
    /* Perform deep copy of the dictionary into the orthogonal dictionay */
    copy_elements(mOrthogonalDictionary,mNewDictionary,szSignal*nAtoms);
    
    
    /* Perform deep copy of the original signal */
    copy_elements(residule,vSignal,szSignal);
    
        
    /* Populate the index array - remember MATLAB index's start at 1 not 0*/         
    for ( int i = 0; i < nAtoms; i++ )
    {
        iOldDictionary[i] = i+1;
    }    

    /* The number of iterations shoule be less the min of the nAtoms
     * and the szSignal. */
    if ( szSignal > nAtoms )
    {
        nIterations = nAtoms;
    } else {
        nIterations = szSignal;
    }

    /*************************************************************/
    /********************* Start of main routine *****************/
    /*************************************************************/
    

    for ( k = 0; k < nIterations; k++)
    {    
        iNewAtom = k*szSignal;
        
        /* Choose the next atom from the unselected atoms.
         */     
        iChosenAtom = choose_atom_oomp(vSignal,&mOrthogonalDictionary[iNewAtom],szSignal,(nAtoms - k));

        
        /* Stopping criterion (coefficient)
         * MATLAB CODE
         * if max_c<tol2
         *      fprintf('%s stopped, max(|<f,q>|/||q||)<= tol2=%g.\n',name,tol2);
         *      break;
         * end
         */
        if ( iChosenAtom == -1 )
        {
            break;            
        }

        /* Add k
         * because we want the index in the mNewDictionary array,
         * not the index in the smaller array we pass to choose_atom.
         */  
         iChosenAtom += k;      
        
        
        /* Populate the matrices with the new atoms
         * MATLAB CODE
         * if q~=k
         *      Q(:,[k q])=Q(:,[q k]); % swap k-th and q-th columns
         *      D(:,[k q])=D(:,[q k]);
         *      Di([k q])=Di([q k]);
         * end
         */
        if ( k != iChosenAtom )
        {
            swap_elements(&iOldDictionary[k],&iOldDictionary[iChosenAtom],1);
            swap_elements(&mNewDictionary[iNewAtom],&mNewDictionary[(iChosenAtom*szSignal)],szSignal);
            swap_elements(&mOrthogonalDictionary[iNewAtom],&mOrthogonalDictionary[iChosenAtom*szSignal],szSignal);
        }
        
        
        /* Re-orthogonalization of Q(:,k)  w.r.t Q(:,1),..., Q(:,k-1)
         * MATLAB CODE
         * if k>1
         *      for p=1:k-1
         *          Q(:,k)=Q(:,k)-(Q(:,p)'*Q(:,k))*Q(:,p);
         *      end	   
         *  end
         */     
        if ( k > 0 )
        {
            reorthognalize(mOrthogonalDictionary,szSignal,k,1);
        }
                 
        
        /* Normalize atom Q(:,k)
         * MATLAB CODE
         * nork=norm(Q(:,k));
         * Q(:,k)=Q(:,k)/nork; %normalization
         */      
        normKthAtom = normalize(&mOrthogonalDictionary[iNewAtom],szSignal);
         
        
        /* Compute biorthogonal functions beta*/
        calc_biorthogonal(mBiorthogonal, &mNewDictionary[iNewAtom], &mOrthogonalDictionary[iNewAtom], normKthAtom, szSignal, k);
        
        
        /* If last atom then don't need to (can't) orthogonalise
         * Q(:,k+1:n) wrt Q(:,k) so exit loop and return values
         * MATLAB CODE
         * if k==N
         *      break; 
         *  end 
         */
        if ( k == nIterations )
        {
            break;
        }
        
        

         /* Orthogonalization of Q(:,k+1:n) wrt Q(:,k) */
        orthogonalize(mOrthogonalDictionary, szSignal, k, nAtoms);
        
        
        /* Calculate the residule */
        calc_residule(residule, vSignal, &mOrthogonalDictionary[iNewAtom], szSignal);
    
        
        /* Calculate the norm of the residule */
        normResidule = sqrt(real_inner_product(residule,residule,szSignal));

        
        /* Stopping criteria (distance)
         * MATLAB CODE - We use the residule check with omp if this is the same
         * if (norm(f'-D(:,1:k)*(f*beta)')*sqrt(delta) < tol) && (tol~= 0)break;end;
         */
        if ( normResidule < *tolerance1 && *tolerance1 != 0 )
        {
            /* Break so will not increment k before exiting the loop */
            k += 1;
            break;                          
        }
    }
    
    
    /*Calculate the coefficients
     * MATLAB CODE
     * c=f*beta;
     */
    for ( i = 0; i < k; i++ )
    {
        vCoefficients[i] = real_inner_product(vSignal,&mBiorthogonal[(i*szSignal)],szSignal);
    }
    
    /* Free the memory */
    delete residule;
    residule = 0; // For safety
    
    return k;
}






double real_inner_product(double *v1, double *v2, int szVector)
{
  double sum = 0;
  for ( int i = 0; i < szVector; i++)
  {
    sum += v1[i]*v2[i];    
  }
  
  return sum;
}

int choose_atom_oomp(double *vSignal, double *mUnselectedAtoms, int m, int n)
{
    

    int maxIndex = 0;
    double temp, normAtom, maxValue = 0, tolerance2 = 1e-12, tolerance3 = 1e-7; // Should be global constant
    double *cc;
    cc = new double[n];    
    
    
    /* MATLAB CODE
     *  for p=k:N
     *      c(p)=norm(Q(:,p));
     *      if c(p)>tol1 
     *          cc(p)=abs(f*Q(:,p))/c(p);
     *      end
     *  end
     *  [max_c,q]=max(cc);
     */    
    vector_matrix_BLAS(vSignal, mUnselectedAtoms, cc, m, n); 
    for ( int i = 0; i < n; i++)
    {
        normAtom = sqrt(real_inner_product(&mUnselectedAtoms[i*m],&mUnselectedAtoms[i*m],m));
        if ( normAtom > tolerance3 )
        {
            temp = abs(cc[i])/normAtom;
            if ( temp > maxValue )
            {
                maxValue = temp;
                maxIndex = i;
            }
        }
    }
    

    /* MATLAB CODE
     * if max_c<tol2 
     *   fprintf('%s stopped, max(|<f,q>|/||q||)<= tol2=%g.\n',name,tol2);
     *   break;
     * end
     */    
    if (maxValue < tolerance2)
    {
        maxIndex = -1;
    }
    
    /* Free the memory */
    delete cc;
    cc = 0; // For safety

    return maxIndex;
  
}


void vector_matrix_BLAS(double *vector, double *matrix, double *returnVector, int m, int n)
{
    /* Special case as we are transposing the matrix */
    char *chn = "T";
    double zero = 0, one = 1;
    int incx = 1;
    dgemv(chn, &m, &n, &one, matrix, &m, vector, &incx, &zero, returnVector, &incx);
}


void reorthognalize(double *mOrthogonalDictionary,int szSignal,int k, int nRepetitions)
{
    double alpha;

     /* MATLAB CODE
      * for l=1:nRepetitions
      *     for p=1:k-1
      *              Q(:,k)=Q(:,k)-(Q(:,p)'*Q(:,k))*Q(:,p);
      *      end
      * end
     */
    for ( int l = 0; l < nRepetitions; l++ )
    {         
        for ( int i = 0; i < k ; i++)
        {
            alpha = real_inner_product(&mOrthogonalDictionary[(i*szSignal)],&mOrthogonalDictionary[(k*szSignal)],szSignal);
            for ( int j = 0; j < szSignal; j++)
            {
                mOrthogonalDictionary[(k*szSignal) + j] = mOrthogonalDictionary[(k*szSignal)+j] - alpha*mOrthogonalDictionary[(i*szSignal) + j];
            }
        }
    }
}


void orthogonalize(double *mOrthogonalDictionary, int szSignal, int k, int nAtoms)
{
    double alpha;
    /* orthogonalization of Q(:,p) wrt Q(:,k)
     * MATLAB CODE     
     *      for p=k+1:N
     *          Q(:,p)=Q(:,p)-(Q(:,k)'*Q(:,p))*Q(:,k);
     *      end
     */  
  
    for ( int i = k + 1; i < nAtoms ; i++)
    {
        alpha = real_inner_product(&mOrthogonalDictionary[(i*szSignal)],&mOrthogonalDictionary[(k*szSignal)],szSignal);
        for ( int j = 0; j < szSignal; j++)
        {
            mOrthogonalDictionary[(i*szSignal) + j] = mOrthogonalDictionary[(i*szSignal) + j] - alpha*mOrthogonalDictionary[(k*szSignal) + j];
        }
    }
    
}


void calc_biorthogonal(double *mBiorthogonal, double *vNewAtom, double *vOrthogonalAtom, double normAtom, int szSignal, int k)
{
    /* Compute biorthogonal functions beta from 1 to k-1
     * MATLAB CODE
     * beta=beta-Q(:,k)*(D(:,k)'*beta)/nork;
     */

    double alpha;
    /* If picking more than 150 atoms may be quicker to use BLAS and the matrix vector multiplication */
    if ( k > 0 )
    {
        for ( int j = 0; j < k ; j++ )
        {
            alpha = real_inner_product(vNewAtom,&mBiorthogonal[(j*szSignal)],szSignal)/normAtom;

            for ( int i = 0; i < szSignal; i++ )
            {		 	  
                mBiorthogonal[(j*szSignal)+i] = mBiorthogonal[(j*szSignal)+i] - alpha*vOrthogonalAtom[i];        
            }
        }
    }
    
    /* Calculate the kth biorthogonal function
     * MATLAB CODE
     * beta(:,k)=Q(:,k)/nork; % kth biorthogonal function
     */        
    
    for ( int i = 0; i< szSignal; i++ )
    {
        mBiorthogonal[(k*szSignal)+i] = vOrthogonalAtom[i]/normAtom;
    } 
            
    
}



void calc_residule(double *residule, double *vSignal, double *mOrthogonalDictionary, int szSignal)
{
    /* Calculate the residule
     * MATLAB CODE
     * Re=Re-f*Q(:,k)*Q(:,k)';
     */

    double alpha;
    
    alpha = real_inner_product(mOrthogonalDictionary,vSignal,szSignal);    
    for ( int i = 0; i < szSignal; i++)
    {
       residule[i] = residule[i] - alpha*mOrthogonalDictionary[i];
    }
    
}


void swap_elements(double *swapA, double *swapB, int nElements)
{
    double swappedElement;
    for ( int i = 0; i < nElements; i++ )
    {
        swappedElement = swapB[i];
        swapB[i] = swapA[i];
        swapA[i] = swappedElement;
    }
}

void copy_elements(double *array1, double *array2, int nElements)
{
    for ( int i = 0; i < nElements; i++ )
    {
        array1[i] = array2[i];
    }
}
        
    
double normalize(double *atom, int szSignal, double normAtom)
{   
    if ( normAtom == 0)
    {
        normAtom = sqrt(real_inner_product(atom,atom,szSignal));
    }
    
    for ( int i = 0; i < szSignal; i++)
    {
        atom[i] = atom[i]/normAtom;
    }
    return normAtom;
}