/*
 * This file is part of the Carpe Diem Active Learning Software,
 * Copyright (C) 2017 Clemens-Alexander Brust (ikosa dot de at gmail dot com).
 *
 * For licensing information, see the LICENSE file included with this project.
 */

#include "Project.h"

#include <string>
#include <sstream>
#include <chrono>
#include <iomanip>
#include <ctime>
#include <cn24.h>

Project::Project(QObject* _parent, Conv::ClassManager::ClassUpdateHandler* class_update_handler, ProjectStateHandler* project_state_handler)
  : QObject(_parent), class_update_handler_(class_update_handler), state_handler_(project_state_handler), state(NOTHING)
{

}

void Project::PredictSegment(ProjectProgressHandler *progress_handler, Conv::Segment *segment, std::vector<std::string> *prediction_filenames, std::vector<std::vector<Conv::BoundingBox> > *prediction_boxes) {
  if(state != Project::LOADED)
    return;

  progress_handler->OnProjectProgressUpdate(0);

  // Switch to testing mode
  graph_->SetIsTesting(true);
  input_layer_->ForceWeightsZero();

  Conv::datum total_sample_count = segment->GetSampleCount();
  Conv::datum running_sample_count = 0;

  for(unsigned int sample = 0; sample < segment->GetSampleCount(); sample++) {
    Conv::JSON& sample_json = segment->GetSample(sample);
    std::vector<Conv::BoundingBox> sample_predictions;

    // Load sample
    input_layer_->ForceLoadDetection(sample_json, 0);
    graph_->FeedForward();

    // Copy predictions
    Conv::DetectionMetadataPointer output_boxes = predicted_metadata_[0];
    for(unsigned int b = 0; b < output_boxes->size(); b++) {
      Conv::BoundingBox bbox = output_boxes->at(b);
      sample_predictions.push_back(bbox);
    }

    // Store predictions
    std::string sample_filename = sample_json["image_rpath"];
    prediction_filenames->push_back(sample_filename);
    prediction_boxes->push_back(sample_predictions);

    running_sample_count += 1.0;
    progress_handler->OnProjectProgressUpdate(running_sample_count / total_sample_count);
  }

  // Done
  progress_handler->OnProjectProgressUpdate(1);
  progress_handler->OnProjectProgressDone();
}

void Project::UpdateModel(ProjectProgressHandler *progress_handler) {
  if(update_set_->GetSampleCount() == 0) {
    progress_handler->OnProjectProgressDone();
    return;
  }
  needs_rescore_ = true;
  // Update input layer settings
  input_layer_->training_sets_.clear();
  input_layer_->training_weights_.clear();
  input_layer_->training_sets_.push_back(known_samples_);
  input_layer_->training_sets_.push_back(update_set_);
  input_layer_->training_weights_.push_back(1);
  input_layer_->training_weights_.push_back(1);
  input_layer_->UpdateDatasets();

  ProjectTrainerProgressHandler trainer_progress_handler(progress_handler);

  progress_handler->OnProjectProgressUpdate(0);
  trainer_->SetUpdateHandler(&trainer_progress_handler);
  trainer_->settings()["epoch_iterations"] = 10 * update_set_->GetSampleCount();
  trainer_->Train(1, false);

  progress_handler->OnProjectProgressUpdate(1);
  Save();
  progress_handler->OnProjectProgressDone();
}

void Project::UpdateScores(ProjectProgressHandler *progress_handler) {
  progress_handler->OnProjectProgressUpdate(0);
  if(!needs_rescore_) {
    progress_handler->OnProjectProgressDone();
    return;
  }

  Conv::JSONNetGraphFactory netgraph_factory(architecture_json, 123123);
  Conv::ActiveLearningPolicy* policy = Conv::YOLOActiveLearningPolicy::CreateWithName(active_learning_policy_,netgraph_factory.GetYOLOConfiguration());

  Conv::NetGraphBuffer& prediction_buffer = graph_->GetOutputNodes()[0]->output_buffers[0];
  Conv::DatasetMetadataPointer* predicted_metadata = prediction_buffer.combined_tensor->metadata;

  input_layer_->ForceWeightsZero();
  graph_->SetIsTesting(true);

  Conv::datum total_sample_count = new_set_->GetSampleCount();
  Conv::datum running_sample_count = 0;

  unsigned int batch_size = prediction_buffer.combined_tensor->data.samples();

  for(unsigned int s = 0; s < new_set_->GetSegmentCount(); s++) {
    Conv::Segment* segment = new_set_->GetSegment(s);
    Conv::datum segment_score = 0;

    for(unsigned int sample = 0; sample < segment->GetSampleCount(); sample+= batch_size) {
      for(unsigned int bindex = 0; bindex < batch_size && (sample+bindex) < (segment->GetSampleCount()); bindex++) {
        Conv::JSON& sample_json = segment->GetSample(sample+bindex);
        input_layer_->ForceLoadDetection(sample_json, bindex);
      }
      graph_->FeedForward();
      for(unsigned int bindex = 0; bindex < batch_size && (sample+bindex) < (segment->GetSampleCount()); bindex++) {
        segment_score += policy->Score(prediction_buffer.combined_tensor->data, predicted_metadata, bindex);
        running_sample_count += 1.0;
        progress_handler->OnProjectProgressUpdate(running_sample_count / total_sample_count);
      }
    }
    segment->score = segment_score;
  }

  delete policy;
  needs_rescore_ = false;
  progress_handler->OnProjectProgressDone();
}

bool Project::Save() {
  bool model_result = SaveModel();
  if(!model_result)
    return false;

  Conv::JSON project_json = Serialize();
  std::string project_filename = project_folder_ + "/project.json";
  std::ofstream project_file(project_filename, std::ios::out);
  if(!project_file.good()) {
    LOGERROR << "Could not open " << project_filename << " for writing";
    return false;
  }

  project_file << project_json.dump(2);
  return true;
}

bool Project::Load(std::string project_folder) {
  needs_rescore_ = true;
  project_folder_ = project_folder;
  std::string project_filename = project_folder_ + "/project.json";
  std::ifstream project_file(project_filename, std::ios::in);
  if(!project_file.good()) {
    LOGERROR << "Could not open " << project_filename << " for reading";
    return false;
  }
  Conv::JSON project_json = Conv::JSON::parse(project_file);
  bool project_result = Deserialize(project_json);
  if(!project_result) {
    SetState(FAILED);
    return false;
  }

  SetState(Project::LOADED);

  // Load Model
  bool model_result = LoadModel();
  if(!model_result) {
    SetState(FAILED);
    return false;
  }

  return true;
}

bool Project::AddSample(std::string filename) {
  needs_rescore_ = true;
  Conv::Segment* target_segment = nullptr;

  // If no new samples have been added, create first segment
  if(new_set_->GetSegmentCount() == 0) {
    std::stringstream ss; ss << "New Data Batch " << 1;
    Conv::Segment* segment = new Conv::Segment(ss.str());
    new_set_->AddSegment(segment);
    target_segment = segment;
  } else {
    // Otherwise, get last segment and check if it has room for new sample
    Conv::Segment* segment = new_set_->GetSegment(new_set_->GetSegmentCount() - 1);
    if(segment->GetSampleCount() < new_batch_size_) {
      target_segment = segment;
    } else {
      // No room, create new segment
      std::stringstream ss; ss << "New Data Batch " << new_set_->GetSegmentCount() + 1;
      Conv::Segment* segment = new Conv::Segment(ss.str());
      new_set_->AddSegment(segment);
      target_segment = segment;
    }
  }

  Conv::JSON sample_json = Conv::JSON::object();
  sample_json["image_filename"] = filename;
  sample_json["boxes"] = Conv::JSON::array();

  return target_segment->AddSample(sample_json);
}

void Project::Predict(std::string image_filename, std::vector<Conv::BoundingBox> &predictions) {
  if(state != Project::LOADED)
    return;

  std::string found_path = Conv::PathFinder::FindPath(image_filename, "");
  if(found_path.length() > 0) {
    Conv::JSON sample_json = Conv::JSON::object();
    sample_json["image_rpath"] = found_path;
    sample_json["boxes"] = Conv::JSON::array();
    input_layer_->ForceWeightsZero();
    input_layer_->ForceLoadDetection(sample_json, 0);

    graph_->SetIsTesting(true);
    graph_->FeedForward();

    Conv::DetectionMetadataPointer output_boxes = predicted_metadata_[0];
    LOGINFO << "Predicted " << output_boxes->size() << " boxes.";

    for(unsigned int b = 0; b < output_boxes->size(); b++) {
      Conv::BoundingBox bbox = output_boxes->at(b);
      predictions.push_back(bbox);
    }
  } else {
    LOGERROR << "Could not find " << image_filename << "!";
  }
}

Conv::JSON Project::Serialize() {
  Conv::JSON project_json = Conv::JSON::object();
  project_json["architecture"] = architecture_json;
  project_json["update_set"] = update_set_->Serialize();
  project_json["known_set"] = known_samples_->Serialize();
  project_json["new_set"] = new_set_->Serialize();
  project_json["name"] = project_name_;
  project_json["new_batch_size"] = new_batch_size_;
  project_json["active_learning_policy"] = active_learning_policy_;
  return project_json;
}

bool Project::Deserialize(Conv::JSON& project_json) {
  if(state != Project::NOTHING) {
    LOGERROR << "Already have a project!";
    return false;
  }

  // Load JSON
  architecture_filename_="_from_json_";
  if(!project_json["architecture"].is_object()) {
    LOGERROR << "Project JSON is missing architecture";
    return false;
  }
  if(!project_json["update_set"].is_object() || !project_json["known_set"].is_object()) {
    LOGERROR << "Project JSON is missing set informations!";
    return false;
  }
  if(project_json.count("new_batch_size") == 1 && project_json["new_batch_size"].is_number()) {
    new_batch_size_ = project_json["new_batch_size"];
  }
  if(project_json.count("active_learning_policy") == 1 && project_json["active_learning_policy"].is_string()) {
    active_learning_policy_ = project_json["active_learning_policy"];
  }

  architecture_json = project_json["architecture"];

  // Create class manager
  class_manager_ = new Conv::ClassManager();
  class_manager_->RegisterClassUpdateHandler(class_update_handler_);

  // Load architecture
  Conv::JSONNetGraphFactory netgraph_factory(architecture_json, 123123);
  graph_ = new Conv::NetGraph();

  // Create dataset input layer
  unsigned int batch_size_parallel = 1;
  if(netgraph_factory.GetHyperparameters().count("batch_size_parallel") == 1 && netgraph_factory.GetHyperparameters()["batch_size_parallel"].is_number()) {
    batch_size_parallel = netgraph_factory.GetHyperparameters()["batch_size_parallel"];
  }
  input_layer_ = new Conv::SegmentSetInputLayer(netgraph_factory.GetDataInput(), Conv::DETECTION, class_manager_, batch_size_parallel, 123923);
  Conv::NetGraphNode* input_node = new Conv::NetGraphNode(input_layer_);
  input_node->is_input = true;

  // Add other layers
  graph_->AddNode(input_node);
  bool result = netgraph_factory.AddLayers(*graph_, class_manager_, 23923);
  if(!result) {
    SetState(Project::FAILED);
    LOGERROR << "Could not construct network!";
    return false;
  }
  graph_->Initialize();
  graph_->InitializeWeights(true);

  // Set helper pointers
  predicted_metadata_ = (Conv::DetectionMetadataPointer*) graph_->GetOutputNodes()[0]->output_buffers[0].combined_tensor->metadata;

  // Set project properties
  project_name_ = project_json["name"];

  // Load trainer
  trainer_ = new Conv::Trainer(*graph_, netgraph_factory.GetHyperparameters());

  // Load samples
  known_samples_ = new Conv::SegmentSet("Known Examples");
  bool deserialization_result = known_samples_->Deserialize(project_json["known_set"]);
  update_set_ = new Conv::SegmentSet("Update Set");
  deserialization_result &= update_set_->Deserialize(project_json["update_set"]);
  new_set_ = new Conv::SegmentSet("New Set");
  deserialization_result &= new_set_->Deserialize(project_json["new_set"]);

  if(!deserialization_result) {
    LOGERROR << "SegmentSet deserialization failed! See log for details.";
    return false;
  }

  return true;
}

bool Project::New(std::string architecture_filename, std::string model_filename, std::string project_name, std::string project_folder) {
  needs_rescore_ = true;
  if(state != Project::NOTHING) {
    LOGERROR << "Already have a project!";
    return false;
  } else {
    // Validate filenames
    std::ifstream architecture_file(architecture_filename, std::ios::in);
    if(!architecture_file.good()) {
      LOGERROR << "Failed to open architecture!";
      return false;
    }

    std::ifstream model_file(model_filename, std::ios::in | std::ios::binary);
    if(!model_file.good()) {
      LOGERROR << "Failed to open model!";
    }

    // Create class manager
    class_manager_ = new Conv::ClassManager();
    class_manager_->RegisterClassUpdateHandler(class_update_handler_);

    // Load architecture
    architecture_json = Conv::JSON::parse(architecture_file);
    Conv::JSONNetGraphFactory netgraph_factory(architecture_json, 123123);
    graph_ = new Conv::NetGraph();

    // Create dataset input layer
    unsigned int batch_size_parallel = 1;
    if(netgraph_factory.GetHyperparameters().count("batch_size_parallel") == 1 && netgraph_factory.GetHyperparameters()["batch_size_parallel"].is_number()) {
      batch_size_parallel = netgraph_factory.GetHyperparameters()["batch_size_parallel"];
    }
    input_layer_ = new Conv::SegmentSetInputLayer(netgraph_factory.GetDataInput(), Conv::DETECTION, class_manager_, batch_size_parallel, 123923);
    Conv::NetGraphNode* input_node = new Conv::NetGraphNode(input_layer_);
    input_node->is_input = true;

    // Add other layers
    graph_->AddNode(input_node);
    bool result = netgraph_factory.AddLayers(*graph_, class_manager_, 23923);
    if(!result) {
      SetState(Project::FAILED);
      LOGERROR << "Could not construct network!";
      return false;
    }
    graph_->Initialize();
    graph_->InitializeWeights(true);


    // Load model
    graph_->DeserializeParameters(model_file);

    // Load trainer
    trainer_ = new Conv::Trainer(*graph_, netgraph_factory.GetHyperparameters());

    // Initialize segment sets
    known_samples_ = new Conv::SegmentSet("Known Examples");
    update_set_ = new Conv::SegmentSet("Update Set");
    new_set_ = new Conv::SegmentSet("New Set");

    // Set helper pointers
    predicted_metadata_ = (Conv::DetectionMetadataPointer*) graph_->GetOutputNodes()[0]->output_buffers[0].combined_tensor->metadata;

    // Set project properties
    project_name_ = project_name;
    project_folder_ = project_folder;

    SetState(Project::LOADED);
    return true;
  }
}

void Project::SaveSnapshot() {
  if(state != LOADED)
    return;

  std::chrono::system_clock::time_point now = std::chrono::system_clock::now();
  std::time_t now_c = std::chrono::system_clock::to_time_t(now);
  std::stringstream ss;
  ss << project_folder_ << "/" << "snapshot-" << std::put_time(std::localtime(&now_c), "%y%m%d-%H%M") << "-" << known_samples_->GetSampleCount() << ".CNParamX";
  std::string model_filename = ss.str();
  std::ofstream model_file(model_filename, std::ios::out | std::ios::binary);
  if(!model_file.good()) {
    LOGERROR << "Could not open " << model_filename << " for serializing the model!";
  }
  graph_->SerializeParameters(model_file);
}

bool Project::LoadSnapshot(std::string &file_name) {
  if(state != LOADED)
    return false;

  needs_rescore_ = true;
  std::ifstream model_file(file_name, std::ios::in | std::ios::binary);
  if(!model_file.good()) {
    LOGERROR << "Could not open " << file_name << " for deserializing the model!";
    return false;
  }
  graph_->DeserializeParameters(model_file);
  return true;
}

bool Project::SaveModel() {
  std::string model_filename = project_folder_ + "/model.CNParamX";
  std::ofstream model_file(model_filename, std::ios::out | std::ios::binary);
  if(!model_file.good()) {
    LOGERROR << "Could not open " << model_filename << " for serializing the model!";
    return false;
  }
  graph_->SerializeParameters(model_file);
  return true;
}

bool Project::LoadModel() {
  needs_rescore_ = true;
  std::string model_filename = project_folder_ + "/model.CNParamX";
  std::ifstream model_file(model_filename, std::ios::in | std::ios::binary);
  if(!model_file.good()) {
    LOGERROR << "Could not open " << model_filename << " for deserializing the model!";
    return false;
  }
  graph_->DeserializeParameters(model_file);
  return true;
}

bool Project::RenameClass(const std::string &org_name, const std::string new_name) {
  // Rename classes
  bool class_result = class_manager_->RenameClass(org_name, new_name);
  if(!class_result) {
    LOGERROR << "Could not rename class!";
    return false;
  }

  // Rename samples
  bool sample_result = known_samples_->RenameClass(org_name, new_name);
  sample_result &= new_set_->RenameClass(org_name, new_name);
  sample_result &= update_set_->RenameClass(org_name, new_name);
  if(!sample_result) {
    LOGERROR << "Could not rename class in samples!";
    SetState(FAILED);
    return false;
  }

  return true;
}