#include "VehiclePlate.h"
#include "vehicle_plate_dr.h"


VehiclePlate::VehiclePlate(/* args */)
{
}

VehiclePlate::~VehiclePlate()
{
    release();
}

int VehiclePlate::init(int dev_id, string model_dir){

    string model_path = model_dir + "/models/vp2/yolo_plate_det310p.om" ;
    string cls_model_path = model_dir + "/models/vp2/plate3cls_310p.om" ;
    string reg_model_path = model_dir + "/models/vp2/plate_reg230316_310p.om" ;

    LOG_INFO("hp 版本：{}  模型路径：{}", vpdr_get_version(), model_path);

    vpd_param dparam;
    char modelNames[100];
    strcpy(modelNames, model_path.c_str());
    dparam.modelNames = modelNames;
    dparam.thresld = 0.2;
    dparam.devId = dev_id;

    vpr_param rparam;
    char cls_modelNames[100];
    strcpy(cls_modelNames, cls_model_path.c_str());
    rparam.cls_modelNames = cls_modelNames;
    char reg_modelNames[100];
    strcpy(reg_modelNames, reg_model_path.c_str());
    rparam.reg_modelNames = reg_modelNames;

    m_devId = dev_id;
    ACL_CALL(aclrtSetDevice(m_devId), ACL_SUCCESS, -1);
    ACL_CALL(aclrtCreateContext(&m_algorthim_ctx, m_devId), ACL_SUCCESS, -1);

    int ret = vpdr_init(&m_handle, dparam, rparam);
    if(ret != 0){
        LOG_ERROR("vc_init error.");
        return -1;
    }

    return 0;
}

vector<VehiclePlateResult> VehiclePlate::process(vector<DeviceMemory*> vec_gpuMem) 
{ // 每个 DeviceMemory 只包含一辆车

    int ret = aclrtSetCurrentContext(m_algorthim_ctx);

    vector<VehiclePlateResult> vec_res;

    int batchsize = vec_gpuMem.size();

    vector<sy_img> batch_img;
    for (int i = 0; i < vec_gpuMem.size(); i++) {
        DeviceMemory* mem = vec_gpuMem[i];

        sy_img img;
        img.w_ = mem->getWidth();     
        img.h_ = mem->getHeight();
        img.c_ = mem->getChannel();   
        img.data_ = mem->getMem();
        batch_img.push_back(img);
    }

    vector<vplates_result> vec_result;
    vec_result.resize(batchsize);
    ret = vpdr_batch(m_handle, batch_img.data(), batchsize, vec_result.data());
    if (SY_SUCCESS != ret) {
        printf("vpr detection process failed!");
        return vec_res;
    }

    for(int batchIdx = 0; batchIdx < batchsize; batchIdx ++){
        auto& plate_info = vec_result[batchIdx].vehicle_plate_infos[0];

        VehiclePlateResult res;
        res.left_ = plate_info.rect.left_;
        res.top_ = plate_info.rect.top_;
        res.width_ = plate_info.rect.width_;
        res.height_ = plate_info.rect.height_;
        res.detect_score = plate_info.detect_score;
        memcpy(res.recg, plate_info.recg, sizeof(PlateNum) * 8);
        res.num_score = plate_info.num_score;
        res.type = plate_info.type;
        res.state = plate_info.state;
        res.state_score = plate_info.state_score;

        vec_res.push_back(res);
    }

    return vec_res;
}

void VehiclePlate::release(){
    if (m_handle){
        vpdr_release(&m_handle);
        m_handle = NULL;
    }
    if(m_algorthim_ctx){
        aclrtDestroyContext(m_algorthim_ctx);
        m_algorthim_ctx = nullptr;
    }
}