/*
//
//                  INTEL CORPORATION PROPRIETARY INFORMATION
//     This software is supplied under the terms of a license agreement or
//     nondisclosure agreement with Intel Corporation and may not be copied
//     or disclosed except in accordance with the terms of that agreement.
//          Copyright(c) 1999-2006 Intel Corporation. All Rights Reserved.
//
//     Intel(R) Integrated Performance Primitives Speech Processing Sample for Windows*
//
//  By downloading and installing this sample, you hereby agree that the
//  accompanying Materials are being provided to you under the terms and
//  conditions of the End User License Agreement for the Intel(R) Integrated
//  Performance Primitives product previously accepted by you. Please refer
//  to the file ippEULA.rtf located in the root directory of your Intel(R) IPP
//  product installation for more information.
//
*/

//---------------------------------------------------------------------
//                  GMM training class
//---------------------------------------------------------------------

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <ipps.h>
#include <ippsr.h>
#include "gmm.h"

#define TRAINCYCLES 3

class MFCC{
public:
    static void readMFCC(char *infile,float **mfccs,int *size,int *step,int *frameNum,int cutEnergy);
};

void MFCC::readMFCC(char *infile,float **mfccs,int *size,int *step,int *frameNum,int cutEnergy){
struct  _MFCChead{
       int   frameNum;
       int   h_framePeriod;
       short frameSize;
       short h_mfccKind;
   } head;
    FILE *fin;
    float *mfcc;
    int   i,step4;
    fin = fopen(infile,"rb");
    if(!fin){
        printf("\n Unable to open file %s \n",infile);
        exit(3);
    }
    if(1!=fread(&head,sizeof(_MFCChead),1,fin)){
        printf("\n Unexpected feature file end \n");
        fclose(fin); exit(3);
    }
    if(cutEnergy) *size = (head.frameSize>>2)-1;
    else          *size = head.frameSize>>2;
    step4 = (*size+3)&(~3);// 16-byte align
    mfcc = ippsMalloc_32f(head.frameNum*step4+4);
    for (i=0;i<head.frameNum;++i){
        if(1!=fread(&mfcc[i*step4],head.frameSize,1,fin)){
            printf("\n Unexpected feature file end \n");
            fclose(fin); exit(3);
        }
    }
    *mfccs = mfcc;
    *step = step4;
    *frameNum = head.frameNum;
    fclose(fin);
    return;
}

typedef struct  {
       short gaussNum;
       short frameSize;
       short cutEnergy;
   }_mixHead;

void GMM::EvalDet(){
  int        j;
    for (j=0;j<gaussianNum;++j){
        ippsUpdateGConst_32f(&cvar[step*j], width,&det[j]);//len*ln(2*p)-sum(ln(cvar[j][i]))
    }
    ippsOutProbPreCalc_32f_I(weight,det,gaussianNum);
}

void GMM::ReadMix(char *infile){
  FILE  *fin;
  _mixHead mhead;
    fin = fopen(infile,"rb");
    if(!fin) printf("Unable to open file %s\n",infile);
    if(1!= fread(&mhead,sizeof(_mixHead),1,fin)){
        printf("\n Unexpected feature file end \n");
        fclose(fin); exit(3);
    }
    width = mhead.frameSize;
    gaussianNum  = mhead.gaussNum;
    cutEnergy = mhead.cutEnergy;
    step = (width+3)&(~3);// 16-byte align

    InitArays(step,gaussianNum);
    if((int)fread(mean,sizeof(float)*step,gaussianNum,fin)<gaussianNum){
        printf("\n Unexpected mixture file end \n");
        fclose(fin); exit(3);
    }
    if((int)fread(cvar,sizeof(float)*step,gaussianNum,fin)<gaussianNum){
        printf("\n Unexpected mixture file end \n");
        fclose(fin); exit(3);
    }
    if((int)fread(weight,sizeof(float),gaussianNum,fin)<gaussianNum){
        printf("\n Unexpected mixture file end \n");
        fclose(fin); exit(3);
    }
    fclose(fin);
    EvalDet();
}

void GMM::WriteMix(char *outfile){
  FILE  *fin;
    _mixHead mhead;
    fin = fopen(outfile,"wb");
    if(!fin) printf("Unable to open file %s\n",outfile);
    mhead.frameSize = width;
    mhead.gaussNum  = gaussianNum;
    mhead.cutEnergy = cutEnergy;
    if( (int)fwrite(&mhead,sizeof(_mixHead),1,fin)<1){
        printf("\n Unexpected mixture file end \n");
        fclose(fin); exit(3);
    }
    if( (int)fwrite(mean,sizeof(float)*step,gaussianNum,fin)<gaussianNum){
        printf("\n Unexpected mixture file end \n");
        fclose(fin); exit(3);
    }
    if( (int)fwrite(cvar,sizeof(float)*step,gaussianNum,fin)<gaussianNum ){
        printf("\n Unexpected mixture file end \n");
        fclose(fin); exit(3);
    }
    if( (int)fwrite(weight,sizeof(float),gaussianNum,fin)<gaussianNum){
        printf("\n Unexpected mixture file end \n");
        fclose(fin); exit(3);
    }
    fclose(fin);
}

void GMM::InitArays(int step,int gaussianNum){
    cvar = ippsMalloc_32f(step*gaussianNum);
    mean = ippsMalloc_32f(step*gaussianNum);
    det  = ippsMalloc_32f(gaussianNum);
    weight = ippsMalloc_32f(gaussianNum);
}

void GMM::InitMixWithCDBK(char *trainf,int maxMixSize,bool CutEnergy){
  float    **featureClass,*tmpBuf,*tmpMean,*tmpVar;
  int        k,i,j,clusters,frames;
  int       *pIndx;
  IppsCdbkState_32f *pCdbk;

    gaussianNum = maxMixSize;
    cutEnergy = CutEnergy;
    MFCC::readMFCC(trainf,&mfccs,&width,&step,&frameNum,cutEnergy);//read train data
    InitArays(step,gaussianNum);
    tmpBuf =ippsMalloc_32f(step*3);
    tmpMean=tmpBuf+step;
    tmpVar =tmpMean+step;
    featureClass=(float**)ippsMalloc_32f(frameNum);
    pIndx = ippsMalloc_32s(frameNum);

    // calculate and invert grand variance
    ippsMeanVarColumn_32f_D2(mfccs,frameNum,step,tmpMean,tmpBuf,width);
    ippsSet_32f(1,tmpVar,width);
    ippsDiv_32f_I(tmpBuf,tmpVar,width);
    // build codebook using Mahalanobis distance with grand variance
    ippsCdbkInitAlloc_WgtL2_32f(&pCdbk,mfccs,tmpVar,width,step,frameNum,gaussianNum,
       IPP_CDBK_KMEANS_LONG);
    // get real number of clusters
    ippsGetCdbkSize_32f(pCdbk,&gaussianNum);
    // quantify training frames
    ippsVQ_32f(mfccs,step,pIndx,frameNum,pCdbk);
    ippsCdbkFree_32f(pCdbk);

    // initialize Gaussian means with non-empty cluster centroids
    // and  Gaussian variances with non-empty cluster variances
    for(clusters=frames=k=0;k<gaussianNum;++k){
        for(j=i=0;i<frameNum;++i){
            if(k==pIndx[i]){
                featureClass[j++]=&mfccs[i*step];
            }
        }
        if (j>1){
            ippsMeanVarColumn_32f_D2L((const float**)featureClass,j,
               &mean[clusters*step],&cvar[clusters*step],width);
            ippsSet_32f(1.0f,tmpBuf,width);
            // invert variances for likelihood calculation, prevent NANs
            if (ippsDiv_32f(&cvar[clusters*step],tmpBuf,&cvar[clusters*step],
               width)==ippStsNoErr) {
               weight[clusters++] = (float)j;
               frames+=j;
            }
        }
    }
    if (clusters<=0) {
       printf("\n Could not initialize mixture \n");
       exit(4);
    } else {
       gaussianNum = clusters;
    }
    // initialize Gaussian weights with cluster sharesm
    ippsDivC_32f_I((float)frames,weight,gaussianNum);
    // log for likelihood calculation
    ippsLn_32f_I(weight,gaussianNum);
    //calculate constants for likelihood calculation
    EvalDet();

    ippsFree(featureClass);
    ippsFree(pIndx);
}

void GMM::TrainMix(int trainCycles){
  float   **postProb,*grad,*tmpWgt,*tmpVar;
  float     max,sum;
  int       j,t,cycles,i,gaussianNum4;
    gaussianNum4 = (gaussianNum+3)&(~3);//16-byte align
    grad     = ippsMalloc_32f(gaussianNum4);
    tmpWgt   = ippsMalloc_32f(frameNum);
    tmpVar   = ippsMalloc_32f(step);
    postProb = (float **)ippsMalloc_32f(frameNum);
    postProb[0] = ippsMalloc_32f(frameNum*gaussianNum4);
    for (i=1;i<frameNum;++i){
        postProb[i] = postProb[0]+gaussianNum4*i;
    }
    for(cycles=0; cycles<trainCycles;++cycles){//training cycle beginning
        for (t=0;t<frameNum;++t){// find posterior probabilities postProb[][]
            // additive term for log likelihood
            ippsCopy_32f(det,postProb[t],gaussianNum);
            // component log likelihoods for frame t
            ippsLogGaussMultiMix_32f_D2(mean,cvar,step,&mfccs[t*step],width,
                postProb[t],gaussianNum);
            // subtract maximum for better exponent
            ippsMax_32f(postProb[t],gaussianNum,&max);
            ippsSubC_32f_I(max,postProb[t],gaussianNum);
            // component likelihoods for frame t - numerators of h[j,t]
            ippsExp_32f_I(postProb[t],gaussianNum);
            // denomerators of h[j,t]
            ippsSum_32f(postProb[t],gaussianNum,&sum,ippAlgHintNone);
            // h[j,t]
            ippsDivC_32f_I(sum,postProb[t],gaussianNum);
        }
        // update weights
        ippsSumColumn_32f_D2((const float*)postProb[0],gaussianNum4,frameNum,grad,
            gaussianNum);
        ippsDivC_32f(grad,(float)frameNum,weight,gaussianNum);
        // log for likelihood calculation
        ippsLn_32f_I(weight,gaussianNum);

        ippsSet_32f(1,cvar,gaussianNum*step);
        for (j=0;j<gaussianNum;++j){
            //transpose column
            for (t=0;t<frameNum;++t)
                tmpWgt[t] = postProb[t][j];
            ippsDivC_32f_I(grad[j],tmpWgt,frameNum);
            //  update mixture mean and variance
            ippsWeightedMeanVarColumn_32f_D2(mfccs,step,tmpWgt,frameNum,&mean[j*step],
               tmpVar,width);
            //inverse for likelihood calculation
            ippsDiv_32f_I(tmpVar,&cvar[j*step],width);
        }
        //calculate constants for likelihood calculation
        EvalDet();
    }//training cycle end
    ippsFree(postProb[0]);
    ippsFree(postProb);
    ippsFree(grad);
    ippsFree(tmpVar);
    ippsFree(tmpWgt);
}

float GMM::LogLHPerFrame(char *TestFile){
  int    j,frames,twidth,step;
  float *testMFCC;
  float  logLH=0,res;
  MFCC::readMFCC(TestFile,&testMFCC,&twidth,&step,&frames,cutEnergy);
    if(twidth!=width || step!=step){
        printf("\nIncompatible feature %s \n",TestFile);
        exit(3);
    }

    for (j=0; j<frames; j++) {
       ippsLogGaussMixture_32f_D2(testMFCC+j*step,mean,cvar,gaussianNum,step,width,det,&res);
       logLH+=res;
    }

    ippsFree(testMFCC);
    return logLH/frames;
}