#include "FeatureAnalysis.h"

#include "../decoder/dvpp/DvppDataMemory.hpp"

struct ObjKey {
        string task_id;
        int object_id;
    };

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

FeatureAnalysis::~FeatureAnalysis()
{
}

int FeatureAnalysis::init(int dev_id, string model_dir)
{
    m_dev_id = dev_id;

    int ret = m_human_feature.init(dev_id, model_dir);
    if (ret < 0){
        return ret;
    }

    ret = m_human_parsing.init(dev_id, model_dir);
    if (ret < 0){
        return ret;
    }

    ret = m_vehicle_color.init(dev_id, model_dir);
    if (ret < 0){
        return ret;
    }

    ret = m_vehicle_plate.init(dev_id, model_dir);
    if (ret < 0){
        return ret;
    }

    ret = m_vpd.init(dev_id, model_dir);
    if (ret < 0){
        return ret;
    }

    return 0;
}

int FeatureAnalysis::update_object_info(ObjectInfo& res_obj, vpc_img_info img_info)
{
    if (img_info.index == 0)
    {
        update_human_info(res_obj, img_info);
    }
    else if (img_info.index >= 4 && img_info.index <= 8) 
    {
        update_car_info(res_obj, img_info);
    }
}

void FeatureAnalysis::update_human_info(ObjectInfo& res_obj, vpc_img_info& img_info){
    vector<DeviceMemory*> vec_Mem;
    unsigned char * pHwData = (unsigned char *)acldvppGetPicDescData(img_info.pic_desc);
    int width = acldvppGetPicDescWidth(img_info.pic_desc);
    int width_stride = acldvppGetPicDescWidthStride(img_info.pic_desc);
    int height = acldvppGetPicDescHeight(img_info.pic_desc);
    int height_stride = acldvppGetPicDescHeightStride(img_info.pic_desc);
    int _size = acldvppGetPicDescSize(img_info.pic_desc);
    DeviceMemory* mem = new DvppDataMemory(width, width_stride, height, height_stride, _size, img_info.task_id, to_string(m_dev_id), false, img_info.task_frame_count, pHwData);
    if (mem) {
        vec_Mem.push_back(mem);
    }

    if (vec_Mem.size() == 1)
    {
        vector<hp_analysis_res> vec_hp_res = m_human_parsing.process(vec_Mem);
        if (vec_hp_res.size() > 0)
        {
            hp_analysis_res hp_res = vec_hp_res[0];
            for (size_t j = 0; j < HP_ATTRI_INDEX_SIZE; j++) {
                res_obj.hp_cls[j].res_index = hp_res.res_objs[j].res_index;
                res_obj.hp_cls[j].res_prob = hp_res.res_objs[j].res_prob;
            }
        }
        
        vector<HFResult> vec_hf_res = m_human_feature.process(vec_Mem);
        if (vec_hf_res.size() > 0)
        {
            HFResult hf_res = vec_hf_res[0];
            memcpy(res_obj.hp_feature, hf_res.feature, sizeof(float) * HUMANREID_FEATURE_SIZE);
        }
    }

    delete mem;
    mem = nullptr;
}

void FeatureAnalysis::update_car_info(ObjectInfo& res_obj, vpc_img_info& img_info){
    vector<DeviceMemory*> vec_Mem;
    unsigned char * pHwData = (unsigned char *)acldvppGetPicDescData(img_info.pic_desc);
    int width = acldvppGetPicDescWidth(img_info.pic_desc);
    int width_stride = acldvppGetPicDescWidthStride(img_info.pic_desc);
    int height = acldvppGetPicDescHeight(img_info.pic_desc);
    int height_stride = acldvppGetPicDescHeightStride(img_info.pic_desc);
    int _size = acldvppGetPicDescSize(img_info.pic_desc);
    DeviceMemory* mem = new DvppDataMemory(width, width_stride, height, height_stride, _size, img_info.task_id, to_string(m_dev_id), false, img_info.task_frame_count, pHwData);
    if (mem) {
        vec_Mem.push_back(mem);
    }

    if (vec_Mem.size() == 1)
    {
        vector<VehicleColorResult> vec_vc_res = m_vehicle_color.process(vec_Mem);
        if (vec_vc_res.size() > 0)
        {
            VehicleColorResult vc_res = vec_vc_res[0];
            res_obj.vehicle_color_index = vc_res.index;
            res_obj.vehicle_color_prob = vc_res.prob;
        }
        
        vector<VPDResult> vec_vpd_res = m_vpd.process(vec_Mem);
        if (vec_vpd_res.size()  > 0)
        {
            VPDResult res = vec_vpd_res[0];
            for (size_t i = 0; i < res.vec_vpd_res.size(); i++)
            {
                auto& info = res.vec_vpd_res[i];
                info.left_ = res_obj.left + info.left_;
                info.top_ = res_obj.top + info.top_;
            }
            
            res_obj.vec_vpd_cls = res.vec_vpd_res;
        }

        vector<VehiclePlateResult> vec_plate_res = m_vehicle_plate.process(vec_Mem);
        if (vec_plate_res.size()  > 0)
        {
            for (size_t i = 0; i < vec_plate_res.size(); i++)
            {
                auto& info = vec_plate_res[i];
                info.left_ = res_obj.left + info.left_;
                info.top_ = res_obj.top + info.top_;
            }

            res_obj.vehicle_plate = vec_plate_res[0];
        }
    }

    delete mem;
    mem = nullptr;
}