#include <iostream>
#include <lt/lt.h>
#include <chrono>
#include <filesystem>
using namespace std::chrono;

std::string exec_path;


#define PBSTR "||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||"
#define PBWIDTH 60

void printProgress(double percentage, float t) {
    int val = (int)(percentage * 100);
    int lpad = (int)(percentage * PBWIDTH);
    int rpad = PBWIDTH - lpad;
    printf("\r%3d%% [%.*s%*s]", val, lpad, PBSTR, rpad, "");
    printf("\t%f (ms)", t);
    fflush(stdout);
}


static int res_theta_i = 1024;
static int res_theta_m = 1024;
static int res_sigma = 1024;

std::string print_float(const float& f) {
    std::string str = std::to_string(f);
    if (str.find('.') == std::string::npos) {
        str += ".0";
    }
    str += "f";
    return str;
}

void write_lut(std::vector<std::vector<float>>& tab_sigma, std::vector<std::vector<float>>& inv_tab_sigma) {

    std::vector<float> thetas_i = lt::linspace<float>(0., lt::pi * 0.5, res_theta_i);
    std::vector<float> thetas_m = lt::linspace<float>(0., lt::pi * 0.5, res_theta_m);

    std::stringstream out;
    out << "#pragma one\n";
    out << "static int res_theta_i = " << res_theta_i << ";\n";
    out << "static int res_theta_m = " << res_theta_m << ";\n";
    out << "static int res_sigma = " << res_sigma << ";\n";

    // Write theta_m array
    out << "static float thetas_m[" << res_theta_m << "]"  << "{\n";
    for (int j = 0; j < res_theta_m; j++) {

        out << print_float(thetas_m[j]);
        if (j < res_theta_m - 1)
            out << ",";
    }
    out << "};\n";

    // Write theta_i array
    out << "static float thetas_i[" << res_theta_i << "]" << "{\n";
    for (int j = 0; j < res_theta_i; j++) {

        out << print_float(thetas_i[j]);
        if (j < res_theta_i - 1)
            out << ",";
    }
    out << "};\n";


    // Write sigma look up table
    out << "static float tab_sigma[" << res_theta_i << "]" << "[" << res_theta_m << "] = " << "{\n";

    for (int i = 0; i < res_theta_i; i++) {
        out << "{";

        for (int j = 0; j < res_theta_m; j++) {

            out << print_float(tab_sigma[i][j]);
            if (j < res_theta_m - 1)
                out << ",";
        }
        out << "}";
        if (i < res_theta_i - 1)
            out << ",";
        out << "\n";
    }

    out << "};\n";


    // Write inverse sigma look up table
    out << "static float inv_tab_sigma[" << res_theta_i << "]" << "[" << res_sigma << "] = " << "{\n";

    for (int i = 0; i < res_theta_i; i++) {
        out << "{";

        for (int j = 0; j < res_sigma; j++) {

            out << print_float(inv_tab_sigma[i][j]);
            if (j < res_sigma - 1)
                out << ",";
        }
        out << "}";
        if (i < res_theta_i - 1)
            out << ",";
        out << "\n";
    }

    out << "};\n";

    // Save to file
    std::ofstream outFile;
    outFile.open(exec_path+"sigma_tab.h");
    outFile << out.rdbuf();
    outFile.close();
    lt::Log(lt::logHighlight) << exec_path << "sigma_tab.h saved";
}

void fill_lut(std::vector<std::vector<float>>& tab_sigma, std::vector<std::vector<float>>& inv_tab_sigma) {

    std::vector<float> thetas_i = lt::linspace<float>( 0., lt::pi * 0.5, res_theta_i);
    std::vector<float> thetas_m = lt::linspace<float>( 0., lt::pi * 0.5, res_theta_m);
    lt::MicrograinMicrosurface ms;

    for (int i = 0; i < res_theta_i; i++) {
        
        float cos_theta_i = std::cos(thetas_i[i]);
        float sin_theta_i = std::sqrt(1. - cos_theta_i * cos_theta_i);
        lt::vec3 wi_u(0, sin_theta_i, cos_theta_i);

        for (int j = 0; j < res_theta_m; j++) {

            float cos_theta_m = std::cos(thetas_m[j]);
            float sin_theta_m = std::sqrt(1. - cos_theta_m * cos_theta_m);
            lt::vec3 wm_u(0, sin_theta_m, cos_theta_m);

            float sigma_b = ms.sigma_base(wm_u);
            float sigma_s = ms.sigma_shadow(wm_u, wi_u);
            float sigma_ = sigma_b + sigma_s;

            tab_sigma[i][j] = sigma_b + sigma_s;
        }

    }

    for (int i = 0; i < res_theta_i; i++) {
        
        float cos_theta_i = std::cos(thetas_i[i]);
        float sin_theta_i = std::sqrt(1. - cos_theta_i * cos_theta_i);
        lt::vec3 wi_u(0, sin_theta_i, cos_theta_i);

        std::vector<float> sigmas = lt::linspace<float>(0, ms.sigma(wi_u), res_sigma, false);
     
        for (int j = 0; j < res_sigma; j++) {
            int idx = lt::binary_search<float>(tab_sigma[i], sigmas[j]);
            inv_tab_sigma[i][j] = thetas_m[idx];
        }

    }

}

void generate_look_up_table() {
    std::vector<std::vector<float>> tab_sigma;
    std::vector<std::vector<float>> inv_tab_sigma;

    tab_sigma.resize(res_theta_i);
    inv_tab_sigma.resize(res_theta_i);

    for (auto& v : tab_sigma)
        v.resize(res_theta_m);
    for (auto& v : inv_tab_sigma)
        v.resize(res_sigma);

    fill_lut(tab_sigma, inv_tab_sigma);
    write_lut(tab_sigma, inv_tab_sigma);
}

int main(int argc, char* argv[]) {
    std::filesystem::path executable_path(argv[0]);
    std::filesystem::path current_path = std::filesystem::current_path();
    exec_path = std::filesystem::relative(executable_path, current_path).remove_filename().generic_string();
    
    lt::State::log_level = lt::logHighlight;
    lt::State::exectuable_path = exec_path;

    lt::Log(lt::logHighlight) << "micrograin: generation of look up tables ...";
    generate_look_up_table();

    
    return EXIT_SUCCESS;

}