//
//  classify.cpp
//  transfer-learning
//
//  Created by Kushashwa Ravi Shrimali on 15/08/19.
//

#include <iostream>
#include <torch/torch.h>
#include <opencv2/opencv.hpp>
#include <torch/script.h>
#include <dirent.h>

// Utility function to load image from given folder
// File type accepted: .jpg
std::vector<std::string> load_images(std::string folder_name) {
    std::vector<std::string> list_images;
    std::string base_name = folder_name;
    DIR* dir;
    struct dirent *ent;
    if((dir = opendir(base_name.c_str())) != NULL) {
        while((ent = readdir(dir)) != NULL) {
            std::string filename = ent->d_name;
            if(filename.length() > 4 && filename.substr(filename.length() - 3) == "jpg") {
                std::string newf = base_name + filename;
                list_images.push_back(newf);
            }
        }
    }
    return list_images;
}

void print_probabilities(std::string loc, std::string model_path, std::string model_path_linear) {
    // Load image with OpenCV.
    cv::Mat img = cv::imread(loc);
    cv::resize(img, img, cv::Size(224, 224), cv::INTER_CUBIC);
    // Convert the image and label to a tensor.
    torch::Tensor img_tensor = torch::from_blob(img.data, {1, img.rows, img.cols, 3}, torch::kByte);
    img_tensor = img_tensor.permute({0, 3, 1, 2}); // convert to CxHxW
    img_tensor = img_tensor.to(torch::kF32);
    
    // Load the model.
    torch::jit::script::Module model;
    model = torch::jit::load(model_path);
    
    torch::nn::Linear model_linear(512, 2);
    torch::load(model_linear, model_path_linear);
    
    // Predict the probabilities for the classes.
    std::vector<torch::jit::IValue> input;
    input.push_back(img_tensor);
    torch::Tensor prob = model.forward(input).toTensor();
    prob = prob.view({prob.size(0), -1});
    prob = model_linear(prob);
    
    std::cout << "Printing for location: " << loc << std::endl;
    std::cout << "Cat prob: " << *(prob.data<float>())*100. << std::endl;
    std::cout << "Dog prob: " << *(prob.data<float>()+1)*100. << std::endl;
}

int main(int arc, char** argv)
{
    // argv[1] should is the test image
    std::string location = argv[1];
    
    // argv[2] contains pre-trained model without last layer
    // argv[3] contains trained last FC layer
    std::string model_path = argv[2];
    std::string model_path_linear = argv[3];
    
    // Load the model.
    // You can also use: auto model = torch::jit::load(model_path);
    torch::jit::script::Module model = torch::jit::load(model_path);
    
    // Print probabilities for dog and cat classes
    print_probabilities(location, model_path, model_path_linear);
    return 0;
}