#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include "ccnn.h"

const int MAX_ALLOC_LAY = 256;

/****************************************************/
void MyCreateNetLays(struct CcnnContext *pCcnnContext)
{
    int inLayArrID[MAX_ALLOC_LAY];

    int ID = 1, maxSeq = 1;
    int tmpInChannels = 3;
    int tmpInHeight = 28;
    int tmpInWidth = 28;

    pCcnnContext->layNum = 0;

    // 1 Input 
    InputLayer *nInputLayer = new InputLayer(ID++, maxSeq, pCcnnContext->batchNum, tmpInChannels, tmpInHeight, tmpInWidth);
    pCcnnContext->pLcfFatherLayer[pCcnnContext->layNum++] = nInputLayer;

    int tmpKerSize = 3;
    int tmpKerNum = 64;
    float tmpLambda = 0.0F;
    float tmpAlpha = 0.9F;
    int tmpPad = 0;
    int tmpStride = 1;
    int tmpDia = 1;
    int tmpStdPad = 0;
    float learnRate = 0.0001F;
    float lrDecay = 0;
    int lrStep = 0;

    // 2.Conv
    ConvLayer *nConvLayer = new ConvLayer(ID++, maxSeq, "CONV", tmpPad, tmpStride, tmpDia, tmpStdPad, tmpKerNum, tmpKerSize, \
                                          tmpLambda, tmpAlpha, learnRate, lrDecay, lrStep);
    inLayArrID[0] = ID - 2;
    nConvLayer->SetInYFromLayersID(inLayArrID, 1);
    pCcnnContext->pLcfFatherLayer[pCcnnContext->layNum++] = nConvLayer;

    float tmpGamma = 1.0F;
    float tmpBeta = 0.0F;

    // 3.BatchNorm
    BatchNormLayer *nBatchNormLayer = new BatchNormLayer(ID++, maxSeq, tmpGamma, tmpBeta, tmpLambda, \
                                                         tmpAlpha, learnRate, lrDecay, lrStep);
    inLayArrID[0] = ID - 2;    
    nBatchNormLayer->SetInYFromLayersID(inLayArrID, 1);
    pCcnnContext->pLcfFatherLayer[pCcnnContext->layNum++] = nBatchNormLayer;	    

    // 4.Active
    ActiveLayer *nActiveLayer = new ActiveLayer(ID++, maxSeq, "RELU");
    inLayArrID[0] = ID - 2;   
    nActiveLayer->SetInYFromLayersID(inLayArrID, 1);
    pCcnnContext->pLcfFatherLayer[pCcnnContext->layNum++] = nActiveLayer;

    tmpKerSize = 3;
    tmpKerNum = 128;           
    tmpPad = 1;
    tmpStride = 1;
    tmpDia = 1;
    tmpStdPad = 0;

    // 5.CBA = Conv + BatchNorm + Active
    CBAGroupLayer *nCBAGroupLayer = new CBAGroupLayer(ID++, maxSeq, "CONV", tmpPad, tmpStride, tmpDia, tmpStdPad, tmpKerNum, \
                                                      tmpKerSize, tmpGamma, tmpBeta, tmpLambda, tmpAlpha, learnRate, lrDecay, \
                                                      lrStep, "RELU");
    inLayArrID[0] = ID - 2;
    nCBAGroupLayer->SetInYFromLayersID(inLayArrID, 1);
    pCcnnContext->pLcfFatherLayer[pCcnnContext->layNum++] = nCBAGroupLayer; 

    // 6.CBA = Conv + BatchNorm + Active
    tmpKerNum = 256;           
    nCBAGroupLayer = new CBAGroupLayer(ID++, maxSeq, "CONV", tmpPad, tmpStride, tmpDia, tmpStdPad, tmpKerNum, \
                                                      tmpKerSize, tmpGamma, tmpBeta, tmpLambda, tmpAlpha, learnRate, lrDecay, \
                                                      lrStep, "RELU");
    inLayArrID[0] = ID - 2;
    nCBAGroupLayer->SetInYFromLayersID(inLayArrID, 1);
    pCcnnContext->pLcfFatherLayer[pCcnnContext->layNum++] = nCBAGroupLayer; 

    tmpKerSize = 2;
    int rtPad = 0;
    int rbPad = 0;
    int clPad = 0;
    int crPad = 0;
            
    // 7.Max POOL
    MaxPoolLayer *nMaxPoolLayer = new MaxPoolLayer(ID++, maxSeq, "MAXPOOL", NULL, tmpKerSize, rtPad, rbPad, clPad, crPad);
    inLayArrID[0] = ID - 2;       
    nMaxPoolLayer->SetInYFromLayersID(inLayArrID, 1);
    pCcnnContext->pLcfFatherLayer[pCcnnContext->layNum++] = nMaxPoolLayer; 

    // 8.Bias
    BiasLayer *nBiasLayer = new BiasLayer(ID++, maxSeq, tmpLambda, tmpAlpha, learnRate, lrDecay, lrStep);
    inLayArrID[0] = ID - 2;
    nBiasLayer->SetInYFromLayersID(inLayArrID, 1);
    pCcnnContext->pLcfFatherLayer[pCcnnContext->layNum++] = nBiasLayer;

    // 9.FullConnect
    int tmpOutputNum = 10;
    FCLayer *nFCLayer = new FCLayer(ID++, maxSeq, tmpOutputNum, tmpLambda, tmpAlpha, learnRate, lrDecay, lrStep);
    inLayArrID[0] = ID - 2;
    nFCLayer->SetInYFromLayersID(inLayArrID, 1);
    pCcnnContext->pLcfFatherLayer[pCcnnContext->layNum++] = nFCLayer;

    //10.SoftMax
    SoftMaxVecLayer *nSoftMaxVecLayer = new SoftMaxVecLayer(ID++, maxSeq, FL);
    inLayArrID[0] = ID - 2;
    nSoftMaxVecLayer->SetInYFromLayersID(inLayArrID, 1);
    pCcnnContext->pLcfFatherLayer[pCcnnContext->layNum++] = nSoftMaxVecLayer;
}

/***************************************************/
struct CcnnContext *CreateNetStructManual(int method)
{
    struct CcnnContext *pCcnnContext = new struct CcnnContext; 

    //分配内存池
    //分别为： GPU计算内存池cPool，主参数内存池mPool，备用内存池bPool（当mPool不够用的时候使用bPool）
    //cPool 类型需要为 GPU_TYPE GPU类型，分配的是显卡显存
    //mPool 类型可以为显卡GPU_TYPE显存，或普通内存CPU_TYPE。当显存还有剩余的时候，推荐配置显存
    LcfNet::AllocGlobalMem(&(pCcnnContext->cPool), &(pCcnnContext->mPool), &(pCcnnContext->bPool), \
                           0.4F, GPU_TYPE, 0.2F, GPU_TYPE, 1.0F, CPU_TYPE);
                            
    pCcnnContext->pLcfFatherLayer = new LcfFatherLayer *[MAX_ALLOC_LAY];

    //网络基础参数
    pCcnnContext->batchNum = 32;
    pCcnnContext->maxIterNum = 10000;
    pCcnnContext->saveIterNum = 100;
    pCcnnContext->minStopError = 0.001F;
    
    //设置学习参数保存文件名，每saveIterNum次，保存文件为：ccnnWeight.dat_1000 ccnnWeight.dat_2000形式
    snprintf(pCcnnContext->weightFile, sizeof(pCcnnContext->weightFile), "%s", "ccnnWeight.dat");
    
    //method = 1 推理模式
    //method = 0 学习模式
    pCcnnContext->method = method;
    if (1 == pCcnnContext->method) {
        pCcnnContext->batchNum = 1;
    }  
    pCcnnContext->ifLoadWeight = 0;
    
    //自己手动配置网络结构：配置各个算子及其参数
    MyCreateNetLays(pCcnnContext);

    //加载网络内存及参数等
    LcfNet::LoadNetMemAndPara(pCcnnContext);
    
    return pCcnnContext;
}

/*************/
void TrainNet()
{
    int picNum, mPicHWLen;
    PicInfo *trainImgs;
    int *labelData;
    
    struct CcnnContext *pCtx = CreateNetStructManual(0);           
    mPicHWLen = LcfNet::ReadInputHeightWidth(pCtx);
        
    LcfNet::ReadImgClassifySamples(&picNum, trainImgs, labelData, mPicHWLen);
    
    //训练图像分类网络
    LcfNet::NetTrainClassify(pCtx, trainImgs, labelData, picNum);
    
    MemBasicManage::RealeasePicInfo(trainImgs);
    MemBasicManage::RealeaseArrayCPUInt(labelData);    
    
    //释放网络结构及内存参数等
    LcfNet::RealeaseCcnnStructMemory(pCtx);
}

/******************************/
int main(int argc, char *argv[])
{
    TrainNet();
    
    return 1;
}
