#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include "ccnn.h"
    
/*****************************/
void PrintUsage(char *exeNname)
{
    printf("# -------------------------------------------------------------------------------#\r\n");
    printf("#                         ccnn-1.2 (2025-03-07)                                  #\r\n");
    printf("# A C++ Lightweight artificial intelligence learning and reasoning framework     #\r\n");
    printf("# Email: zqgrmkl@sohu.com                                                        #\r\n");
    printf("# ifLoadWeight is 1 for load Pre-training weight train. 0 is new train.          #\r\n");
    printf("# netType is 1 for Image classify.                                               #\r\n");
    printf("# netType is 2 for Image Semantic segmentation.                                  #\r\n");
    printf("# netType is 3 for Image OneStage object detection.                              #\r\n");
    printf("# netType is 4 for vector Multiple regression.                                   #\r\n");
    printf("# netType is 5 for vector classify.                                              #\r\n");
    printf("# netType is 6 for GRU RNN.                                                      #\r\n");
    printf("# netType is 7 for NLP + Attention.                                              #\r\n");
    printf("# -------------------------------------------------------------------------------#\r\n");
    printf("%s ifLoadWeight netType confFile\r\n", exeNname);
}

/******************************************************************/
void MainTrainClassify(const char *mCnfXMLName, int ifLoadPreWeight)
{
    int picNum, mPicHWLen;
    PicInfo *trainImgs;
    int *labelData;
    
    struct CcnnContext *pCtx = LcfNet::CreateNetFromXML(mCnfXMLName, 0, ifLoadPreWeight);           
    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);
}

/*************************************************************/
void MainTrainSeg(const char *mCnfXMLName, int ifLoadPreWeight)
{
    int picNum, mPicHWLen, segClsNum;
    PicInfo *trainImgs, *labelData;
    
    struct CcnnContext *pCtx = LcfNet::CreateNetFromXML(mCnfXMLName, 0, ifLoadPreWeight);       
    mPicHWLen = LcfNet::ReadInputHeightWidth(pCtx);
    
    LcfNet::ReadSemanticSamples(&picNum, trainImgs, labelData, mPicHWLen);
    
    ReadPixCnfClassify mReadPixCnfClassify;
    PixClassify *pSeg = mReadPixCnfClassify.GetClassifyFromCnfFile(mCnfXMLName, &segClsNum);    
    LcfNet::NetTrainSeg(pCtx, pSeg, segClsNum, trainImgs, labelData, picNum);
    delete []pSeg;
    
    MemBasicManage::RealeasePicInfo(trainImgs);
    MemBasicManage::RealeasePicInfo(labelData);    
    LcfNet::RealeaseCcnnStructMemory(pCtx);
}

/*****************************************************************/
void MainTrainObjsDet(const char *mCnfXMLName, int ifLoadPreWeight)
{
    int picNum, mPicHWLen;
    struct PicGroundTrueBox *gtBox;
    PicInfo *trainImgs;
    
    struct CcnnContext *pCtx = LcfNet::CreateNetFromXML(mCnfXMLName, 0, ifLoadPreWeight);   
    mPicHWLen = LcfNet::ReadInputHeightWidth(pCtx);
            
    char *strNameList = LcfNet::ReadSampleClassNameStrList(mCnfXMLName);
    LcfNet::ReadObjectsDetSamples(&picNum, gtBox, trainImgs, mPicHWLen, strNameList);

    LcfNet::NetTrainObjectsDet(pCtx, trainImgs, gtBox, picNum);

    MemBasicManage::RealeasePicInfo(trainImgs);
    MemBasicManage::RealeaseArrayPicGroundTrueBox(gtBox);
    MemBasicManage::RealeaseArrayCPUChar(strNameList); 
    LcfNet::RealeaseCcnnStructMemory(pCtx);
}

/**********************************************************************/
void MainTrainMultiRegress(const char *mCnfXMLName, int ifLoadPreWeight)
{
    int sampleNum, inputDim, outputDim;
    float *trainData, *lableNum;
    
    struct CcnnContext *pCtx = LcfNet::CreateNetFromXML(mCnfXMLName, 0, ifLoadPreWeight);       
    inputDim = LcfNet::ReadInputChannels(pCtx);
    outputDim = LcfNet::ReadOutputChannels(pCtx);
            
    LcfNet::ReadMultiRegressSamples(&sampleNum, trainData, lableNum, inputDim, outputDim, 1, 1);
        
    LcfNet::NetTrainMultiRegress(pCtx, trainData, lableNum, sampleNum, inputDim, outputDim);

    MemBasicManage::RealeaseArrayCPUFloat(trainData);
    MemBasicManage::RealeaseArrayCPUFloat(lableNum);
    LcfNet::RealeaseCcnnStructMemory(pCtx);
}

/*******************************************************************/
void MainTrainVectorCls(const char *mCnfXMLName, int ifLoadPreWeight)
{
    int sampleNum, inputDim, outputDim;
    float *trainData, *lableNum;

    struct CcnnContext* pCtx = LcfNet::CreateNetFromXML(mCnfXMLName, 0, ifLoadPreWeight);
    inputDim = LcfNet::ReadInputChannels(pCtx);
    outputDim = 1;

    LcfNet::ReadMultiRegressSamples(&sampleNum, trainData, lableNum, inputDim, outputDim, 1, 0);

    LcfNet::NetTrainVectorCls(pCtx, trainData, lableNum, sampleNum, inputDim);

    MemBasicManage::RealeaseArrayCPUFloat(trainData);
    MemBasicManage::RealeaseArrayCPUFloat(lableNum);
    LcfNet::RealeaseCcnnStructMemory(pCtx);
}

/********************************************************************/
void MainTrainRNNRegress(const char *mCnfXMLName, int ifLoadPreWeight)
{
    struct CcnnContext *pCtx = LcfNet::CreateNetFromXML(mCnfXMLName, 0, ifLoadPreWeight);
    int inputDim = LcfNet::ReadInputChannels(pCtx);
    int tmpSeqNum = LcfNet::ReadInputMaxSeqNum(pCtx);

    struct RNNDataInfo *tmpStruct = LcfNet::ReadRNNInputData(inputDim, tmpSeqNum, pCtx->mBatchNum, 1, 1);

    LcfNet::MainTrainRNNRegress(pCtx, tmpStruct->intputData, tmpStruct->labelData, tmpSeqNum, tmpStruct->batchCount, inputDim, inputDim);

    LcfNet::RealeaseRNNDDataStruct(tmpStruct);
    LcfNet::RealeaseCcnnStructMemory(pCtx);  
}

/*************************************************************/
void MainTrainNLP(const char *mCnfXMLName, int ifLoadPreWeight)
{
    struct CcnnContext *pCtx = LcfNet::CreateNetFromXML(mCnfXMLName, 0, ifLoadPreWeight);
   
    int sentenceCount, *dataBuff, *seqArr, *labelBuff, *labelSeqArr;
    ReadNlpSample mReadNlpSample;
    mReadNlpSample.LoadSeqNumSampleAndLablel("cnVector.txt", "enVector.txt", &sentenceCount, dataBuff, seqArr, labelBuff, labelSeqArr);
           
    int inputDim = LcfNet::ReadInputChannels(pCtx);
    int outputDim = LcfNet::ReadOutputChannels(pCtx);
    LcfNet::MainTrainNLP(pCtx, dataBuff, labelBuff, sentenceCount, \
                         seqArr, labelSeqArr, inputDim, outputDim);

    LcfNet::RealeaseCcnnStructMemory(pCtx);
    MemBasicManage::RealeaseArrayCPUInt(dataBuff); 
    MemBasicManage::RealeaseArrayCPUInt(labelBuff); 
    MemBasicManage::RealeaseArrayCPUInt(seqArr); 
    MemBasicManage::RealeaseArrayCPUInt(labelSeqArr); 
}

/******************************/
int main(int argc, char *argv[])
{    
    if (argc != 4) {
        PrintUsage(argv[0]);
        return 1;
    }

    int ifLoadPreWeight = atoi(argv[1]);
    int netType = atoi(argv[2]);
 
    if (1 == netType) {
        MainTrainClassify(argv[3], ifLoadPreWeight);
    } else if (2 == netType) {
        MainTrainSeg(argv[3], ifLoadPreWeight);
    } else if (3 == netType) {
        MainTrainObjsDet(argv[3], ifLoadPreWeight);
    } else if (4 == netType) {
        MainTrainMultiRegress(argv[3], ifLoadPreWeight);
    } else if (5 == netType) {
        MainTrainVectorCls(argv[3], ifLoadPreWeight);
    } else if (6 == netType) {
        MainTrainRNNRegress(argv[3], ifLoadPreWeight);
    } else if (7 == netType) {
        MainTrainNLP(argv[3], ifLoadPreWeight);
    } else {
        PrintUsage(argv[0]);
    }

    return 0;
}
