// -*- c++ -*-

// This file is part of the Collective Variables module (Colvars).
// The original version of Colvars and its updates are located at:
// https://github.com/Colvars/colvars
// Please update all Colvars source files before making any changes.
// If you wish to distribute your changes, please submit them to the
// Colvars repository at GitHub.

#include "colvarmodule.h"
#include "colvarvalue.h"
#include "colvar.h"
#include "colvarcomp.h"
#include "colvar_neuralnetworkcompute.h"

using namespace neuralnetworkCV;


colvar::neuralNetwork::neuralNetwork()
{
    set_function_type("neuralNetwork");
}


int colvar::neuralNetwork::init(std::string const &conf)
{
    int error_code = linearCombination::init(conf);
    if (error_code != COLVARS_OK) return error_code;
    // the output of neural network consists of multiple values
    // read "output_component" key to determine it
    get_keyval(conf, "output_component", m_output_index);
    // read weight files
    bool has_weight_files = true;
    size_t num_layers_weight = 0;
    std::vector<std::string> weight_files;
    while (has_weight_files) {
        std::string lookup_key = std::string{"layer"} + cvm::to_str(num_layers_weight + 1) + std::string{"_WeightsFile"};
        if (key_lookup(conf, lookup_key.c_str())) {
            std::string weight_filename;
            get_keyval(conf, lookup_key.c_str(), weight_filename, std::string(""));
            weight_files.push_back(weight_filename);
            cvmodule->log(std::string{"Will read layer["} + cvm::to_str(num_layers_weight + 1) + std::string{"] weights from "} + weight_filename + '\n');
            ++num_layers_weight;
        } else {
            has_weight_files = false;
        }
    }
    // read bias files
    bool has_bias_files = true;
    size_t num_layers_bias = 0;
    std::vector<std::string> bias_files;
    while (has_bias_files) {
        std::string lookup_key = std::string{"layer"} + cvm::to_str(num_layers_bias + 1) + std::string{"_BiasesFile"};
        if (key_lookup(conf, lookup_key.c_str())) {
            std::string bias_filename;
            get_keyval(conf, lookup_key.c_str(), bias_filename, std::string(""));
            bias_files.push_back(bias_filename);
            cvmodule->log(std::string{"Will read layer["} + cvm::to_str(num_layers_bias + 1) + std::string{"] biases from "} + bias_filename + '\n');
            ++num_layers_bias;
        } else {
            has_bias_files = false;
        }
    }
    // read activation function strings
    bool has_activation_functions = true;
    size_t num_activation_functions = 0;
    // pair(is_custom_function, function_string)
    std::vector<std::pair<bool, std::string>> activation_functions;
    while (has_activation_functions) {
        std::string lookup_key = std::string{"layer"} + cvm::to_str(num_activation_functions + 1) + std::string{"_activation"};
        std::string lookup_key_custom = std::string{"layer"} + cvm::to_str(num_activation_functions + 1) + std::string{"_custom_activation"};
        if (key_lookup(conf, lookup_key.c_str())) {
            // Ok, this is not a custom function
            std::string function_name;
            get_keyval(conf, lookup_key.c_str(), function_name, std::string(""));
            if (activation_function_map.find(function_name) == activation_function_map.end()) {
                return cvmodule->error("Unknown activation function name: \"" + function_name + "\".\n",
                               COLVARS_INPUT_ERROR);
            }
            activation_functions.push_back(std::make_pair(false, function_name));
            cvmodule->log(std::string{"The activation function for layer["} + cvm::to_str(num_activation_functions + 1) + std::string{"] is "} + function_name + '\n');
            ++num_activation_functions;
#ifdef LEPTON
        } else if (key_lookup(conf, lookup_key_custom.c_str())) {
            std::string function_expression;
            get_keyval(conf, lookup_key_custom.c_str(), function_expression, std::string(""));
            activation_functions.push_back(std::make_pair(true, function_expression));
            cvmodule->log(std::string{"The custom activation function for layer["} + cvm::to_str(num_activation_functions + 1) + std::string{"] is "} + function_expression + '\n');
            ++num_activation_functions;
#endif
        } else {
            has_activation_functions = false;
        }
    }
    // expect the three numbers are equal
    if ((num_layers_weight != num_layers_bias) || (num_layers_bias != num_activation_functions)) {
        return cvmodule->error(
            "Error: the numbers of weights, biases and activation functions do not match.\n",
            COLVARS_INPUT_ERROR);
    }
//     nn = std::make_unique<neuralnetworkCV::neuralNetworkCompute>();
    // std::make_unique is only available in C++14
    if (nn) nn.reset();
    nn = std::unique_ptr<neuralnetworkCV::neuralNetworkCompute>(new neuralnetworkCV::neuralNetworkCompute());
    for (size_t i_layer = 0; i_layer < num_layers_weight; ++i_layer) {
        denseLayer d;
#ifdef LEPTON
        if (activation_functions[i_layer].first) {
            // use custom function as activation function
            try {
                d = denseLayer(weight_files[i_layer], bias_files[i_layer], activation_functions[i_layer].second, cvmodule);
            } catch (std::exception &ex) {
                return cvmodule->error("Error on initializing layer " + cvm::to_str(i_layer) +
                                           " (" + ex.what() + ")\n",
                                       COLVARS_INPUT_ERROR);
            }
        } else {
#endif
            // query the map of supported activation functions
            const auto& f = activation_function_map[activation_functions[i_layer].second].first;
            const auto& df = activation_function_map[activation_functions[i_layer].second].second;
            try {
                d = denseLayer(weight_files[i_layer], bias_files[i_layer], f, df, cvmodule);
            } catch (std::exception &ex) {
                return cvmodule->error("Error on initializing layer " + cvm::to_str(i_layer) +
                                           " (" + ex.what() + ")\n",
                                       COLVARS_INPUT_ERROR);
            }
#ifdef LEPTON
        }
#endif
        // add a new dense layer to network
        if (nn->addDenseLayer(d)) {
            if (cvm::debug()) {
                // show information about the neural network
                cvmodule->log("Layer " + cvm::to_str(i_layer) + " : has " + cvm::to_str(d.getInputSize()) + " input nodes and " + cvm::to_str(d.getOutputSize()) + " output nodes.\n");
                for (size_t i_output = 0; i_output < d.getOutputSize(); ++i_output) {
                    for (size_t j_input = 0; j_input < d.getInputSize(); ++j_input) {
                        cvmodule->log("    weights[" + cvm::to_str(i_output) + "][" + cvm::to_str(j_input) + "] = " + cvm::to_str(d.getWeight(i_output, j_input)));
                    }
                    cvmodule->log("    biases[" + cvm::to_str(i_output) + "] = " + cvm::to_str(d.getBias(i_output)) + "\n");
                }
            }
        } else {
            return cvmodule->error("Error: error on adding a new dense layer.\n", COLVARS_INPUT_ERROR);
        }
    }
    nn->input().resize(cv.size());
    return error_code;
}

colvar::neuralNetwork::~neuralNetwork() {
}

void colvar::neuralNetwork::calc_value() {
    x.reset();
    for (size_t i_cv = 0; i_cv < cv.size(); ++i_cv) {
        cv[i_cv]->calc_value();
        const colvarvalue& current_cv_value = cv[i_cv]->value();
        // for current nn implementation we have to assume taht types are always scaler
        if (current_cv_value.type() == colvarvalue::type_scalar) {
            nn->input()[i_cv] = cv[i_cv]->sup_coeff * (cvm::pow(current_cv_value.real_value, cv[i_cv]->sup_np));
        } else {
            cvmodule->error("Error: using of non-scaler component.\n");
            return;
        }
    }
    nn->compute();
    x = nn->getOutput(m_output_index);
}

void colvar::neuralNetwork::calc_gradients() {
    for (size_t i_cv = 0; i_cv < cv.size(); ++i_cv) {
        cv[i_cv]->calc_gradients();
        if (cv[i_cv]->is_enabled(f_cvc_explicit_gradient)) {
            const cvm::real factor = nn->getGradient(m_output_index, i_cv);
            const cvm::real factor_polynomial = getPolynomialFactorOfCVGradient(i_cv);
            for (size_t j_elem = 0; j_elem < cv[i_cv]->value().size(); ++j_elem) {
                for (size_t k_ag = 0 ; k_ag < cv[i_cv]->atom_groups.size(); ++k_ag) {
                    for (size_t l_atom = 0; l_atom < (cv[i_cv]->atom_groups)[k_ag]->size(); ++l_atom) {
                        cv[i_cv]->atom_groups[k_ag]->grad_x(l_atom) *= factor_polynomial * factor;
                        cv[i_cv]->atom_groups[k_ag]->grad_y(l_atom) *= factor_polynomial * factor;
                        cv[i_cv]->atom_groups[k_ag]->grad_z(l_atom) *= factor_polynomial * factor;
                    }
                }
            }
        }
    }
}

void colvar::neuralNetwork::apply_force(colvarvalue const &force) {
    for (size_t i_cv = 0; i_cv < cv.size(); ++i_cv) {
        // If this CV us explicit gradients, then atomic gradients is already calculated
        // We can apply the force to atom groups directly
        if (cv[i_cv]->is_enabled(f_cvc_explicit_gradient)) {
            for (size_t k_ag = 0 ; k_ag < cv[i_cv]->atom_groups.size(); ++k_ag) {
                (cv[i_cv]->atom_groups)[k_ag]->apply_colvar_force(force.real_value);
            }
        } else {
            // Compute factors for polynomial combinations
            const cvm::real factor_polynomial = getPolynomialFactorOfCVGradient(i_cv);
            const cvm::real factor = nn->getGradient(m_output_index, i_cv);;
            colvarvalue cv_force = force.real_value * factor * factor_polynomial;
            cv[i_cv]->apply_force(cv_force);
        }
    }
}


cvm::real colvar::neuralNetwork::dist2(colvarvalue const &x1, colvarvalue const &x2) const
{
  return x1.dist2(x2);
}


colvarvalue colvar::neuralNetwork::dist2_lgrad(colvarvalue const &x1, colvarvalue const &x2) const
{
  return x1.dist2_grad(x2);
}


colvarvalue colvar::neuralNetwork::dist2_rgrad(colvarvalue const &x1, colvarvalue const &x2) const
{
  return x2.dist2_grad(x1);
}



void colvar::neuralNetwork::wrap(colvarvalue & /* x_unwrapped */) const {}
