#include "yolo.h"

#include <fstream>

#include "yolo.hpp"

namespace csdk_wrap{
    sdkyolo::sdkyolo()
    :objInfo_(NULL)
    ,objCount_(0)
    ,image_(NULL)
    ,width_(0)
    ,height_(0)
    ,channel_(3)
    {}

    sdkyolo::~sdkyolo(){
        for(auto i : detors_){
            delete i;
        }
        if (objInfo_) free(objInfo_);
        if (image_){
            free(image_->data);
            delete image_;
        }
    }

    int sdkyolo::buffer_to_image(const unsigned char *data, const int w, const int h, const int color_channel){
        int size = w*h;
        int size2 = size*2;
    
        int c = color_channel;
        if (w != width_ || h != height_ || color_channel != channel_){
            if (image_){
                free(image_->data);
                delete image_;
            }
            image_ = new image_t;
            image_->h = h;
            image_->w = w;
            image_->c = c;
            image_->data = (float*)calloc(h*w*c, sizeof(float));
        }

        // image im = make_image(w, h, c);
        const unsigned char *srcData = data;
    
        int count = 0;
        switch(c){
            case 1:{
                for (; count < size; ){
                    image_->data[count] = 
                    image_->data[w*h + count] = 
                    image_->data[w*h*2 + count] = 
                    (float)(srcData[count])/255.0;
    
                    ++count;
                }
                break;
            }
            case 3:{
                float* desData = image_->data;
    
                for(int i = 0;i<size;i++){
                    *(desData) = *(srcData + 2) /255.0f;
                    *(desData+size) = *(srcData + 1) /255.0f;
                    *(desData+size2) = *(srcData) /255.0f;
    
                    desData++;
                    srcData+=3;
                }
                break;
            }
    
            default:
                printf("Channel number not supported.\n");
                break;
        }

	    return 0;
    }

    void sdkyolo::objects_names_from_file(std::string filename){
        std::ifstream file(filename);

        if (!file.is_open()){
            printf("open %s file error\n", filename.c_str());
            return;
        }
        for(std::string line; getline(file, line);) names_.push_back(line);

        printf("names count %d\n", names_.size());
    }

    const char* sdkyolo::yolo_obj_name_by_type(const int typ){
        if(names_.empty() || typ < 0 || typ >= names_.size()) return NULL;
        return names_.at(typ).c_str();
    }

    void* sdkyolo::init_yolo(const char *cfg, const char *weights, const char *name, const int gi){
        if(!cfg || !weights || !name){
            printf("init Detector error\n");
            return NULL;
        }
    
        if(names_.empty())
            objects_names_from_file(name);

        auto det = new Detector(cfg, weights, gi);
        detors_.push_back(det);

        return det;
    }

    int sdkyolo::yolo_detect(void *handle, const cIMAGE *img, const float thrsh, const bool use_mean, void **objs, int *objCount){
        Detector *det = (Detector*)handle;

        const int color_channel = img->channel;
        buffer_to_image(img->data, img->width, img->height, color_channel);

        
        std::vector<bbox_t> result_vec = det->detect(*image_, thrsh, use_mean);
		// det->free_image(*im);
        // delete im;

        if (objCount_ < result_vec.size()){
            free(objInfo_);

            objCount_ = result_vec.size();
            objInfo_ = (cObjInfo*)malloc(objCount_ * sizeof(cObjInfo));
        }

        int count = 0;
        for(auto &i : result_vec){
            
            cObjInfo info;
            info.typ = i.obj_id;
            info.prob = i.prob;
            info.rcObj.left = i.x;
			info.rcObj.top = i.y;
			info.rcObj.right = i.x+i.w;
            info.rcObj.bottom = i.y+i.h;
            
            objInfo_[count++] = info;
        }
        
        *objCount = count;
        *objs = objInfo_;

        return objCount_;
    }
}