#include "VehicleHeadTail.h"
#include "../utils/logger.hpp"

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

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

int VehicleHeadTail::init(head_tail_param param){

    ACL_CALL(aclrtCreateContext(&m_ctx, param.devId), ACL_SUCCESS, SY_FAILED);
    

    LOG_INFO("head_tail version: head_tail_256.2024.12.18");

    string model_path = param.sdk_root + "/models/car_head_tail/head_tail_256_241220_310P.om";

    m_cnn_cls = new CnnCls();
    int ret = m_cnn_cls->Init(model_path.c_str());
    if (ret != SY_SUCCESS) {
        delete m_cnn_cls;
        m_cnn_cls = nullptr;
        LOG_ERROR("sy_hcp model init failed!");
        return SY_FAILED;
    }

    max_batch = param.max_batch;

    ACL_CALL(aclrtCreateStream(&stream), ACL_SUCCESS, SY_FAILED);
    m_dvpp = new DvppProcessx();
    ret = m_dvpp->InitResource(stream);
    if (ret != SY_SUCCESS) {
        delete m_dvpp;
        m_dvpp = nullptr;
        LOG_ERROR("dvpp init failed!");
        return SY_FAILED;
    }

    return SY_SUCCESS;
}

int VehicleHeadTail::detect(vector<sy_img> img_data_array, vector<HeadTailResult>& result){

    if (m_cnn_cls == NULL) {
        LOG_ERROR("HumanCarParse get null handle!");
	    return SY_FAILED;
    }

    ACL_CALL(aclrtSetCurrentContext(m_ctx), ACL_SUCCESS, SY_FAILED);	

    int inputW = m_cnn_cls->GetInputWidth();
    int inputH = m_cnn_cls->GetInputHeight();

    for (int i = 0; i < img_data_array.size(); i++) {
        if (img_data_array[i].data_ == NULL || img_data_array[i].w_ == 0 || img_data_array[i].h_ == 0) {
            LOG_ERROR("HumanCarParse Get null input ptr!");
	        return SY_FAILED;
        }
        ImageData resizeImg, src;
        ACL_CALL(Utils::CopysyImageDataToDvppV2(src, img_data_array[i]), SY_SUCCESS, SY_FAILED);
        ACL_CALL(m_dvpp->CropAndPaste(resizeImg, src, inputW, inputH), SY_SUCCESS, SY_FAILED);
        int ret = m_cnn_cls->Inference(resizeImg);
        if (ret != SY_SUCCESS) {
            LOG_ERROR( "HumanCarParse process error!");
            return SY_MODEL_FORWARD_ERROR;
        }

        vector<float> hcpRes;
        ret = m_cnn_cls->PostProcess(hcpRes);
        if (ret != SY_SUCCESS) {
            LOG_ERROR("HumanCarParse postprocess error!");
            return SY_MODEL_GETRESULT_ERROR; 
        }
        if (hcpRes.size() != 2) {
            LOG_ERROR("HumanCarParse postprocess error!");
            return SY_MODEL_GETRESULT_ERROR;
        }

        HeadTailResult one_result;
        one_result.cls = hcpRes[0];
        one_result.confidence = hcpRes[1];

        result.push_back(one_result);
    }

    return SY_SUCCESS;
}

int VehicleHeadTail::release() {

    ACL_CALL(aclrtSetCurrentContext(m_ctx), ACL_ERROR_NONE, SY_FAILED);

    if (m_cnn_cls) {
        delete m_cnn_cls;
        m_cnn_cls = nullptr;
    }

    if (m_dvpp) {
        delete m_dvpp;
        m_dvpp = nullptr;
    }

    if (stream != nullptr) {
        int ret = aclrtDestroyStream(stream);
        if (ret != ACL_SUCCESS) {
            LOG_ERROR("destroy stream failed");
        }
        stream = nullptr;
    }

    aclrtDestroyContext(m_ctx);
}