#include "mex.h"
#include <stdio.h>
#include <stdlib.h>
#include "cuda.h"
#include "cublas.h"
#include "../include/cuknns.h"
#include "../include/shared_functions.h"

#define DIMENTIONS 128

char* errMsg = "Only 128 dimentional data curentrly suported\n";

void mergeRes(knntype *data, knntype *idx, int k, int Q, int numStreams){

  for(int s=0; s<Q; s++){

    knntype cmax = -FLT_MAX;
    int maxid = 0;
    for(int i=0; i<k; i++){
      knntype tmp = data[s*k+i];
      if(tmp>cmax){
        cmax = tmp;
        maxid = i;
      }
    }

    for(int i=1; i<numStreams; i++){
      for(int j=0; j<k; j++){
        knntype tmp = data[s*k + i*Q*k + j];
        if(tmp<cmax){
          data[s*k + maxid] = tmp;
          idx[s*k + maxid] = idx[s*k + i*Q*k + j];
          cmax = -FLT_MAX;
          for(int p=0; p<k; p++){
            knntype tmp2 = data[s*k + p];
            if(tmp2>cmax){
              cmax = tmp2;
              maxid = p;
            }
          }

        }
      }
    }

  }


}


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

  /*=========MATLAB pointers=========*/
  knntype *data;//, *dp;
  knntype *query;
  knntimes TimesOut;
  size_t memory_free, memory_total;
  double *knnTime, *dstTime, *srchTime;

  TimesOut.dst_time = 0;
  TimesOut.srch_time = 0;
  TimesOut.knn_time = 0;

  /*========= Devece pointers ======*/
  knntype *d_data, *d_query, *d_labels, *d_dist, *d_dotp;

  float elapsedTime;
  cudaEvent_t start, stop;
  cudaEventCreate(&start);
  cudaEventCreate(&stop);

  cublasHandle_t handle;

  cublasCreate(&handle);

  query = (knntype*)mxGetData(prhs[0]);
  int numQueries = mxGetN(prhs[0]);
  data = (knntype*)mxGetData(prhs[1]);
  int objects = mxGetN(prhs[1]);
  //int objects = (int)mxGetScalar(prhs[2]);
  int attributes = mxGetM(prhs[0]);
  int k = (int)(mxGetScalar(prhs[2]));
  int numStreams = (int)mxGetScalar(prhs[3]);

  int qk = (int)(log((float)k) / log(2.0));

  //mexPrintf("numStreams = %d\n", numStreams);

  plhs[2] = mxCreateDoubleMatrix(1, 1, mxREAL);
  plhs[3] = mxCreateDoubleMatrix(1, 1, mxREAL);
  plhs[4] = mxCreateDoubleMatrix(1, 1, mxREAL);
  knnTime = mxGetPr(plhs[2]);
  dstTime = mxGetPr(plhs[3]);
  srchTime = mxGetPr(plhs[4]);

  knntype *d_dotB;

  cudaMalloc((void**)&d_query, numQueries*attributes*sizeof(knntype));
  cudaMalloc((void**)&d_dotB, numQueries*sizeof(knntype));


  /* find maximum corpus block that can be processed */
  cuMemGetInfo(&memory_free, &memory_total);
  int memory_req = (4*numQueries + attributes + 1)*sizeof(knntype);

  int maxObjects = (int)ceil((float)memory_free*0.9 / (float)memory_req);
  //printf("\n\n maxObjects: %d, Log2 max Corpus: %d\n\n", maxObjects, (int)floor(log((double)maxObjects)/log(2)));

  maxObjects = (1 << (int)floor(log((double)maxObjects)/log(2)));

  int reqStreams = 1;

  int CorpusBlocks = (int)ceil((float)objects/(float)maxObjects);

  numStreams = max(numStreams, reqStreams);

  //mexPrintf("General inter:\n");
  //mexPrintf("Free memory: %d\n", memory_free);
  //mexPrintf("Objects : %d\n", objects);
  //mexPrintf("max Corpus size Size: %d\n", maxObjects);
  //mexPrintf("Blocks of the Corpus: %d\n", CorpusBlocks);
  //mexPrintf("numStreams = %d\n", numStreams);

  maxObjects = min(maxObjects, objects);

  //mexPrintf("New maxObjets: %d\n", maxObjects);
  //mexPrintf("Mex function general\n");

  /* initialize distance functions */
  distFunctParam dstfunc;
  dstfunc.distF = pdist_NT;
  switch(attributes){
  case 128:
    dstfunc.dotP.dotQ = &dot4;
    dstfunc.dotP.nthreadsQ = 128;
    dstfunc.dotP.dotCrp = &crpdot2_128;
    dstfunc.dotP.nthreadsCrp = 128;
    break;
  case 1024: 
    dstfunc.dotP.dotQ = &dot4_1024;
    dstfunc.dotP.nthreadsQ = 512;
    dstfunc.dotP.dotCrp = &crpdot2_1024;
    dstfunc.dotP.nthreadsCrp = 512;
    break;
  case 2048:
    dstfunc.dotP.dotQ = &dot4_2048;
    dstfunc.dotP.nthreadsQ = 512;
    dstfunc.dotP.dotCrp = &crpdot2_2048;
    dstfunc.dotP.nthreadsCrp = 512;
    break;
  }


  /*===========Initialize the output===================*/
#ifndef __DOUBLE__
  plhs[0] = mxCreateNumericMatrix(k, numQueries, mxSINGLE_CLASS, mxREAL);
  plhs[1] = mxCreateNumericMatrix(k, numQueries, mxSINGLE_CLASS, mxREAL);
#else
  plhs[0] = mxCreateDoubleMatrix(k, numQueries, mxREAL);
  plhs[1] = mxCreateDoubleMatrix(k, numQueries, mxREAL);
#endif
  knntype* values = (knntype*)mxGetData(plhs[0]);
  knntype* indices = (knntype*)mxGetData(plhs[1]);

  knntype* outbuff_dist = (knntype*)malloc(k*numStreams*CorpusBlocks*numQueries*sizeof(knntype));
  knntype* outbuff_idx = (knntype*)malloc(k*numStreams*CorpusBlocks*numQueries*sizeof(knntype));

  CUstream *stream = (CUstream*)malloc(numStreams*sizeof(CUstream));

  for(int i=0; i<numStreams; i++){
    cuStreamCreate(&stream[i], 0);
  }

  /*Memory allocation*/
  knntype *distbuff, *idxbuff;
  cudaMalloc((void**)&d_data, maxObjects*attributes*sizeof(knntype));
  cudaMalloc((void**)&d_dotp, maxObjects*sizeof(knntype));
  cudaMalloc((void**)&d_dist, maxObjects*numQueries*sizeof(knntype));
  cudaMalloc((void**)&d_labels, maxObjects*numQueries*sizeof(knntype));
  cudaMalloc((void**)&distbuff, numQueries*maxObjects*sizeof(knntype));
  cudaMalloc((void**)&idxbuff, numQueries*maxObjects*sizeof(knntype));

  cudaMemcpyAsync(d_query, query, numQueries*attributes*sizeof(knntype), cudaMemcpyHostToDevice, stream[0]);

  /*compute the dot product of the queries*/
  dstfunc.dotP.dotQ<<<numQueries, dstfunc.dotP.nthreadsQ, 0, stream[0]>>>(d_dotB, d_query);

  cudaEventRecord(start, 0);

  for(int ii=0, c = 0; ii<objects; ii+=maxObjects, c++){

    int CorpusBlockSize = min(maxObjects, objects-ii);
    int StreamSize = CorpusBlockSize / numStreams;
    for(int jj=0; jj<numStreams; jj++){
      cudaMemcpyAsync(d_data + jj*StreamSize*attributes, data + ii*attributes + jj*StreamSize*attributes, StreamSize*attributes*sizeof(knntype), cudaMemcpyHostToDevice, stream[jj]);
    }

    for(int jj=0; jj<numStreams; jj++){
      cuknnsBitonicSTR(d_dist + jj*StreamSize*numQueries, d_data + jj*StreamSize*attributes, d_query, d_labels + jj*StreamSize*numQueries, d_dotp+jj*StreamSize, d_dotB, distbuff + jj*StreamSize*numQueries, idxbuff + jj*StreamSize*numQueries, StreamSize, attributes, numQueries, k, handle, stream[jj], &TimesOut, c*numStreams + jj, &dstfunc);
    }


    for(int jj=0; jj<numStreams; jj++){
      cudaMemcpyAsync(outbuff_dist + jj*k*numQueries + c*numStreams*k*numQueries, d_dist + jj*StreamSize*numQueries, k*numQueries*sizeof(knntype), cudaMemcpyDeviceToHost, stream[jj]);
      cudaMemcpyAsync(outbuff_idx + jj*k*numQueries + c*numStreams*k*numQueries, d_labels + jj*StreamSize*numQueries, k*numQueries*sizeof(knntype), cudaMemcpyDeviceToHost, stream[jj]);
    }
  }

  cuCtxSynchronize();


  cudaEventRecord(stop, 0);
  cudaEventSynchronize(stop);

  cudaEventElapsedTime(&elapsedTime, start, stop);
  TimesOut.knn_time = elapsedTime / 1000;
  knnTime[0] = TimesOut.knn_time;
  dstTime[0] = TimesOut.dst_time / 1000;
  srchTime[0] = TimesOut.srch_time / 1000;

  //printf("reordering....\n");

  if(CorpusBlocks*numStreams>1){
    mergeRes(outbuff_dist, outbuff_idx, k, numQueries, numStreams*CorpusBlocks);
  }

  //printf("done reordering....\n");
  
  memcpy(values, outbuff_dist, k*numQueries*sizeof(knntype));
  memcpy(indices, outbuff_idx, k*numQueries*sizeof(knntype));

  for(int i=0; i<numStreams; i++){
    cuStreamDestroy(stream[i]);
  }

  free(stream);
  free(outbuff_dist);
  free(outbuff_idx);

  /*======clean======*/
  cublasDestroy(handle);
  cudaFree(d_data);
  cudaFree(d_dotp);
  cudaFree(d_dotB);
  cudaFree(d_query);
  cudaFree(d_dist);
  cudaFree(d_labels);
  cudaFree(distbuff);
  cudaFree(idxbuff);
  cudaDeviceReset();
  /////////
}








