#include "facecomparison.h"

bool FaceComparison::init_model(string face_det_model,string face_landm_model,string face_rec_model,InitModelError &model_error){
    
    face_det = RetinaFace();
    face_det.confidence_threshold = confidence_threshold;
    face_det.is_bbox_process = is_bbox_process;
    face_det.num_thread=num_thread;
    face_det.forward_type = forward_type;
    bool init_det_code = face_det.init_model(face_det_model);
    model_error.init_error.push_back(init_det_code);
    face_landm =FaceLandmarks();
    bool init_landm_code = face_landm.init_model(face_landm_model);
    face_landm.num_thread=num_thread;
    face_landm.forward_type=forward_type;
    model_error.init_error.push_back(init_landm_code);
    face_rec =FaceRecognize();
    bool init_rec_code = face_rec.init_model(face_rec_model);
    face_rec.num_thread=num_thread;
    face_rec.forward_type=forward_type;
    model_error.init_error.push_back(init_rec_code);
    if(init_det_code&&init_landm_code&&init_rec_code){
        return true;
    }else{
        return false;
    }
}

int FaceComparison::inference(string image1_path,string image2_path,OutputInfo &output_info){
    bool result=false;
    cv::Mat image1=cv::imread(image1_path);
    cv::Mat image2=cv::imread(image2_path);
    if(image1.empty()){
        cout<<"Please check if the first picture path exists"<<endl;
        output_info.error_code=9001;
        return -1;
    }
    if(image2.empty()){
        cout<<"Please check if the second picture path exists"<<endl;
        output_info.error_code=9002;
        return -1;
    }
    vector<Bbox> box1=face_det.inference(image1);
    vector<Bbox> box2=face_det.inference(image2);
    if(box1.empty()){
        cout<<"No faces are detected in the first picture"<<endl;
        output_info.error_code=9003;
        return -1;
    }
    if(box2.empty()){
        cout<<"No faces are detected in the second picture"<<endl;
        output_info.error_code=9004;
        return -1;
    }
    int max_box1=0;
    double max_area1=0,max_area2=0;    
    for(int i=0;i<box1.size();++i){
        double tmp_area1=(box1[i].ymax-box1[i].ymin)*(box1[i].xmax-box1[i].xmin);
        if(tmp_area1>max_area1){
            max_box1=i;
            max_area1=tmp_area1;
        }
    }
    Rect rect1=Rect(box1[max_box1].xmin,box1[max_box1].ymin,box1[max_box1].xmax-box1[max_box1].xmin,box1[max_box1].ymax-box1[max_box1].ymin);
    Mat face_area1=image1(rect1);
    vector<vector<float>> landms1=face_landm.inference(face_area1);
    vector<vector<float>> land1={
        {float(landms1[104][0]),float(landms1[104][1])},
        {float(landms1[105][0]),float(landms1[105][1])},
        {float(landms1[46][0]),float(landms1[46][1])},
        {float(landms1[84][0]),float(landms1[84][1])},
        {float(landms1[90][0]),float(landms1[90][1])}
    };
    Mat align_resize_image1=face_rec.preprocess_face(face_area1,land1);
    for(int j=0;j<box2.size();++j){
        Rect rect2=Rect(box2[j].xmin,box2[j].ymin,box2[j].xmax-box2[j].xmin,box2[j].ymax-box2[j].ymin);
        Mat face_area2=image2(rect2);
        vector<vector<float>> landms2=face_landm.inference(face_area2);
        vector<vector<float>> land2={
            {float(landms2[104][0]),float(landms2[104][1])},
            {float(landms2[105][0]),float(landms2[105][1])},
            {float(landms2[46][0]),float(landms2[46][1])},
            {float(landms2[84][0]),float(landms2[84][1])},
            {float(landms2[90][0]),float(landms2[90][1])}
        };
        Mat align_resize_image2=face_rec.preprocess_face(face_area2,land2);
        double samilar_score=face_rec.inference(align_resize_image1,align_resize_image2);
        if(samilar_score>face_recongnize_thr){
            output_info.error_code=0;
            output_info.result=true;
        }else{
            output_info.error_code=0;
            output_info.result=false;
        }
    }
    return result;
}

int FaceComparison::inference(Mat image1,Mat image2,OutputInfo &output_info){
    if(image1.empty()){
        cout<<"Please check if the first picture path exists"<<endl;
        output_info.error_code=9001;
        return -1;
    }
    if(image2.empty()){
        cout<<"Please check if the second picture path exists"<<endl;
        output_info.error_code=9002;
        return -1;
    }
    vector<Bbox> box1=face_det.inference(image1);
    vector<Bbox> box2=face_det.inference(image2);
    if(box1.empty()){
        cout<<"No faces are detected in the first picture"<<endl;
        output_info.error_code=9003;
        return -1;
    }
    if(box2.empty()){
        cout<<"No faces are detected in the second picture"<<endl;
        output_info.error_code=9004;
        return -1;
    }
    int max_box1=0;
    double max_area1=0,max_area2=0;    
    for(int i=0;i<box1.size();++i){
        double tmp_area1=(box1[i].ymax-box1[i].ymin)*(box1[i].xmax-box1[i].xmin);
        if(tmp_area1>max_area1){
            max_box1=i;
            max_area1=tmp_area1;
        }
    }
    Rect rect1=Rect(box1[max_box1].xmin,box1[max_box1].ymin,box1[max_box1].xmax-box1[max_box1].xmin,box1[max_box1].ymax-box1[max_box1].ymin);
    Mat face_area1=image1(rect1);
    vector<vector<float>> landms1=face_landm.inference(face_area1);
    vector<vector<float>> land1={
        {float(landms1[104][0]),float(landms1[104][1])},
        {float(landms1[105][0]),float(landms1[105][1])},
        {float(landms1[46][0]),float(landms1[46][1])},
        {float(landms1[84][0]),float(landms1[84][1])},
        {float(landms1[90][0]),float(landms1[90][1])}
    };
    Mat align_resize_image1=face_rec.preprocess_face(face_area1,land1);
    for(int j=0;j<box2.size();++j){
        Rect rect2=Rect(box2[j].xmin,box2[j].ymin,box2[j].xmax-box2[j].xmin,box2[j].ymax-box2[j].ymin);
        Mat face_area2=image2(rect2);
        vector<vector<float>> landms2=face_landm.inference(face_area2);
        vector<vector<float>> land2={
            {float(landms2[104][0]),float(landms2[104][1])},
            {float(landms2[105][0]),float(landms2[105][1])},
            {float(landms2[46][0]),float(landms2[46][1])},
            {float(landms2[84][0]),float(landms2[84][1])},
            {float(landms2[90][0]),float(landms2[90][1])}
        };
        Mat align_resize_image2=face_rec.preprocess_face(face_area2,land2);
        double samilar_score=face_rec.inference(align_resize_image1,align_resize_image2);
        if(samilar_score>face_recongnize_thr){
            output_info.error_code=0;
            output_info.result=true;
        }else{
            output_info.error_code=0;
            output_info.result=false;
        }
    }
    return 0;
}