#include <cwr_app.h>
#include <boost/filesystem.hpp>
#include <cblas.h>

using namespace std;
using namespace std::chrono;
using namespace caffe;
using namespace boost::filesystem;

bool CwrApp::caffe_initialized = false;
CwrApp::CwrApp() = default;

CwrApp::CwrApp(const std::string &solver_path,
               const std::string &initial_weights_path,
               int classes_n,
               const std::string &initial_prediction_layer,
               const std::vector<std::string> &intial_labels,
               const cv::Mat& initial_mean_image) {
    std::stringstream debug_str;

    this->training_disabled = true;
    if(!caffe_initialized) {
        caffe_initialized = true;
        google::InitGoogleLogging("");
        google::InstallFailureSignalHandler();
        Caffe::set_mode(Caffe::CPU);

        debug_str << "OpenBLAS threads: " << openblas_get_num_threads();
        log_android_debug(debug_str);
    }

    this->solver_path = solver_path;
    this->weights_path = initial_weights_path;
    {
        SolverParameter solver_param;
        ReadSolverParamsFromTextFileOrDie(solver_path, &solver_param);

        solver = caffe::SolverRegistry<float>::CreateSolver(solver_param);
    }

    net = solver->net();
    net->CopyTrainedLayersFrom(initial_weights_path);

    test_net = solver->test_nets()[0];
    test_net->ShareTrainedLayersWith(net.get());

    this->train_layer_name_to_id = generate_layers_name_to_id_mapping(this->net);
    this->test_layer_name_to_id = generate_layers_name_to_id_mapping(this->test_net);
    this->train_layer_names_to_param_id = generate_layers_name_to_param_id_mapping(this->net);

    this->preallocate_memory();

    //test_net->CopyTrainedLayersFrom(initial_test_weights_path);

    extract_minibatch_size_from_prototxt_with_input_layers(net, test_net,
                                                           &minibatch_size,
                                                           &inference_minibatch_size);

    labels = intial_labels;
    training_batch_images = {};
    training_batch_features = {};
    mean_image = initial_mean_image;
    prediction_layer = initial_prediction_layer;
    this->epochs = 0;
}

CwrApp::CwrApp(
        const std::string &solver_path,
        const std::string &initial_weights_path,
        const std::vector<float> &initial_class_updates,
        int train_epochs,
        int initial_classes_n,
        const std::vector<std::string> &cwr_layers,
        const std::string &initial_prediction_layer,
        const std::vector<std::string> &intial_labels,
        const Mat& initial_mean_image,
        const std::string &training_pre_extraction_layer,
        const std::string &backward_stop_layer) {
    std::stringstream debug_str;

    //TODO: add parameters checks
    this->training_disabled = false;
    this->feature_extraction_layer = training_pre_extraction_layer;
    this->backward_stop_layer = backward_stop_layer;
    if(!caffe_initialized) {
        caffe_initialized = true;
        google::InitGoogleLogging("");
        google::InstallFailureSignalHandler();
        Caffe::set_mode(Caffe::CPU);
        debug_str << "OpenBLAS threads: " << openblas_get_num_threads();

        log_android_debug(debug_str);
    }

    this->solver_path = solver_path;
    this->weights_path = initial_weights_path;
    {
        SolverParameter solver_param;
        ReadSolverParamsFromTextFileOrDie(solver_path, &solver_param);

        solver = caffe::SolverRegistry<float>::CreateSolver(solver_param);
    }

    net = solver->net();
    net->CopyTrainedLayersFrom(initial_weights_path);

    test_net = solver->test_nets()[0];
    test_net->ShareTrainedLayersWith(net.get());
    //test_net->CopyTrainedLayersFrom(initial_test_weights_path);

    this->train_layer_name_to_id = generate_layers_name_to_id_mapping(this->net);
    this->test_layer_name_to_id = generate_layers_name_to_id_mapping(this->test_net);
    this->train_layer_names_to_param_id = generate_layers_name_to_param_id_mapping(this->net);

    this->preallocate_memory();

    cwr = Cwr(initial_classes_n, initial_class_updates, net, cwr_layers, true);

    cwr.zeros_non_cwr_layers_lr();
    cwr.zeros_cwr_layer_bias_lr(-1.0f);
    this->set_brn_past_weight(20000.0);

    extract_minibatch_size_from_prototxt_with_input_layers(net, test_net,
                                                           &minibatch_size,
                                                           &inference_minibatch_size);

    bool need_target = net_use_target_vectors(net);
    if(!need_target) {
        throw std::runtime_error("Net has no target layer!");
    }

    labels = intial_labels;
    training_batch_images = {};
    training_batch_features = {};
    mean_image = initial_mean_image;
    prediction_layer = initial_prediction_layer;
    this->epochs = train_epochs;
}

Mat CwrApp::create_mean_image(float B, float G, float R) {
    Mat mean(128, 128, CV_32FC3);
    auto mean_data = (float*) mean.data;

    for (int size = 0; size < 3 * 128 * 128; size += 3) { // 3 = channels
        mean_data[size + 0] = B;
        mean_data[size + 1] = G;
        mean_data[size + 2] = R;
    }

    return mean;
}

int CwrApp::get_minibatch_size() {
    return minibatch_size;
}

void CwrApp::add_batch_image(const Mat& new_image) {
    if(this->training_disabled) {
        throw std::runtime_error("Training disabled!");
    }

    stringstream debug_str;

    this->mtx->lock();
    if(this->must_pre_extract_batch_features()) {  // True
        debug_str << "Pre-extracting features from layer " << this->feature_extraction_layer << endl;
        std::chrono::high_resolution_clock::time_point t1, t2;
        auto duration = duration_cast<milliseconds>( t2 - t1 ).count(); // Trick used for duration "auto" definition
        t1 = std::chrono::high_resolution_clock::now();

        this->training_batch_features.push_back(
                std::make_shared<vector<float>>(this->inference(new_image, this->feature_extraction_layer)));

        t2 = std::chrono::high_resolution_clock::now();
        duration = duration_cast<milliseconds>( t2 - t1 ).count();
        debug_str << "Features extraction took: " << duration << " ms" << endl;
    } else {
        debug_str << "Adding pattern without features pre-extraction" << endl;
        this->training_batch_images.push_back(new_image);
    }

    this->mtx->unlock();

    log_android_debug(debug_str);
}

caffe::Solver<float>* CwrApp::get_solver() {
    return this->solver;
}

void CwrApp::reset_batch() {
    if(this->training_disabled) {
        throw std::runtime_error("Training disabled!");
    }

    this->mtx->lock();
    this->training_batch_features.clear();
    this->training_batch_images.clear();
    this->mtx->unlock();
}

boost::shared_ptr<Net<float>> CwrApp::get_net() {
    return this->net;
}

void CwrApp::cwr_execute_step(int label,
        ModelStatusReporter* reporter,
        TrainingProgressListener* callback) {
    if(this->training_disabled) {
        throw std::runtime_error("Training disabled!");
    }

    this->mtx->lock();
    int train_minibatch_size = this->get_minibatch_size();
    if(this->training_threads > 0) {
        int max_th = min(this->training_threads, train_minibatch_size);
        cout << "Using " << max_th << " training threads" << endl;
        openblas_set_num_threads(max_th);
    }

    cout << "Using " << openblas_get_num_threads() << " training threads" << endl;

    if(this->must_pre_extract_batch_features()) {  // True, as we're running latent replay
        this->cwr_execute_step_with_features(label, reporter, callback);
    } else {
        this->cwr_execute_step_with_images(label, reporter, callback);
    }

    this->mtx->unlock();
}

cv::Mat CwrApp::subtract_mean_image(const cv::Mat &image) {
    return image - this->mean_image;
}

std::vector<float> CwrApp::inference(const cv::Mat &image) {
    return this->inference(image, this->prediction_layer);
}

std::vector<float> CwrApp::inference(const cv::Mat &image,
        const std::string &output_layer_name) {
    mtx->lock();

    if(this->feature_extraction_threads > 0) {
        int max_th = min(this->feature_extraction_threads, this->get_inference_minibatch_size());
        openblas_set_num_threads(max_th);
    }

    feed_image_layer(test_net, "data", {image}, 0, 1);
    test_net->ForwardTo(test_layer_name_to_id.at(output_layer_name));

    boost::shared_ptr<Blob<float>> pred_layer = test_net->blob_by_name(output_layer_name);
    vector<float> result(pred_layer->shape(1));

    int start_offset = pred_layer->offset({0, 0, 0, 0});
    memcpy(result.data(), &pred_layer->cpu_data()[start_offset], pred_layer->shape(1)*sizeof(float));

    mtx->unlock();

    return result;
}

void CwrApp::save_everything_to_disk(const std::string &folder) {
    mtx->lock();
    std::stringstream debug_str;
    debug_str << "save_everything_to_disk:" << endl;

    // CWR status
    std::vector<std::string> cwr_layers_names = cwr.get_cwr_layers_name();
    uint32_t max_classes = cwr.get_max_classes();
    uint32_t current_classes_n = cwr.get_current_classes_n();
    std::vector<float> class_updates = cwr.get_class_updates();

    // App status
    const std::string& solver_path = this->solver_path;
    const std::string& weights_path = this->weights_path;
    const std::vector<string>& labels = this->labels;
    const cv::Mat& mean_image = this->mean_image;
    const std::string& prediction_layer = this->prediction_layer;
    const std::string& feature_extraction_layer = this->feature_extraction_layer;
    const uint32_t epochs = this->epochs;
    const std::string& rehearsal_layer = this->rehearsal_layer;
    const std::string& backward_stop_layer = this->backward_stop_layer;
    const std::shared_ptr<RehearsalMemory<cwra_rehe_t>> rehearsal_memory = this->rehearsal_memory;

    //Serialization
    boost::filesystem::path serialization_folder{folder};
    boost::filesystem::create_directory(serialization_folder);
    boost::filesystem::path serialization_file = serialization_folder / "cwr_app_status";
    std::string path = serialization_file.make_preferred().string();

    debug_str << "saving to: " << path << endl;

    auto out_file = std::fstream(path, std::ios::out | std::ios::binary);

    serialize_string_vector(cwr_layers_names, out_file);
    debug_str << cwr_layers_names << endl;

    serialize_uint32t(current_classes_n, out_file);
    debug_str << current_classes_n << ", ";

    serialize_float_vector(class_updates, out_file);
    debug_str << class_updates << endl;

    serialize_string(solver_path, out_file);
    debug_str << solver_path << ", ";

    serialize_string(weights_path, out_file);
    debug_str << weights_path << endl;

    serialize_string_vector(labels, out_file);
    debug_str << labels << ", " << endl;

    serialize_Mat(mean_image, out_file);
    debug_str << hash_generic((void*) mean_image.data, mean_image.total() * mean_image.elemSize()) << ", " << endl;

    serialize_string(prediction_layer, out_file);
    debug_str << prediction_layer << ", ";

    serialize_string(feature_extraction_layer, out_file);
    debug_str << feature_extraction_layer << ", ";

    serialize_uint32t(epochs, out_file);
    debug_str << epochs << ", ";

    serialize_string(rehearsal_layer, out_file);
    debug_str << rehearsal_layer << ", ";

    serialize_string(backward_stop_layer, out_file);
    debug_str << backward_stop_layer << ", ";

    if(!rehearsal_layer.empty()) {
        rehearsal_memory->save_everything(out_file);
        debug_str << "ReheMem: " << hash_rehearsal_memory(rehearsal_memory->getSamplesX()) << endl;
    }

    out_file.close();

    if(out_file.fail()) {
        debug_str << "Failed to write state to disk" << endl;
    } else {
        this->overwrite_model_on_disk();
    }

    /*log_android_debug( debug_str );

    debug_str << "Params hash (net):" << endl << hash_net_params(this->net) << endl;
    log_android_debug( debug_str );

    debug_str << "Params hash (test_net):" << endl << hash_net_params(this->test_net) << endl;*/

    log_android_debug( debug_str );

    mtx->unlock();
}

void CwrApp::overwrite_model_on_disk() {
    mtx->lock();
    std::string net_path = this->weights_path;
    this->serialize_model_on_disk(net_path);
    mtx->unlock();
}

void CwrApp::serialize_model_on_disk(const std::string &path) {
    mtx->lock();
    //this->test_net->Forward();
    //this->net->Forward();

    //__android_log_print(ANDROID_LOG_DEBUG, "NativePrint", "serialize_model_on_disk net %s", hash_net(this->net).c_str());
    //__android_log_print(ANDROID_LOG_DEBUG, "NativePrint", "serialize_model_on_disk test net %s", hash_net(this->test_net).c_str());

    string net_path = path;
    //string test_net_path = path + "_test";
    NetParameter net_param;
    this->net->ToProto(&net_param, false);
    WriteProtoToBinaryFile(net_param, net_path);

    /*NetParameter test_net_param;
    this->test_net->ToProto(&test_net_param, false);
    WriteProtoToBinaryFile(test_net_param, test_net_path);*/
    mtx->unlock();
}

bool CwrApp::state_exists(const std::string &folder) {
    boost::filesystem::path serialization_folder{folder};
    boost::filesystem::path serialization_file = serialization_folder / "cwr_app_status";
    std::string path = serialization_file.make_preferred().string();

    if (FILE *file = fopen(path.c_str(), "r")) {
        fclose(file);
        return true;
    } else {
        return false;
    }
}

CwrApp CwrApp::read_from_disk(const std::string &folder) {
    std::stringstream debug_str;
    debug_str << "read_from_disk:" << endl;

    // CWR status
    std::vector<std::string> cwr_layers_names;
    uint32_t max_classes;
    uint32_t current_classes_n;
    std::vector<float> class_updates;

    // App status
    std::string solver_path;
    std::string weights_path;
    std::vector<string> labels;
    cv::Mat mean_image;
    std::string prediction_layer;
    std::string feature_extraction_layer;
    uint32_t epochs;
    std::string rehearsal_layer;
    std::string backward_stop_layer;
    std::shared_ptr<RehearsalMemory<cwra_rehe_t>> rehearsal_memory;

    //Serialization
    boost::filesystem::path serialization_folder{folder};
    boost::filesystem::path serialization_file = serialization_folder / "cwr_app_status";
    std::string path = serialization_file.make_preferred().string();

    debug_str << "loading from: " << path << endl;

    auto in_file = std::fstream(path, std::ios::in | std::ios::binary);

    log_android_debug( debug_str );

    cwr_layers_names = deserialize_string_vector(in_file); // OK
    debug_str << cwr_layers_names << endl;

    current_classes_n = deserialize_uint32t(in_file);
    debug_str << current_classes_n << ", ";

    class_updates = deserialize_float_vector(in_file); // OK
    debug_str << class_updates << endl;

    solver_path = deserialize_string(in_file); // OK
    debug_str << solver_path << ", ";

    weights_path = deserialize_string(in_file); // OK
    debug_str << weights_path << endl;

    labels = deserialize_string_vector(in_file); // OK
    debug_str << labels << ", " << endl;

    log_android_debug( debug_str );

    mean_image = deserialize_Mat(in_file); // OK
    debug_str << hash_generic((void*) mean_image.data, mean_image.total() * mean_image.elemSize()) << ", " << endl;

    prediction_layer = deserialize_string(in_file); // OK
    debug_str << prediction_layer << ", ";

    feature_extraction_layer = deserialize_string(in_file); // OK
    debug_str << feature_extraction_layer << ", ";

    epochs = deserialize_uint32t(in_file); // OK
    debug_str << epochs << ", ";

    rehearsal_layer = deserialize_string(in_file);
    debug_str << rehearsal_layer << ", ";

    backward_stop_layer = deserialize_string(in_file);
    debug_str << backward_stop_layer << ", ";

    if(!rehearsal_layer.empty()) {
        rehearsal_memory = load_rehe_from_snapshot(in_file);
        debug_str << "ReheMem: " << hash_rehearsal_memory(rehearsal_memory->getSamplesX()) << endl;
    }

    in_file.close(); // OK

    CwrApp result = CwrApp(solver_path, weights_path, class_updates, epochs,
                  current_classes_n, cwr_layers_names, prediction_layer, labels, mean_image,
                  feature_extraction_layer, backward_stop_layer);
    if(!rehearsal_layer.empty()) {
        result.set_rehearsal_layer(rehearsal_layer);
        result.set_rehearsal_memory(rehearsal_memory);
    }

    /*debug_str << "Net hash: " << hash_net(result.net) << endl;
    debug_str << "Test net hash: " << hash_net(result.test_net) << endl;

    debug_str << "Params hash (net):" << endl << hash_net_params(result.net) << endl;
    log_android_debug( debug_str );

    debug_str << "Params hash (test_net):" << endl << hash_net_params(result.test_net) << endl;*/

    log_android_debug( debug_str );

    return result;
}

void CwrApp::preallocate_memory() {
    bool has_forward_sensitive_params = false;

    for( auto it = this->train_layer_name_to_id.begin(); it != this->train_layer_name_to_id.end(); ++it )
    {
        string layer_name = it->first;
        int layer_id = it->second;
        boost::shared_ptr<caffe::Layer<float>> layer = this->net->layers()[layer_id];
        string layer_type_str(layer->type());

        if(layer_type_str == "BatchReNorm") {
            has_forward_sensitive_params = true;
            break;
        }
    }

    if(!has_forward_sensitive_params) {
        this->net->Forward();
        this->test_net->Forward();
    }
}

bool CwrApp::is_training_disabled() const {
    return training_disabled;
}

const vector<string> &CwrApp::get_labels() const {
    return labels;
}

const Mat &CwrApp::get_mean_image() const {
    return mean_image;
}

const string &CwrApp::get_prediction_layer() const {
    return prediction_layer;
}

int CwrApp::get_epochs() const {
    return epochs;
}

bool CwrApp::is_caffe_initialized() {
    return caffe_initialized;
}

int CwrApp::add_category(const std::string &category_label) {
    mtx->lock();
    //TODO: add checks (max class number)
    this->labels.push_back(category_label);
    int result = this->cwr.increment_class_number();
    mtx->unlock();
    return result;
}

int CwrApp::get_current_categories_count() {
    return this->cwr.get_current_classes_n();
}

int CwrApp::get_max_categories() {
    return this->cwr.get_max_classes();
}

void CwrApp::set_brn_past_weight(float weight) {
    cwr.set_brn_past_weight(weight);
}

bool CwrApp::must_pre_extract_batch_features() {
    return !this->feature_extraction_layer.empty();
}

void CwrApp::cwr_execute_step_with_images(
        int label,
        ModelStatusReporter* reporter,
        TrainingProgressListener* callback) {
    std::stringstream debug_stream;

    std::chrono::high_resolution_clock::time_point t1, t2;
    auto duration = duration_cast<milliseconds>( t2 - t1 ).count(); // Trick used for duration "auto" definition
    int forward_ms = 0, backward_ms = 0, update_ms = 0;

    std::string backward_stop = this->feature_extraction_layer;
    if(!this->backward_stop_layer.empty()) {
        backward_stop = this->backward_stop_layer;
    }

    int train_minibatch_size = this->get_minibatch_size();
    int train_iterations_per_epoch = (int) ceil((double)training_batch_images.size() / (double)train_minibatch_size);
    int train_iterations = train_iterations_per_epoch * this->epochs;
    std::vector<cv::Mat> train_x = this->training_batch_images;

    std::vector<int> train_y(this->training_batch_images.size(), label);
    std::vector<std::vector<float>> target_y = compute_one_hot_vectors(train_y, this->net, this->prediction_layer);
    int iteration;

    pad_data_single(train_x, train_minibatch_size);
    pad_data_single(train_y, train_minibatch_size);
    pad_data_single(target_y, train_minibatch_size);

    std::set<int> unique_y_set(train_y.begin(), train_y.end());

    std::vector<cwra_rehe_t> rehe_x;
    std::vector<int> rehe_y;
    std::vector<std::vector<float>> rehe_t;
    int reha_iters_per_epoch = 0;
    int reha_in_minibatch = 0;

    if(this->rehearsal_memory != nullptr) {
        rehe_x = this->rehearsal_memory->getSamplesX();
        rehe_y = this->rehearsal_memory->getSamplesY();
        rehe_t = compute_one_hot_vectors(rehe_y, this->net, this->prediction_layer);
        cout << "Will use rehearsal memory: " << rehe_x.size() << " patterns";
        unique_y_set.insert(rehe_y.begin(), rehe_y.end());

        reha_in_minibatch = this->get_net()->blob_by_name("data_reha")->shape()[0];
        reha_iters_per_epoch = pad_data_single(rehe_x, reha_in_minibatch);
        pad_data_single(rehe_y, reha_in_minibatch);
        pad_data_single(rehe_t, reha_in_minibatch);
        cout << " (" << rehe_x.size() << " after padding)" << endl;
    }

    std::vector<int> unique_y(unique_y_set.begin(), unique_y_set.end());
    std::sort(unique_y.begin(), unique_y.end());
    std::vector<int> y_freq(unique_y.size(), 0);
    for(auto &train_label : train_y) {
        for(int i = 0; i < unique_y.size(); i++) {
            if(unique_y[i] == train_label) {
                y_freq[i]++;
            }
        }
    }

    for(auto &train_label : rehe_y) {
        for(int i = 0; i < unique_y.size(); i++) {
            if(unique_y[i] == train_label) {
                y_freq[i]++;
            }
        }
    }

    if(reporter) reporter->startChangePhase("Reset weights");
    cwr.reset_weights();
    if(reporter) reporter->onPhaseFinished("Reset weights");
    if(reporter) reporter->startChangePhase("Load weights nic");
    cwr.load_weights_nic(unique_y);
    if(reporter) reporter->onPhaseFinished("Load weights nic");
    //shuffle_in_unison(train_x, train_y, target_y, 0);

    if(rehe_x.size() > 1) {
        //shuffle_in_unison(rehe_x, rehe_t, target_y, 0);
    }

    iteration = 0;

    int backward_stop_layer_id = train_layer_name_to_id.at(backward_stop);

    while (iteration < train_iterations) {
        debug_stream << "Iteration " << iteration << endl;
        log_android_debug(debug_stream);

        if(callback != nullptr) {
            callback->updateProgress((float) iteration / train_iterations);
        }

        int it_mod = iteration % train_iterations_per_epoch;
        int start = it_mod * train_minibatch_size;
        int end = (it_mod + 1) * train_minibatch_size;

        feed_image_layer(this->get_net(), "data", train_x, start, end);
        feed_label_layer(this->get_net(), "label", train_y, start, end);
        feed_target_layer(this->get_net(), "target", target_y, start, end);

        if(!rehe_x.empty()) {
            int it_mod_reha = iteration % reha_iters_per_epoch;
            int reha_start = it_mod_reha * reha_in_minibatch;
            int reha_end = (it_mod_reha + 1) * reha_in_minibatch;

            feed_feature_layer(this->get_net(), "data_reha", rehe_x, reha_start, reha_end);
            feed_label_layer(this->get_net(), "label", rehe_y, reha_start, reha_end, train_minibatch_size);
            feed_target_layer(this->get_net(), "target", rehe_t, reha_start, reha_end, train_minibatch_size);
        }

        /*debug_stream << "Data hash: " << hash_caffe_blob(this->get_net(), "data") << endl;
        if(!rehe_x.empty()) {
            debug_stream << "Rehe data hash: " << hash_caffe_blob(this->get_net(), "data_reha") << endl;
        }
        debug_stream << "Labels hash: " << hash_caffe_blob(this->get_net(), "label") << endl;
        debug_stream << "Target hash: " << hash_caffe_blob(this->get_net(), "target") << endl;*/

        if(reporter) reporter->startChangePhase("Forward");
        t1 = std::chrono::high_resolution_clock::now();
        net->ClearParamDiffs();
        net->Forward();
        t2 = std::chrono::high_resolution_clock::now();
        duration = duration_cast<milliseconds>( t2 - t1 ).count();
        forward_ms += duration;
        if(reporter) reporter->onPhaseFinished("Forward");

        if(reporter) reporter->startChangePhase("Backward");
        t1 = std::chrono::high_resolution_clock::now();
        net->BackwardTo(backward_stop_layer_id);
        t2 = high_resolution_clock::now();
        duration = duration_cast<milliseconds>( t2 - t1 ).count();
        backward_ms += duration;
        if(reporter) reporter->onPhaseFinished("Backward");

        if(reporter) reporter->startChangePhase("ApplyUpdate");
        t1 = std::chrono::high_resolution_clock::now();
        solver->ApplyUpdate();
        t2 = high_resolution_clock::now();
        duration = duration_cast<milliseconds>( t2 - t1 ).count();
        update_ms += duration;
        if(reporter) reporter->onPhaseFinished("ApplyUpdate");

        if (iteration == (train_iterations - 1)) {
            if(reporter) reporter->startChangePhase("Cwr consolidation");
            cout << unique_y << endl;
            cout << y_freq << endl;

            if(reporter) reporter->startChangePhase("Consolidate weights CWR+");
            cwr.consolidate_weights_cwr_plus(unique_y, y_freq);
            if(reporter) reporter->onPhaseFinished("Consolidate weights CWR+");

            // increase weights of trained classes
            vector<float> class_updates = cwr.get_class_updates();

            for(int i = 0; i < unique_y.size(); i++) {
                class_updates[unique_y[i]] = class_updates[unique_y[i]] + y_freq[i];
            }

            cwr.set_class_updates(class_updates);

            if(reporter) reporter->startChangePhase("Load weights");
            cwr.load_weights();
            if(reporter) reporter->onPhaseFinished("Load weights");
            if(reporter) reporter->onPhaseFinished("Cwr consolidation");
        }

        iteration++;
    }

    debug_stream << "Forward time: " << forward_ms << endl
                 << "Backward time: " << backward_ms << endl
                 << "Update time: " << update_ms << endl;

    if(this->rehearsal_memory != nullptr) {
        debug_stream << "Updating rehearsal memory" << endl;

        /* This is the non-balanced version in which 10 randomly chosen instances are
         * inserted in the replay buffer by removing 10 randomly chosen replay instances.

        vector<int> to_be_added_idx = indexes_range(0, train_x.size()-1, 10);
        debug_stream << "Using patterns at indexes: " << to_be_added_idx << endl;
        vector<cwra_rehe_t> to_be_added_x;
        vector<int> to_be_added_y;
        for(int i = 0; i < to_be_added_idx.size(); i++) {
            vector<float> extracted_features = this->inference(train_x[to_be_added_idx[i]], this->rehearsal_layer);

            to_be_added_x.push_back(std::make_shared<vector<float>>(extracted_features));
            to_be_added_y.push_back(train_y[to_be_added_idx[i]]);
        }
        this->rehearsal_memory->update_memory(to_be_added_x, to_be_added_y, 1, 10);
         */

        /* This is the balanced version in which all instances are considered for insertion.
         * The replay buffer will be class-balanced.
         *
         * The selection is based on random choice over the union of old+new instances.
         */
        debug_stream << "Using all patterns for balanced replay: " << endl;
        vector<cwra_rehe_t> to_be_added_x;
        vector<int> to_be_added_y;
        for(int i = 0; i < train_x.size(); i++) {
            vector<float> extracted_features = this->inference(train_x[i], this->rehearsal_layer);

            to_be_added_x.push_back(std::make_shared<vector<float>>(extracted_features));
            to_be_added_y.push_back(train_y[i]);
        }
        this->rehearsal_memory->update_memory_balanced(to_be_added_x, to_be_added_y, 1, -1);
    }

    if(callback != nullptr) {
        callback->updateProgress(1.0f);
    }

    log_android_debug(debug_stream);
}

void CwrApp::cwr_execute_step_with_features(
        int label,
        ModelStatusReporter* reporter,
        TrainingProgressListener* callback) {
    // Define variables used for timing and debugging
    std::stringstream debug_stream;

    std::chrono::high_resolution_clock::time_point t1, t2;
    auto duration = duration_cast<milliseconds>( t2 - t1 ).count(); // Trick used for duration "auto" definition
    int preparation_ms = 0, data_feed_ms = 0, forward_ms = 0, backward_ms = 0, update_ms = 0, consolidation_ms = 0;

    t1 = std::chrono::high_resolution_clock::now();

    // Get the backward stop layer
    // We want to keep the part of the model before the latent replay layer frozen, which means
    // that we can stop the backward pass early.
    std::string backward_stop = this->feature_extraction_layer;
    if(!this->backward_stop_layer.empty()) {
        backward_stop = this->backward_stop_layer;
    }

    // Compute the amount of train iterations
    int train_minibatch_size = this->get_features_train_minibatch_size();
    int train_iterations_per_epoch = (int) ceil((double)training_batch_features.size() / (double)train_minibatch_size);
    int train_iterations = train_iterations_per_epoch * this->epochs;
    std::vector<std::shared_ptr<std::vector<float>>> train_x = this->training_batch_features;

    // Compute one-hot vectors
    std::vector<int> train_y(this->training_batch_features.size(), label);
    std::vector<std::vector<float>> target_y = compute_one_hot_vectors(train_y, this->net, this->prediction_layer);
    int iteration;

    // If (n_train_instances % mb_size) != 0, then pad the batch by duplicating some instances
    pad_data_single(train_x, train_minibatch_size);
    pad_data_single(train_y, train_minibatch_size);
    pad_data_single(target_y, train_minibatch_size);

    std::set<int> unique_y_set(train_y.begin(), train_y.end());

    // Manage the replay buffer
    std::vector<cwra_rehe_t> rehe_x;
    std::vector<int> rehe_y;
    std::vector<std::vector<float>> rehe_t;
    int reha_iters_per_epoch = 0;
    int reha_in_minibatch = 0;

    if(this->rehearsal_memory != nullptr) {
        rehe_x = this->rehearsal_memory->getSamplesX();
        rehe_y = this->rehearsal_memory->getSamplesY();
        rehe_t = compute_one_hot_vectors(rehe_y, this->net, this->prediction_layer);
        cout << "Will use rehearsal memory: " << rehe_x.size() << " patterns";
        unique_y_set.insert(rehe_y.begin(), rehe_y.end());

        reha_in_minibatch = this->get_net()->blob_by_name("data_reha")->shape()[0];
        reha_iters_per_epoch = pad_data_single(rehe_x, reha_in_minibatch);
        pad_data_single(rehe_y, reha_in_minibatch);
        pad_data_single(rehe_t, reha_in_minibatch);
        cout << " (" << rehe_x.size() << " after padding)" << endl;
    }

    // Compute the amount of instances per category (to be used in the CWR consolidation procedure)
    std::vector<int> unique_y(unique_y_set.begin(), unique_y_set.end());
    std::sort(unique_y.begin(), unique_y.end());
    std::vector<int> y_freq(unique_y.size(), 0);
    for(auto &train_label : train_y) {
        for(int i = 0; i < unique_y.size(); i++) {
            if(unique_y[i] == train_label) {
                 y_freq[i]++;
            }
        }
    }

    for(auto &train_label : rehe_y) {
        for(int i = 0; i < unique_y.size(); i++) {
            if(unique_y[i] == train_label) {
                y_freq[i]++;
            }
        }
    }

    // CWR initialization steps
    debug_stream << "Resetting weights" << endl;
    log_android_debug(debug_stream);

    if(reporter) reporter->startChangePhase("Reset weights");
    cwr.reset_weights();
    if(reporter) reporter->onPhaseFinished("Reset weights");

    debug_stream << "Loading nic ones" << endl;
    log_android_debug(debug_stream);
    if(reporter) reporter->startChangePhase("Load weights nic");
    cwr.load_weights_nic(unique_y);
    if(reporter) reporter->onPhaseFinished("Load weights nic");

    //shuffle_in_unison(train_x, train_y, target_y, 0);
    if(rehe_x.size() > 1) {
        //shuffle_in_unison(rehe_x, rehe_t, target_y, 0);
    }

    iteration = 0;

    // Get the ID of the feature extraction layer
    debug_stream << "Getting id of feature extraction layer " << this->feature_extraction_layer << endl;
    log_android_debug(debug_stream);
    int features_layer_id = train_layer_name_to_id.at(this->feature_extraction_layer);

    // Get the ID of the backward stop layer
    debug_stream << "Getting id of backward stop layer  " << backward_stop << endl;
    log_android_debug(debug_stream);
    int backward_stop_layer_id = train_layer_name_to_id.at(backward_stop);
    int backward_stop_layer_param_id = train_layer_names_to_param_id.at(backward_stop); // Used for partial apply update

    t2 = std::chrono::high_resolution_clock::now();
    duration = duration_cast<milliseconds>( t2 - t1 ).count();
    preparation_ms += duration;

    // Training loop
    while (iteration < train_iterations) {
        debug_stream << "Iteration " << iteration << endl;
        log_android_debug(debug_stream);

        if(callback != nullptr) {
            callback->updateProgress((float) iteration / train_iterations);
        }

        int it_mod = iteration % train_iterations_per_epoch;
        int start = it_mod * train_minibatch_size;
        int end = (it_mod + 1) * train_minibatch_size;

        t1 = std::chrono::high_resolution_clock::now();

        feed_feature_layer(this->get_net(), this->feature_extraction_layer, train_x, start, end);
        feed_label_layer(this->get_net(), "label", train_y, start, end);
        feed_target_layer(this->get_net(), "target", target_y, start, end);

        if(!rehe_x.empty()) {
            int it_mod_reha = iteration % reha_iters_per_epoch;
            int reha_start = it_mod_reha * reha_in_minibatch;
            int reha_end = (it_mod_reha + 1) * reha_in_minibatch;

            feed_feature_layer(this->get_net(), "data_reha", rehe_x, reha_start, reha_end);
            feed_label_layer(this->get_net(), "label", rehe_y, reha_start, reha_end, train_minibatch_size);
            feed_target_layer(this->get_net(), "target", rehe_t, reha_start, reha_end, train_minibatch_size);
        }

        t2 = std::chrono::high_resolution_clock::now();
        duration = duration_cast<milliseconds>( t2 - t1 ).count();
        data_feed_ms += duration;

        if(reporter) reporter->startChangePhase("Forward");
        t1 = std::chrono::high_resolution_clock::now();
        net->ClearParamDiffs();
        net->ForwardFrom(features_layer_id);
        t2 = std::chrono::high_resolution_clock::now();
        duration = duration_cast<milliseconds>( t2 - t1 ).count();
        forward_ms += duration;
        if(reporter) reporter->onPhaseFinished("Forward");

        if(reporter) reporter->startChangePhase("Backward");
        t1 = std::chrono::high_resolution_clock::now();
        net->BackwardTo(backward_stop_layer_id);
        t2 = high_resolution_clock::now();
        duration = duration_cast<milliseconds>( t2 - t1 ).count();
        backward_ms += duration;
        if(reporter) reporter->onPhaseFinished("Backward");

        if(reporter) reporter->startChangePhase("ApplyUpdate");
        t1 = std::chrono::high_resolution_clock::now();
        //solver->ApplyUpdate();
        partial_apply_update(dynamic_cast<caffe::SGDSolver<float>*>(solver), backward_stop_layer_param_id, -1);
        t2 = high_resolution_clock::now();
        duration = duration_cast<milliseconds>( t2 - t1 ).count();
        update_ms += duration;
        if(reporter) reporter->onPhaseFinished("ApplyUpdate");

        // CWR consolidation step
        if (iteration == (train_iterations - 1)) {
            t1 = std::chrono::high_resolution_clock::now();

            if(reporter) reporter->startChangePhase("Cwr consolidation");
            cout << unique_y << endl;
            cout << y_freq << endl;

            if(reporter) reporter->startChangePhase("Consolidate weights CWR+");
            cwr.consolidate_weights_cwr_plus(unique_y, y_freq);
            if(reporter) reporter->onPhaseFinished("Consolidate weights CWR+");

            // increase weights of trained classes
            vector<float> class_updates = cwr.get_class_updates();

            for(int i = 0; i < unique_y.size(); i++) {
                class_updates[unique_y[i]] = class_updates[unique_y[i]] + y_freq[i];
            }

            cwr.set_class_updates(class_updates);

            if(reporter) reporter->startChangePhase("Load weights");
            cwr.load_weights(); // This prepares the net for inference
            if(reporter) reporter->onPhaseFinished("Load weights");
            if(reporter) reporter->onPhaseFinished("Cwr consolidation");

            t2 = high_resolution_clock::now();
            duration = duration_cast<milliseconds>( t2 - t1 ).count();
            consolidation_ms += duration;
        }

        iteration++;
    }

    debug_stream << "Preparation: " << preparation_ms << endl
                 << "Data feed: " << data_feed_ms << endl
                 << "Forward time: " << forward_ms << endl
                 << "Backward time: " << backward_ms << endl
                 << "Update time: " << update_ms << endl
                 << "Consolidation: " << consolidation_ms << endl;

    // Update the replay buffer
    if(this->rehearsal_memory != nullptr) {
        debug_stream << "Updating rehearsal memory" << endl;

        /* This is the non-balanced version in which 10 randomly chosen instances are
         * inserted in the replay buffer by removing 10 randomly chosen replay instances.

        vector<int> to_be_added_idx = indexes_range(0, train_x.size()-1, 10);
        debug_stream << "Using patterns at indexes: " << to_be_added_idx << endl;
        vector<cwra_rehe_t> to_be_added_x;
        vector<int> to_be_added_y;
        for(int i = 0; i < to_be_added_idx.size(); i++) {
            to_be_added_x.push_back(train_x[to_be_added_idx[i]]);
            to_be_added_y.push_back(train_y[to_be_added_idx[i]]);
        }
        this->rehearsal_memory->update_memory_balanced(to_be_added_x, to_be_added_y, 1, 10);
        */

        /* This is the balanced version in which all instances are considered for insertion.
         * The replay buffer will be class-balanced.
         *
         * The selection is based on random choice over the union of old+new instances.
         */
        debug_stream << "Using all patterns for balanced replay: " << endl;
        vector<cwra_rehe_t> to_be_added_x;
        vector<int> to_be_added_y;
        for(int i = 0; i < train_x.size(); i++) {
            to_be_added_x.push_back(train_x[i]);
            to_be_added_y.push_back(train_y[i]);
        }
        this->rehearsal_memory->update_memory_balanced(to_be_added_x, to_be_added_y, 1, -1);
    }

    if(callback != nullptr) {
        callback->updateProgress(1.0f);
    }

    log_android_debug(debug_stream);
}

void CwrApp::set_training_threads(int threads) {
    this->training_threads = threads;
}

void CwrApp::set_feature_extraction_threads(int threads) {
    this->feature_extraction_threads = threads;
}

int CwrApp::get_inference_minibatch_size() {
    return this->inference_minibatch_size;
}

void CwrApp::set_rehearsal_memory(const std::shared_ptr<RehearsalMemory<cwra_rehe_t>> &new_rehearsal_memory) {
    this->rehearsal_memory = new_rehearsal_memory;
}

std::shared_ptr<RehearsalMemory<cwra_rehe_t>> CwrApp::get_rehearsal_memory() {
    return this->rehearsal_memory;
}

void CwrApp::set_backward_stop_layer(const std::string &new_backward_stop_layer) {
    this->backward_stop_layer = new_backward_stop_layer;
}

void CwrApp::set_rehearsal_layer(const std::string &new_rehearsal_layer) {
    this->rehearsal_layer = new_rehearsal_layer;
}

std::string CwrApp::get_backward_stop_layer() {
    return this->backward_stop_layer;
}

std::string CwrApp::get_rehearsal_layer() {
    return this->rehearsal_layer;
}

boost::shared_ptr<caffe::Net<float>> CwrApp::get_test_net() {
    return this->test_net;
}

int CwrApp::get_features_train_minibatch_size() {
    if(this->feature_extraction_layer.empty()) {
        return 0;
    }

    return this->net->blob_by_name(this->feature_extraction_layer)->shape()[0];
}

void serialize_string(const std::string &str,
                             std::basic_fstream<char, std::char_traits<char>> &out_file) {
    serialize_uint32t(str.length(), out_file);
    out_file.write(str.c_str(), str.length());
}

std::string deserialize_string(std::basic_fstream<char, std::char_traits<char>> &in_file) {
    uint32_t str_len = deserialize_uint32t(in_file);;

    std::vector<char> buffer(str_len);

    in_file.read(&buffer[0], str_len);
    return std::string(&buffer[0], str_len);
}

void serialize_float_vector(const std::vector<float> &float_vec,
                             std::basic_fstream<char, std::char_traits<char>> &out_file) {
    serialize_uint32t(float_vec.size(), out_file);
    for(auto& fl : float_vec) {
        out_file.write((const char*) &fl, sizeof(float));
    }
}

std::vector<float> deserialize_float_vector(std::basic_fstream<char, std::char_traits<char>> &in_file) {
    uint32_t result_size = deserialize_uint32t(in_file);
    std::vector<float> result;
    std::vector<char> buffer(sizeof(float));

    for(uint32_t i = 0; i < result_size; i++) {
        in_file.read(&buffer[0], sizeof(float));

        result.push_back(*(float*)(&buffer[0]));
    }

    return result;
}

void serialize_rehearsal_blob(const std::vector<cwra_rehe_t> &rehe_vec,
                              std::basic_fstream<char, std::char_traits<char>> &out_file) {
    serialize_uint32t(rehe_vec.size(), out_file);
    for(auto& r_vec : rehe_vec) {
        serialize_float_vector(*r_vec, out_file);
    }
}

std::vector<cwra_rehe_t> deserialize_rehearsal_blob(std::basic_fstream<char, std::char_traits<char>> &in_file) {
    uint32_t result_size = deserialize_uint32t(in_file);
    std::vector<cwra_rehe_t> result;

    for(uint32_t i = 0; i < result_size; i++) {
        result.push_back(std::make_shared<vector<float>>(deserialize_float_vector(in_file)));
    }

    return result;
}

void serialize_int_vector(const std::vector<int> &int_vec,
                            std::basic_fstream<char, std::char_traits<char>> &out_file) {
    serialize_uint32t(int_vec.size(), out_file);
    for(auto& fl : int_vec) {
        int64_t int64val = fl;
        out_file.write((const char*) &int64val, sizeof(int64_t));
    }
}

std::vector<int> deserialize_int_vector(std::basic_fstream<char, std::char_traits<char>> &in_file) {
    uint32_t result_size = deserialize_uint32t(in_file);
    std::vector<int> result;
    std::vector<char> buffer(sizeof(int64_t));

    for(uint32_t i = 0; i < result_size; i++) {
        in_file.read(&buffer[0], sizeof(int64_t));
        result.push_back((int) *(int64_t*)(&buffer[0]));
    }

    return result;
}

void serialize_string_vector(const std::vector<std::string> &str_vec,
                             std::basic_fstream<char, std::char_traits<char>> &out_file) {
    serialize_uint32t(str_vec.size(), out_file);
    for(auto& str : str_vec) {
        serialize_string(str, out_file);
    }
}

std::vector<std::string> deserialize_string_vector(std::basic_fstream<char, std::char_traits<char>> &in_file) {
    uint32_t result_size = deserialize_uint32t(in_file);
    std::vector<std::string> result;

    for(uint32_t i = 0; i < result_size; i++) {
        result.push_back(deserialize_string(in_file));
    }

    return result;
}

void serialize_Mat(const cv::Mat &image,
                             std::basic_fstream<char, std::char_traits<char>> &out_file) {
    // Implicit: CV_32FC3
    serialize_uint32t(image.cols, out_file);
    serialize_uint32t(image.rows, out_file);
    serialize_uint32t(image.total() * image.elemSize(), out_file);
    out_file.write((const char*) image.data, image.total() * image.elemSize());
}

cv::Mat deserialize_Mat(std::basic_fstream<char, std::char_traits<char>> &in_file) {
    uint32_t cols, rows, byte_size;
    // Implicit: CV_32FC3

    cols = deserialize_uint32t(in_file);
    rows = deserialize_uint32t(in_file);
    byte_size = deserialize_uint32t(in_file);
    std::vector<char> buffer;
    buffer.reserve(byte_size);
    in_file.read(&buffer[0], byte_size);
    cv::Mat result(rows, cols, CV_32FC3);
    memcpy(result.data, &buffer[0], byte_size);
    return result;
}

void serialize_uint32t(uint32_t value, std::basic_fstream<char, std::char_traits<char>> &out_file) {
    char* value_ptr = (char*) &value;
    for(int i = 0; i < 4; i++) {
        out_file.write(&value_ptr[i], 1);
    }
}

uint32_t deserialize_uint32t(std::basic_fstream<char, std::char_traits<char>> &out_file) {
    uint32_t result = 0;
    char* value_ptr = (char*) &result;
    for(int i = 0; i < 4; i++) {
        out_file.read(&value_ptr[i], 1);
    }

    return result;
}

std::shared_ptr<RehearsalMemory<cwra_rehe_t>> load_rehe_from_snapshot(std::basic_fstream<char, std::char_traits<char>> &in_stream) {
    std::stringstream debug_str;
    debug_str << "rehe memory, load_everything:" << endl;

    uint32_t current_num_patterns;
    uint32_t max_num_patterns;
    vector<cwra_rehe_t> external_memory_x;
    vector<int> external_memory_y;

    current_num_patterns = deserialize_uint32t(in_stream);
    debug_str << current_num_patterns << ", ";

    max_num_patterns = deserialize_uint32t(in_stream);
    debug_str << max_num_patterns << ", ";

    external_memory_x = deserialize_rehearsal_blob(in_stream);
    debug_str << "External memory x hash: " << hash_rehearsal_memory(external_memory_x) << "," << endl;

    external_memory_y = deserialize_int_vector(in_stream);
    debug_str << "External memory y hash: " << hash_vector(external_memory_y) << "," << endl;

    std::shared_ptr<RehearsalMemory<cwra_rehe_t>> result =
            std::make_shared<RehearsalMemory<cwra_rehe_t>>(max_num_patterns);
    result->load_memory(external_memory_x, external_memory_y);

    log_android_debug(debug_str);
    return result;
}


// World Reference Frame Notes

// Structure:
// left handed
// +y is up
// -x is right
// -z is forward
// it is using column major storage
// distance unit is meters

// OpenVR:
// right handed
// +y is up
// +x is right
// -z is forward
// row major storage
// distance unit is meters
#pragma once

#include <utility>

#include <Eigen/Dense>
#include <ST/XRTypes.h>
#include <openvr_driver.h>

#include "math/Types.hpp"

namespace northstar {
    namespace math {
        class IWorldAdapter {
        public:
            virtual types::SPose FromStructureSensorPoseToOpenVRSpace(const ST::XRPose& Pose) = 0;
            virtual types::Vector3d FromStructureSensorLinearVectorArrayToOpenVRSpace(const std::array<double, 3>& Vector) = 0;
            virtual types::Vector3d FromStructureSensorAngularVectorArrayToOpenVRSpace(const std::array<double, 3>& Vector) = 0;
            virtual types::AffineMatrix4d ConversionMatrixFromLeapMotionTrackingSpaceToHMDRelativeSpace(const types::Vector3d& v3dHeadRelativeLeapPosition, const types::Quaterniond& qdHeadRelativeLeapOrientation) = 0;
            virtual types::AffineMatrix4d ConversionMatrixFromHMDSpaceToOpenVRWorldSpace(const vr::DriverPose_t& sOVRPose) = 0;
            virtual types::Vector3d FromLeapMotionVelocityToOpenVRVelocity(const types::AffineMatrix4d& m4dConversionMatrix, const types::Vector3d& v3dLeapVelocity) = 0;
            virtual types::Quaterniond FromUnityQuaternionToOpenVRQuaternion(const types::Quaterniond& qdUnityRotation) = 0;
            virtual types::Vector3d FromUnityPositionToOpenVRPosition(const types::Vector3d& v3dUnityPosition) = 0;
            virtual types::Vector4d FromUnityProjectionExtentsLRTBToOpenVRProjectionExtentsLRTB(const types::Vector4d& v4dUnityProjectionExtentsLRTB) = 0;
            virtual types::AffineMatrix4d FromUnityMatrix4dToOpenVRMatrix4d(const types::AffineMatrix4d& v4dUnityPose) = 0;
            virtual types::Vector2d FromOpenGLUVToUnityUV(const types::Vector2d& v2dOpenGLUV) = 0;
            virtual types::Vector2d FromUnityUVToOpenGLUV(const types::Vector2d& v2dUnityUV) = 0;
        };
    }
}


/*
 Copyright (C) 2016-2020 Quaternion Risk Management Ltd
 All rights reserved.

 This file is part of ORE, a free-software/open-source library
 for transparent pricing and risk analysis - http://opensourcerisk.org

 ORE is free software: you can redistribute it and/or modify it
 under the terms of the Modified BSD License.  You should have received a
 copy of the license along with this program.
 The license is also available online at <http://opensourcerisk.org>

 This program is distributed on the basis that it will form a useful
 contribution to risk analytics and model standardisation, but WITHOUT
 ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 FITNESS FOR A PARTICULAR PURPOSE. See the license for more details.
*/

#include <orea/aggregation/postprocess.hpp>
#include <orea/aggregation/exposureallocator.hpp>
#include <orea/aggregation/dimcalculator.hpp>
#include <orea/aggregation/dimregressioncalculator.hpp>
#include <orea/aggregation/dynamiccreditxvacalculator.hpp>
#include <orea/aggregation/xvacalculator.hpp>
#include <orea/aggregation/staticcreditxvacalculator.hpp>
#include <orea/aggregation/cvaspreadsensitivitycalculator.hpp>
#include <ored/utilities/log.hpp>
#include <ored/utilities/vectorutils.hpp>
#include <ql/errors.hpp>
#include <ql/time/calendars/weekendsonly.hpp>
#include <ql/version.hpp>

#include <ql/math/distributions/normaldistribution.hpp>
#include <ql/math/generallinearleastsquares.hpp>
#include <ql/math/kernelfunctions.hpp>
#include <ql/methods/montecarlo/lsmbasissystem.hpp>
#include <ql/time/daycounters/actualactual.hpp>

#include <qle/math/nadarayawatson.hpp>
#include <qle/math/stabilisedglls.hpp>

#include <boost/accumulators/accumulators.hpp>
#include <boost/accumulators/statistics/error_of_mean.hpp>
#include <boost/accumulators/statistics/mean.hpp>
#include <boost/accumulators/statistics/stats.hpp>

using namespace std;
using namespace QuantLib;

using namespace boost::accumulators;

namespace ore {
namespace analytics {

PostProcess::PostProcess(
    const boost::shared_ptr<Portfolio>& portfolio, const boost::shared_ptr<NettingSetManager>& nettingSetManager,
    const boost::shared_ptr<Market>& market, const std::string& configuration, const boost::shared_ptr<NPVCube>& cube,
    const boost::shared_ptr<AggregationScenarioData>& scenarioData, const map<string, bool>& analytics,
    const string& baseCurrency, const string& allocMethod, Real marginalAllocationLimit, Real quantile,
    const string& calculationType, const string& dvaName, const string& fvaBorrowingCurve,
    const string& fvaLendingCurve,const boost::shared_ptr<DynamicInitialMarginCalculator>& dimCalculator,
    const boost::shared_ptr<CubeInterpretation>& cubeInterpretation, bool fullInitialCollateralisation,
    vector<Period> cvaSensiGrid, Real cvaSensiShiftSize,
    Real kvaCapitalDiscountRate, Real kvaAlpha, Real kvaRegAdjustment, Real kvaCapitalHurdle, Real kvaOurPdFloor,
    Real kvaTheirPdFloor, Real kvaOurCvaRiskWeight, Real kvaTheirCvaRiskWeight, const boost::shared_ptr<NPVCube>& cptyCube,
    const string& flipViewBorrowingCurvePostfix, const string& flipViewLendingCurvePostfix)
    : portfolio_(portfolio), nettingSetManager_(nettingSetManager), market_(market), configuration_(configuration),
      cube_(cube), cptyCube_(cptyCube), scenarioData_(scenarioData), analytics_(analytics), baseCurrency_(baseCurrency), quantile_(quantile),
      calcType_(parseCollateralCalculationType(calculationType)), dvaName_(dvaName),
      fvaBorrowingCurve_(fvaBorrowingCurve), fvaLendingCurve_(fvaLendingCurve), dimCalculator_(dimCalculator),
      cubeInterpretation_(cubeInterpretation), fullInitialCollateralisation_(fullInitialCollateralisation),
      cvaSpreadSensiGrid_(cvaSensiGrid), cvaSpreadSensiShiftSize_(cvaSensiShiftSize),
      kvaCapitalDiscountRate_(kvaCapitalDiscountRate), kvaAlpha_(kvaAlpha), kvaRegAdjustment_(kvaRegAdjustment),
      kvaCapitalHurdle_(kvaCapitalHurdle), kvaOurPdFloor_(kvaOurPdFloor), kvaTheirPdFloor_(kvaTheirPdFloor),
      kvaOurCvaRiskWeight_(kvaOurCvaRiskWeight), kvaTheirCvaRiskWeight_(kvaTheirCvaRiskWeight) {

    // set a default value for the cube interpretation object if it is NULL
    if (!cubeInterpretation_) {
        WLOG("cube interpretation is not set, use regular");
        cubeInterpretation_ = boost::make_shared<RegularCubeInterpretation>();
    }
    boost::shared_ptr<RegularCubeInterpretation> regularCubeInterpretation =
        boost::dynamic_pointer_cast<RegularCubeInterpretation>(cubeInterpretation_);
    bool isRegularCubeStorage = (regularCubeInterpretation != NULL);

    LOG("cube storage is regular: " << isRegularCubeStorage);
    LOG("cube dates: " << cube->dates().size());

    QL_REQUIRE(marginalAllocationLimit > 0.0, "positive allocationLimit expected");

    // check portfolio and cube have the same trade ids, in the same order
    QL_REQUIRE(portfolio->size() == cube_->ids().size(),
               "PostProcess::PostProcess(): portfolio size ("
                   << portfolio->size() << ") does not match cube trade size (" << cube_->ids().size() << ")");
    for (Size i = 0; i < portfolio->size(); ++i) {
        QL_REQUIRE(portfolio->trades()[i]->id() == cube_->ids()[i], "PostProcess::PostProcess(): portfolio trade #"
                                                                        << i << " (id=" << portfolio->trades()[i]->id()
                                                                        << ") does not match cube trade id ("
                                                                        << cube_->ids()[i]);
    }

    if (analytics_["dynamicCredit"]) {
        QL_REQUIRE(cptyCube_, "cptyCube cannot be null when dynamicCredit is ON");
        // check portfolio and cptyCube have the same counterparties, in the same order
        QL_REQUIRE(portfolio->counterparties().size() + 1 == cptyCube_->ids().size(),
                   "PostProcess::PostProcess(): portfolio counterparty size ("
                   << portfolio->counterparties().size() << ") does not match cpty cube trade size ("
                   << cptyCube_->ids().size() << ")");
        for (Size i = 0; i < portfolio->counterparties().size(); ++i) {
            QL_REQUIRE(portfolio->counterparties()[i] == cptyCube_->ids()[i],
                       "PostProcess::PostProcess(): portfolio counterparty #"
                       << i << " (id=" << portfolio->counterparties()[i]
                       << ") does not match cube name id ("
                       << cptyCube_->ids()[i]);
        }
        QL_REQUIRE(dvaName == cptyCube_->ids().back(),
                       "PostProcess::PostProcess(): dvaName (" << dvaName
                       << ") does not match cube name id ("
                       << cptyCube_->ids().back());
    }

    ExposureAllocator::AllocationMethod allocationMethod = parseAllocationMethod(allocMethod);

    /***********************************************
     * Step 0: Netting as of today
     * a) Compute the netting set NPV as of today
     * b) Find the final maturity of the netting set
     */
    LOG("Compute netting set NPVs as of today and netting set maturity");
    // Don't use Settings::instance().evaluationDate() here, this has moved to simulation end date.
    Date today = market->asofDate();
    LOG("AsOfDate = " << QuantLib::io::iso_date(today));

    /***************************************************************
     * Step 1: Dynamic Initial Margin calculation
     * Fills DIM cube per netting set that can be
     * - returned to be further analysed
     * - used in collateral calculation
     * - used in MVA calculation
     */
    if (analytics_["dim"] || analytics_["mva"]) {
        QL_REQUIRE(dimCalculator_, "DIM calculator not set");
        dimCalculator_->build();
    }

    /************************************************************
     * Step 2: Trade Exposure and Netting
     * a) Aggregation across scenarios per trade and date
     *    This yields single trade exposure profiles, EPE and ENE
     * b) Aggregation of NPVs within netting sets per date
     *    and scenario. This prepares the netting set exposure
     *    calculation below
     */
    exposureCalculator_ =
        boost::make_shared<ExposureCalculator>(
            portfolio, cube_, cubeInterpretation_,
            market_, analytics_["exerciseNextBreak"], baseCurrency_, configuration_,
            quantile_, calcType_, analytics_["dynamicCredit"], analytics_["flipViewXVA"]
        );
    exposureCalculator_->build();

    /******************************************************************
     * Step 3: Netting set exposure and allocation to trades
     *
     * a) Compute all netting set exposure profiles EPE and ENE using
     *    collateral if CSAs are given and active.
     * b) Compute the expected collateral balance for each netting set.
     * c) Allocate each netting set's exposure profile to the trade
     *    level such that the trade exposures add up to the netting
     *    set exposure.
     *    Reference:
     *    Michael Pykhtin & Dan Rosen, Pricing Counterparty Risk
     *    at the Trade Level and CVA Allocations, October 2010
     */
    nettedExposureCalculator_ =
        boost::make_shared<NettedExposureCalculator>(
            portfolio_, market_, cube_, baseCurrency, configuration_, quantile_,
            calcType_, analytics_["dynamicCredit"], nettingSetManager_,
	        exposureCalculator_->nettingSetDefaultValue(),
	        exposureCalculator_->nettingSetCloseOutValue(),
            scenarioData_, cubeInterpretation_, analytics_["dim"],
            dimCalculator_, fullInitialCollateralisation_,
            allocationMethod == ExposureAllocator::AllocationMethod::Marginal, marginalAllocationLimit,
            exposureCalculator_->exposureCube(), ExposureCalculator::allocatedEPE, ExposureCalculator::allocatedENE,
            analytics_["flipViewXVA"]
        );
    nettedExposureCalculator_->build();

    /********************************************************
     * Update Stand Alone XVAs
     * needed for some of the simple allocation methods below
     */
    if (analytics_["dynamicCredit"]) {
        cvaCalculator_ = boost::make_shared<DynamicCreditXvaCalculator>(
            portfolio_, market_, configuration_,baseCurrency_, dvaName_,
            fvaBorrowingCurve_, fvaLendingCurve_, analytics_["dim"],
            dimCalculator, exposureCalculator_->exposureCube(),
            nettedExposureCalculator_->exposureCube(), cptyCube_,
            ExposureCalculator::ExposureIndex::EPE,
            ExposureCalculator::ExposureIndex::ENE,
            NettedExposureCalculator::ExposureIndex::EPE,
            NettedExposureCalculator::ExposureIndex::ENE, 0, analytics_["flipViewXVA"],
            flipViewBorrowingCurvePostfix, flipViewLendingCurvePostfix);
    } else {
        cvaCalculator_ = boost::make_shared<StaticCreditXvaCalculator>(
            portfolio_, market_, configuration_,baseCurrency_, dvaName_,
            fvaBorrowingCurve_, fvaLendingCurve_, analytics_["dim"],
            dimCalculator, exposureCalculator_->exposureCube(),
            nettedExposureCalculator_->exposureCube(),
            ExposureCalculator::ExposureIndex::EPE,
            ExposureCalculator::ExposureIndex::ENE,
            NettedExposureCalculator::ExposureIndex::EPE,
            NettedExposureCalculator::ExposureIndex::ENE, analytics_["flipViewXVA"],
            flipViewBorrowingCurvePostfix, flipViewLendingCurvePostfix);
    }
    cvaCalculator_->build();

    /***************************
     * Simple allocation methods
     */
    boost::shared_ptr<ExposureAllocator> exposureAllocator;
    if (allocationMethod == ExposureAllocator::AllocationMethod::Marginal) {
        DLOG("Marginal Calculation handled in NettedExposureCalculator");
    }
    else if (allocationMethod == ExposureAllocator::AllocationMethod::RelativeFairValueNet)
        exposureAllocator = boost::make_shared<RelativeFairValueNetExposureAllocator>(
            portfolio, exposureCalculator_->exposureCube(),
            nettedExposureCalculator_->exposureCube(), cube_,
            ExposureCalculator::allocatedEPE, ExposureCalculator::allocatedENE,
            ExposureCalculator::EPE, ExposureCalculator::ENE,
            NettedExposureCalculator::EPE, NettedExposureCalculator::ENE);
    else if (allocationMethod == ExposureAllocator::AllocationMethod::RelativeFairValueGross)
        exposureAllocator = boost::make_shared<RelativeFairValueGrossExposureAllocator>(
            portfolio, exposureCalculator_->exposureCube(),
            nettedExposureCalculator_->exposureCube(), cube_,
            ExposureCalculator::allocatedEPE, ExposureCalculator::allocatedENE,
            ExposureCalculator::EPE, ExposureCalculator::ENE,
            NettedExposureCalculator::EPE, NettedExposureCalculator::ENE);
    else if (allocationMethod == ExposureAllocator::AllocationMethod::RelativeXVA)
        exposureAllocator = boost::make_shared<RelativeXvaExposureAllocator>(
            portfolio, exposureCalculator_->exposureCube(),
            nettedExposureCalculator_->exposureCube(), cube_,
            cvaCalculator_->tradeCva(), cvaCalculator_->tradeDva(),
            cvaCalculator_->nettingSetSumCva(), cvaCalculator_->nettingSetSumDva(),
            ExposureCalculator::allocatedEPE, ExposureCalculator::allocatedENE,
            ExposureCalculator::EPE, ExposureCalculator::ENE,
            NettedExposureCalculator::EPE, NettedExposureCalculator::ENE);
    else if (allocationMethod == ExposureAllocator::AllocationMethod::None)
        exposureAllocator = boost::make_shared<NoneExposureAllocator>(
            portfolio, exposureCalculator_->exposureCube(),
            nettedExposureCalculator_->exposureCube());
    else
        QL_FAIL("allocationMethod " << allocationMethod << " not available");
    if(exposureAllocator)
        exposureAllocator->build();

    /********************************************************
     * Update Allocated XVAs
     */
    if (analytics_["dynamicCredit"]) {
        allocatedCvaCalculator_ = boost::make_shared<DynamicCreditXvaCalculator>(
            portfolio_, market_, configuration_, baseCurrency_, dvaName_, fvaBorrowingCurve_, fvaLendingCurve_,
            analytics_["dim"], dimCalculator, exposureCalculator_->exposureCube(),
            nettedExposureCalculator_->exposureCube(), cptyCube_, ExposureCalculator::ExposureIndex::allocatedEPE,
            ExposureCalculator::ExposureIndex::allocatedENE, NettedExposureCalculator::ExposureIndex::EPE,
            NettedExposureCalculator::ExposureIndex::ENE, 0, analytics_["flipViewXVA"], flipViewBorrowingCurvePostfix,
            flipViewLendingCurvePostfix);
    } else {
        allocatedCvaCalculator_ = boost::make_shared<StaticCreditXvaCalculator>(
            portfolio_, market_, configuration_, baseCurrency_, dvaName_, fvaBorrowingCurve_, fvaLendingCurve_,
            analytics_["dim"], dimCalculator, exposureCalculator_->exposureCube(),
            nettedExposureCalculator_->exposureCube(), ExposureCalculator::ExposureIndex::allocatedEPE,
            ExposureCalculator::ExposureIndex::allocatedENE, NettedExposureCalculator::ExposureIndex::EPE,
            NettedExposureCalculator::ExposureIndex::ENE, analytics_["flipViewXVA"], flipViewBorrowingCurvePostfix,
            flipViewLendingCurvePostfix);
    }
    allocatedCvaCalculator_->build();

    /********************************************************
     * Cache average EPE and ENE
     */
    for (auto tradeId : tradeIds()) {
        tradeEPE_[tradeId] = exposureCalculator_->epe(tradeId);
        tradeENE_[tradeId] = exposureCalculator_->ene(tradeId);
        allocatedTradeEPE_[tradeId] = exposureCalculator_->allocatedEpe(tradeId);
        allocatedTradeENE_[tradeId] = exposureCalculator_->allocatedEne(tradeId);
    }
    for (auto nettingSetId : nettingSetIds()) {
        netEPE_[nettingSetId] = nettedExposureCalculator_->epe(nettingSetId);
        netENE_[nettingSetId] = nettedExposureCalculator_->ene(nettingSetId);
    }

    /********************************************************
     * Calculate netting set KVA-CCR and KVA-CVA
     */
    updateNettingSetKVA();

    /***************************************
     * Calculate netting set CVA sensitivity
     */
    updateNettingSetCvaSensitivity();
}

void PostProcess::updateNettingSetKVA() {

    // Loop over all netting sets
    for (auto nettingSetId : nettingSetIds()) {
        // Init results
        ourNettingSetKVACCR_[nettingSetId] = 0.0;
        theirNettingSetKVACCR_[nettingSetId] = 0.0;
        ourNettingSetKVACVA_[nettingSetId] = 0.0;
        theirNettingSetKVACVA_[nettingSetId] = 0.0;
    }

    if (!analytics_["kva"])
        return;

    LOG("Update netting set KVA");

    vector<Date> dateVector = cube_->dates();
    Size dates = dateVector.size();
    Date today = market_->asofDate();
    Handle<YieldTermStructure> discountCurve = market_->discountCurve(baseCurrency_, configuration_);
    DayCounter dc = ActualActual();

    // Loop over all netting sets
    for (auto nettingSetId : nettingSetIds()) {
        string cid;
        if (analytics_["flipViewXVA"]) {
            cid = dvaName_;
        } else {
            cid = nettedExposureCalculator_->counterparty(nettingSetId);
        }
        LOG("KVA for netting set " << nettingSetId);

        // Main input are the EPE and ENE profiles, previously computed
        vector<Real> epe = netEPE_[nettingSetId];
        vector<Real> ene = netENE_[nettingSetId];

        // PD from counterparty Dts, floored to avoid 0 ...
        // Today changed to today+1Y to get the one-year PD
        Handle<DefaultProbabilityTermStructure> cvaDts = market_->defaultCurve(cid, configuration_);
        QL_REQUIRE(!cvaDts.empty(), "Default curve missing for counterparty " << cid);
        Real cvaRR = market_->recoveryRate(cid, configuration_)->value();
        Real PD1 = std::max(cvaDts->defaultProbability(today + 1 * Years), 0.000000000001);
        Real LGD1 = (1 - cvaRR);

        // FIXME: if flipViewXVA is sufficient, then all code for their KVA-CCR could be discarded here...
        Handle<DefaultProbabilityTermStructure> dvaDts;
        Real dvaRR = 0.0;
        Real PD2 = 0;
        if (analytics_["flipViewXVA"]) {
            dvaName_ = nettedExposureCalculator_->counterparty(nettingSetId);
        }
        if (dvaName_ != "") {
            dvaDts = market_->defaultCurve(dvaName_, configuration_);
            dvaRR = market_->recoveryRate(dvaName_, configuration_)->value();
            PD2 = std::max(dvaDts->defaultProbability(today + 1 * Years), 0.000000000001);
        } else {
            ALOG("dvaName not specified, own PD set to zero for their KVA calculation");
        }
        Real LGD2 = (1 - dvaRR);

        // Granularity adjustment, Gordy (2004):
        Real rho1 = 0.12 * (1 - std::exp(-50 * PD1)) / (1 - std::exp(-50)) +
                    0.24 * (1 - (1 - std::exp(-50 * PD1)) / (1 - std::exp(-50)));
        Real rho2 = 0.12 * (1 - std::exp(-50 * PD2)) / (1 - std::exp(-50)) +
                    0.24 * (1 - (1 - std::exp(-50 * PD2)) / (1 - std::exp(-50)));

        // Basel II internal rating based (IRB) estimate of worst case PD:
        // Large homogeneous pool (LHP) approximation of Vasicek (1997)
        InverseCumulativeNormal icn;
        CumulativeNormalDistribution cnd;
        Real PD99_1 = cnd((icn(PD1) + std::sqrt(rho1) * icn(0.999)) / (std::sqrt(1 - rho1))) - PD1;
        Real PD99_2 = cnd((icn(PD2) + std::sqrt(rho2) * icn(0.999)) / (std::sqrt(1 - rho2))) - PD2;

        // KVA regulatory PD, worst case PD, floored at 0.03 for corporates and banks, not floored for sovereigns
        Real kva99PD1 = std::max(PD99_1, kvaTheirPdFloor_);
        Real kva99PD2 = std::max(PD99_2, kvaOurPdFloor_);

        // Factor B(PD) for the maturity adjustment factor, B(PD) = (0.11852 - 0.05478 * ln(PD)) ^ 2
        Real kvaMatAdjB1 = std::pow((0.11852 - 0.05478 * std::log(PD1)), 2.0);
        Real kvaMatAdjB2 = std::pow((0.11852 - 0.05478 * std::log(PD2)), 2.0);

        DLOG("Our KVA-CCR " << nettingSetId << ": PD=" << PD1);
        DLOG("Our KVA-CCR " << nettingSetId << ": LGD=" << LGD1);
        DLOG("Our KVA-CCR " << nettingSetId << ": rho=" << rho1);
        DLOG("Our KVA-CCR " << nettingSetId << ": PD99=" << PD99_1);
        DLOG("Our KVA-CCR " << nettingSetId << ": PD Floor=" << kvaTheirPdFloor_);
        DLOG("Our KVA-CCR " << nettingSetId << ": Floored PD99=" << kva99PD1);
        DLOG("Our KVA-CCR " << nettingSetId << ": B(PD)=" << kvaMatAdjB1);

        DLOG("Their KVA-CCR " << nettingSetId << ": PD=" << PD2);
        DLOG("Their KVA-CCR " << nettingSetId << ": LGD=" << LGD2);
        DLOG("Their KVA-CCR " << nettingSetId << ": rho=" << rho2);
        DLOG("Their KVA-CCR " << nettingSetId << ": PD99=" << PD99_2);
        DLOG("Their KVA-CCR " << nettingSetId << ": PD Floor=" << kvaOurPdFloor_);
        DLOG("Their KVA-CCR " << nettingSetId << ": Floored PD99=" << kva99PD2);
        DLOG("Their KVA-CCR " << nettingSetId << ": B(PD)=" << kvaMatAdjB2);

        for (Size j = 0; j < dates; ++j) {
            Date d0 = j == 0 ? today : cube_->dates()[j - 1];
            Date d1 = cube_->dates()[j];

            // Preprocess:
            // 1) Effective maturity from effective expected exposure as of time j
            //    Index _1 corresponds to our perspective, index _2 to their perspective.
            // 2) Basel EEPE as of time j, i.e. as time averge over EEE, starting at time j
            // More accuracy may be achieved here by using a Longstaff-Schwartz method / regression
            Real eee_kva_1 = 0.0, eee_kva_2 = 0.0;
            Real effMatNumer1 = 0.0, effMatNumer2 = 0.0;
            Real effMatDenom1 = 0.0, effMatDenom2 = 0.0;
            Real eepe_kva_1 = 0, eepe_kva_2 = 0.0;
            Size kmax = j, count = 0;
            // Cut off index for EEPE/EENE calculation: One year ahead
            while (dateVector[kmax] < dateVector[j] + 1 * Years + 4 * Days && kmax < dates - 1)
                kmax++;
            Real sumdt = 0.0, eee1_b = 0.0, eee2_b = 0.0;
            for (Size k = j; k < dates; ++k) {
                Date d2 = cube_->dates()[k];
                Date prevDate = k == 0 ? today : dateVector[k - 1];

                eee_kva_1 = std::max(eee_kva_1, epe[k + 1]);
                eee_kva_2 = std::max(eee_kva_2, ene[k + 1]);

                // Components of the KVA maturity adjustment MA as of time j
                if (dc.yearFraction(d1, d2) > 1.0) {
                    effMatNumer1 += epe[k + 1] * dc.yearFraction(prevDate, d2);
                    effMatNumer2 += ene[k + 1] * dc.yearFraction(prevDate, d2);
                }
                if (dc.yearFraction(d1, d2) <= 1.0) {
                    effMatDenom1 += eee_kva_1 * dc.yearFraction(prevDate, d2);
                    effMatDenom2 += eee_kva_2 * dc.yearFraction(prevDate, d2);
                }

                if (k < kmax) {
                    Real dt = dc.yearFraction(cube_->dates()[k], cube_->dates()[k + 1]);
                    sumdt += dt;
                    Real epe_b = epe[k + 1] / discountCurve->discount(dateVector[k]);
                    Real ene_b = ene[k + 1] / discountCurve->discount(dateVector[k]);
                    eee1_b = std::max(epe_b, eee1_b);
                    eee2_b = std::max(ene_b, eee2_b);
                    eepe_kva_1 += eee1_b * dt;
                    eepe_kva_2 += eee2_b * dt;
                    count++;
                }
            }

            // Normalize EEPE/EENE calculation
            eepe_kva_1 = count > 0 ? eepe_kva_1 / sumdt : 0.0;
            eepe_kva_2 = count > 0 ? eepe_kva_2 / sumdt : 0.0;

            // KVA CCR using the IRB risk weighted asset method and IMM:
            // KVA effective maturity of the nettingSet, capped at 5
            Real kvaNWMaturity1 = std::min(1.0 + (effMatDenom1 == 0.0 ? 0.0 : effMatNumer1 / effMatDenom1), 5.0);
            Real kvaNWMaturity2 = std::min(1.0 + (effMatDenom2 == 0.0 ? 0.0 : effMatNumer2 / effMatDenom2), 5.0);

            // Maturity adjustment factor for the RWA method:
            // MA(PD, M) = (1 + (M - 2.5) * B(PD)) / (1 - 1.5 * B(PD)), capped at 5, floored at 1, M = effective
            // maturity
            Real kvaMatAdj1 =
                std::max(std::min((1.0 + (kvaNWMaturity1 - 2.5) * kvaMatAdjB1) / (1.0 - 1.5 * kvaMatAdjB1), 5.0), 1.0);
            Real kvaMatAdj2 =
                std::max(std::min((1.0 + (kvaNWMaturity2 - 2.5) * kvaMatAdjB2) / (1.0 - 1.5 * kvaMatAdjB2), 5.0), 1.0);

            // CCR Capital: RC = EAD x LGD x PD99.9 x MA(PD, M); EAD = alpha x EEPE(t) (approximated by EPE here);
            Real kvaRC1 = kvaAlpha_ * eepe_kva_1 * LGD1 * kva99PD1 * kvaMatAdj1;
            Real kvaRC2 = kvaAlpha_ * eepe_kva_2 * LGD2 * kva99PD2 * kvaMatAdj2;

            // Expected risk capital discounted at capital discount rate
            Real kvaCapitalDiscount = 1 / std::pow(1 + kvaCapitalDiscountRate_, dc.yearFraction(today, d0));
            Real kvaCCRIncrement1 =
                kvaRC1 * kvaCapitalDiscount * dc.yearFraction(d0, d1) * kvaCapitalHurdle_ * kvaRegAdjustment_;
            Real kvaCCRIncrement2 =
                kvaRC2 * kvaCapitalDiscount * dc.yearFraction(d0, d1) * kvaCapitalHurdle_ * kvaRegAdjustment_;

            ourNettingSetKVACCR_[nettingSetId] += kvaCCRIncrement1;
            theirNettingSetKVACCR_[nettingSetId] += kvaCCRIncrement2;

            DLOG("Our KVA-CCR for " << nettingSetId << ": " << j << " EEPE=" << setprecision(2) << eepe_kva_1
                                    << " EPE=" << epe[j] << " RC=" << kvaRC1 << " M=" << setprecision(6)
                                    << kvaNWMaturity1 << " MA=" << kvaMatAdj1 << " Cost=" << setprecision(2)
                                    << kvaCCRIncrement1 << " KVA=" << ourNettingSetKVACCR_[nettingSetId]);
            DLOG("Their KVA-CCR for " << nettingSetId << ": " << j << " EENE=" << eepe_kva_2 << " ENE=" << ene[j]
                                      << " RC=" << kvaRC2 << " M=" << setprecision(6) << kvaNWMaturity2
                                      << " MA=" << kvaMatAdj2 << " Cost=" << setprecision(2) << kvaCCRIncrement2
                                      << " KVA=" << theirNettingSetKVACCR_[nettingSetId]);

            // CVA Capital
            // effective maturity without cap at 5, DF set to 1 for IMM banks
            // TODO: Set MA in CCR capital calculation to 1
            Real kvaCvaMaturity1 = 1.0 + (effMatDenom1 == 0.0 ? 0.0 : effMatNumer1 / effMatDenom1);
            Real kvaCvaMaturity2 = 1.0 + (effMatDenom2 == 0.0 ? 0.0 : effMatNumer2 / effMatDenom2);
            Real scva1 = kvaTheirCvaRiskWeight_ * kvaCvaMaturity1 * eepe_kva_1;
            Real scva2 = kvaOurCvaRiskWeight_ * kvaCvaMaturity2 * eepe_kva_2;
            Real kvaCVAIncrement1 =
                scva1 * kvaCapitalDiscount * dc.yearFraction(d0, d1) * kvaCapitalHurdle_ * kvaRegAdjustment_;
            Real kvaCVAIncrement2 =
                scva2 * kvaCapitalDiscount * dc.yearFraction(d0, d1) * kvaCapitalHurdle_ * kvaRegAdjustment_;

            DLOG("Our KVA-CVA for " << nettingSetId << ": " << j << " EEPE=" << eepe_kva_1 << " SCVA=" << scva1
                                    << " Cost=" << kvaCVAIncrement1);
            DLOG("Their KVA-CVA for " << nettingSetId << ": " << j << " EENE=" << eepe_kva_2 << " SCVA=" << scva2
                                      << " Cost=" << kvaCVAIncrement2);

            ourNettingSetKVACVA_[nettingSetId] += kvaCVAIncrement1;
            theirNettingSetKVACVA_[nettingSetId] += kvaCVAIncrement2;
        }
    }

    LOG("Update netting set KVA done");
}

void PostProcess::updateNettingSetCvaSensitivity() {

    if (!analytics_["cvaSensi"])
        return;

    LOG("Update netting set CVA sensitvities");

    Handle<YieldTermStructure> discountCurve = market_->discountCurve(baseCurrency_, configuration_);

    for (auto n : netEPE_) {
        string nettingSetId = n.first;
        vector<Real> epe = netEPE_[nettingSetId];
        string cid;
        Handle<DefaultProbabilityTermStructure> cvaDts;
        Real cvaRR;
        if (analytics_["flipViewXVA"]) {
            cid = dvaName_;
        } else {
            cid = nettedExposureCalculator_->counterparty(nettingSetId);
        }
        cvaDts = market_->defaultCurve(cid);
        QL_REQUIRE(!cvaDts.empty(), "Default curve missing for counterparty " << cid);
        cvaRR = market_->recoveryRate(cid, configuration_)->value();

	    bool cvaSensi = analytics_["cvaSensi"];
	    LOG("CVA Sensitivity: " << cvaSensi);
	    if (cvaSensi) {
	        boost::shared_ptr<CVASpreadSensitivityCalculator> cvaSensiCalculator = boost::make_shared<CVASpreadSensitivityCalculator>(
	             nettingSetId, market_->asofDate(), epe, cube_->dates(), cvaDts, cvaRR, discountCurve, cvaSpreadSensiGrid_, cvaSpreadSensiShiftSize_);

	        for (Size i = 0; i < cvaSensiCalculator->shiftTimes().size(); ++i) {
	            DLOG("CVA Sensi Calculator: t=" << cvaSensiCalculator->shiftTimes()[i]
		         << " h=" << cvaSensiCalculator->hazardRateSensitivities()[i]
		         << " s=" << cvaSensiCalculator->cdsSpreadSensitivities()[i]);
	        }

	        netCvaHazardRateSensi_[nettingSetId] = cvaSensiCalculator->hazardRateSensitivities();
	        netCvaSpreadSensi_[nettingSetId] = cvaSensiCalculator->cdsSpreadSensitivities();
	        cvaSpreadSensiTimes_ = cvaSensiCalculator->shiftTimes();
	    } else {
	        cvaSpreadSensiTimes_ = vector<Real>();
	        netCvaHazardRateSensi_[nettingSetId] = vector<Real>();
	        netCvaSpreadSensi_[nettingSetId] = vector<Real>();
	    }
    }

    LOG("Update netting set CVA sensitivities done");
}

const vector<Real>& PostProcess::tradeEPE(const string& tradeId) {
    QL_REQUIRE(tradeEPE_.find(tradeId) != tradeEPE_.end(), "Trade " << tradeId << " not found in exposure map");
    return tradeEPE_[tradeId];
}

const vector<Real>& PostProcess::tradeENE(const string& tradeId) {
    QL_REQUIRE(tradeENE_.find(tradeId) != tradeENE_.end(), "Trade " << tradeId << " not found in exposure map");
    return tradeENE_[tradeId];
}

const vector<Real>& PostProcess::tradeEE_B(const string& tradeId) {
    return exposureCalculator_->ee_b(tradeId);
}

const Real& PostProcess::tradeEPE_B(const string& tradeId) {
    return exposureCalculator_->epe_b(tradeId);
}

const vector<Real>& PostProcess::tradeEEE_B(const string& tradeId) {
    return exposureCalculator_->eee_b(tradeId);
}

const Real& PostProcess::tradeEEPE_B(const string& tradeId) {
    return exposureCalculator_->eepe_b(tradeId);
}

const vector<Real>& PostProcess::tradePFE(const string& tradeId) {
    return exposureCalculator_->pfe(tradeId);
}

const vector<Real>& PostProcess::netEPE(const string& nettingSetId) {
    QL_REQUIRE(netEPE_.find(nettingSetId) != netEPE_.end(),
               "Netting set " << nettingSetId << " not found in exposure map");
    return netEPE_[nettingSetId];
}

const vector<Real>& PostProcess::netENE(const string& nettingSetId) {
    QL_REQUIRE(netENE_.find(nettingSetId) != netENE_.end(),
               "Netting set " << nettingSetId << " not found in exposure map");
    return netENE_[nettingSetId];
}

vector<Real> PostProcess::netCvaHazardRateSensitivity(const string& nettingSetId) {
    if (netCvaHazardRateSensi_.find(nettingSetId) != netCvaHazardRateSensi_.end())
        return netCvaHazardRateSensi_[nettingSetId];
    else
        return vector<Real>();
}

vector<Real> PostProcess::netCvaSpreadSensitivity(const string& nettingSetId) {
    if (netCvaSpreadSensi_.find(nettingSetId) != netCvaSpreadSensi_.end())
        return netCvaSpreadSensi_[nettingSetId];
    else
        return vector<Real>();
}

const vector<Real>& PostProcess::netEE_B(const string& nettingSetId) {
    return nettedExposureCalculator_->ee_b(nettingSetId);
}

const Real& PostProcess::netEPE_B(const string& nettingSetId) {
    return nettedExposureCalculator_->epe_b(nettingSetId);
}

const vector<Real>& PostProcess::netEEE_B(const string& nettingSetId) {
    return nettedExposureCalculator_->eee_b(nettingSetId);
}

const Real& PostProcess::netEEPE_B(const string& nettingSetId) {
    return nettedExposureCalculator_->eepe_b(nettingSetId);
}

const vector<Real>& PostProcess::netPFE(const string& nettingSetId) {
    return nettedExposureCalculator_->pfe(nettingSetId);
}

const vector<Real>& PostProcess::expectedCollateral(const string& nettingSetId) {
    return nettedExposureCalculator_->expectedCollateral(nettingSetId);
}

const vector<Real>& PostProcess::colvaIncrements(const string& nettingSetId) {
    return nettedExposureCalculator_->colvaIncrements(nettingSetId);
}

const vector<Real>& PostProcess::collateralFloorIncrements(const string& nettingSetId) {
    return nettedExposureCalculator_->collateralFloorIncrements(nettingSetId);
}

const vector<Real>& PostProcess::allocatedTradeEPE(const string& tradeId) {
    QL_REQUIRE(allocatedTradeEPE_.find(tradeId) != allocatedTradeEPE_.end(),
               "Trade " << tradeId << " not found in exposure map");
    return allocatedTradeEPE_[tradeId];
}

const vector<Real>& PostProcess::allocatedTradeENE(const string& tradeId) {
    QL_REQUIRE(allocatedTradeENE_.find(tradeId) != allocatedTradeENE_.end(),
               "Trade " << tradeId << " not found in exposure map");
    return allocatedTradeENE_[tradeId];
}

Real PostProcess::tradeCVA(const string& tradeId) {
    return cvaCalculator_->tradeCva(tradeId);
}

Real PostProcess::tradeDVA(const string& tradeId) {
    return cvaCalculator_->tradeDva(tradeId);
}

Real PostProcess::tradeMVA(const string& tradeId) {
    return cvaCalculator_->tradeMva(tradeId);
}

Real PostProcess::tradeFBA(const string& tradeId) {
    return cvaCalculator_->tradeFba(tradeId);
}

Real PostProcess::tradeFCA(const string& tradeId) {
    return cvaCalculator_->tradeFca(tradeId);
}

Real PostProcess::tradeFBA_exOwnSP(const string& tradeId) {
    return cvaCalculator_->tradeFba_exOwnSp(tradeId);
}

Real PostProcess::tradeFCA_exOwnSP(const string& tradeId) {
    return cvaCalculator_->tradeFca_exOwnSp(tradeId);
}

Real PostProcess::tradeFBA_exAllSP(const string& tradeId) {
    return cvaCalculator_->tradeFba_exAllSp(tradeId);
}

Real PostProcess::tradeFCA_exAllSP(const string& tradeId) {
    return cvaCalculator_->tradeFca_exAllSp(tradeId);
}

Real PostProcess::nettingSetCVA(const string& nettingSetId) {
    return cvaCalculator_->nettingSetCva(nettingSetId);
}

Real PostProcess::nettingSetDVA(const string& nettingSetId) {
    return cvaCalculator_->nettingSetDva(nettingSetId);
}

Real PostProcess::nettingSetMVA(const string& nettingSetId) {
    return cvaCalculator_->nettingSetMva(nettingSetId);

}

Real PostProcess::nettingSetFBA(const string& nettingSetId) {
    return cvaCalculator_->nettingSetFba(nettingSetId);

}

Real PostProcess::nettingSetFCA(const string& nettingSetId) {
    return cvaCalculator_->nettingSetFca(nettingSetId);

}

Real PostProcess::nettingSetOurKVACCR(const string& nettingSetId) {
    QL_REQUIRE(ourNettingSetKVACCR_.find(nettingSetId) != ourNettingSetKVACCR_.end(),
               "NettingSetId " << nettingSetId << " not found in nettingSet KVACCR map");
    return ourNettingSetKVACCR_[nettingSetId];
}

Real PostProcess::nettingSetTheirKVACCR(const string& nettingSetId) {
    QL_REQUIRE(theirNettingSetKVACCR_.find(nettingSetId) != theirNettingSetKVACCR_.end(),
               "NettingSetId " << nettingSetId << " not found in nettingSet KVACCR map");
    return theirNettingSetKVACCR_[nettingSetId];
}

Real PostProcess::nettingSetOurKVACVA(const string& nettingSetId) {
    QL_REQUIRE(ourNettingSetKVACVA_.find(nettingSetId) != ourNettingSetKVACVA_.end(),
               "NettingSetId " << nettingSetId << " not found in nettingSet KVACVA map");
    return ourNettingSetKVACVA_[nettingSetId];
}

Real PostProcess::nettingSetTheirKVACVA(const string& nettingSetId) {
    QL_REQUIRE(theirNettingSetKVACVA_.find(nettingSetId) != theirNettingSetKVACVA_.end(),
               "NettingSetId " << nettingSetId << " not found in nettingSet KVACVA map");
    return theirNettingSetKVACVA_[nettingSetId];
}

Real PostProcess::nettingSetFBA_exOwnSP(const string& nettingSetId) {
    return cvaCalculator_->nettingSetFba_exOwnSp(nettingSetId);
}

Real PostProcess::nettingSetFCA_exOwnSP(const string& nettingSetId) {
    return cvaCalculator_->nettingSetFca_exOwnSp(nettingSetId);
}

Real PostProcess::nettingSetFBA_exAllSP(const string& nettingSetId) {
    return cvaCalculator_->nettingSetFba_exAllSp(nettingSetId);
}

Real PostProcess::nettingSetFCA_exAllSP(const string& nettingSetId) {
    return cvaCalculator_->nettingSetFca_exAllSp(nettingSetId);
}

Real PostProcess::allocatedTradeCVA(const string& allocatedTradeId) {
    return allocatedCvaCalculator_->tradeCva(allocatedTradeId);
}

Real PostProcess::allocatedTradeDVA(const string& allocatedTradeId) {
    return allocatedCvaCalculator_->tradeDva(allocatedTradeId);
}

Real PostProcess::nettingSetCOLVA(const string& nettingSetId) {
    return nettedExposureCalculator_->colva(nettingSetId);
}

Real PostProcess::nettingSetCollateralFloor(const string& nettingSetId) {
    return nettedExposureCalculator_->collateralFloor(nettingSetId);
}

void PostProcess::exportDimEvolution(ore::data::Report& dimEvolutionReport) {
    dimCalculator_->exportDimEvolution(dimEvolutionReport);
}

void PostProcess::exportDimRegression(const std::string& nettingSet, const std::vector<Size>& timeSteps,
                                      const std::vector<boost::shared_ptr<ore::data::Report>>& dimRegReports) {

    boost::shared_ptr<RegressionDynamicInitialMarginCalculator> regCalc =
        boost::dynamic_pointer_cast<RegressionDynamicInitialMarginCalculator>(dimCalculator_);

    if (regCalc)
        regCalc->exportDimRegression(nettingSet, timeSteps, dimRegReports);
}

} // namespace analytics
} // namespace ore


#ifndef __SDT_HH
#define __SDT_HH
#include <armadillo>
#include <functional>
#include "Module.hh"
#include "nxbus_engine.h"
class SDT : public Sensor {
  TRICK_INTERFACE(SDT);

 public:
  SDT()
      : VECTOR_INIT(PHI, 3),
        VECTOR_INIT(DELTA_VEL, 3) {};
  virtual ~SDT(){};

  virtual void init() = 0;
  virtual void algorithm(double int_step) = 0;

  std::function<arma::vec3()> grab_WBICB;
  std::function<arma::vec3()> grab_FSPCB;
  std::function<arma::vec3()> grab_CONING;
  std::function<arma::vec3()> grab_GHIGH;
  std::function<arma::vec3()> grab_GLOW;
  std::function<arma::vec3()> grab_AHIGH;
  std::function<arma::vec3()> grab_ALOW;

  virtual arma::vec3 get_PHI() { return PHI; }
  virtual arma::vec3 get_DELTA_VEL() { return DELTA_VEL; }
  virtual int write_to_(const char *bus_name) { return 0; }

 protected:

  VECTOR(PHI, 3); /* *o (r)       PHI = DELTA_ALPHA + 0.5 * DELTA_BETA */

  VECTOR(DELTA_VEL, 3); /* *o (m/s)     Delta V */
};

#endif