/* Distributed under the Apache License, Version 2.0.
   See accompanying NOTICE file for details.*/

#include <iostream>
#include <fstream>
#include "LBM.h"
#include "RunUtils.h"
#include "ITKUtils.h"
#ifdef VTK_EXTENSION
  #include "VTKUtils.h"
#endif

int main(int argc, char* argv[])
{
  Logger log("./LBMITKDriver.log");
  // Specify which CUDA device to use, or leave black to automatically choose the largest device found.
  if (!cuda_device_check(log))
  {
    log.Error("CUDA device check failed, ensure you have specified a valid CUDA device id.");
    return 1;
  }

  LBM::RunConfig runset;
  std::string out_dir = "./itk_runs/";

  std::vector<LBM::RunIO> runs;
  if (argc >= 2)
  {
    for (int i = 1; i < argc; i++)
    {
      if (std::string(argv[i]).find(".set") != std::string::npos)
      { 
        std::string arg = runset.GetITKDir() + "/" + argv[i];
        log.Info("Processing : " + arg);
        std::string line;
        std::ifstream cFile(arg, std::ios::in);
        while (std::getline(cFile, line))
        {
          if (line.empty() || line.find("#") != std::string::npos)
            continue;// Skip a comment
          for (LBM::RunIO run : runset.GetITKDirRuns("/" + line))
            runs.push_back(run);
        }
      }
      else
      {// Must be a directory, look for files in it
        log.Info("Processing : " + std::string(argv[i]));
        for (LBM::RunIO run : runset.GetITKRuns(argv[i]))
          runs.push_back(run);
      }
    }
  }
  else
  {
    // You can set this string to be a specific run directory under the verification directory
    // This will run that specific run only, if you want to debug it
    std::string specific_run = "";//"/cylinder0";
    for (LBM::RunIO run : runset.GetITKDirRuns(specific_run))
      runs.push_back(run);
  }

  try
  {
    for (LBM::RunIO run : runs)
    {
      LBM::Run lbm;
      LBM::ITK::Configuration cfg;
      std::string name = run.output_base_path.substr(0, run.output_base_path.length());
      name = name.substr(name.find_last_of("/") + 1);
      log.Info("Running " + name);
      if (!LBM::ITK::LoadConfiguration(run.input, lbm, cfg))
      {
        log.Error("Unable to load configuration " + run.input);
        continue;
      }
      if (lbm.cfg.name.empty())
        lbm.cfg.name = name;

      // Provide a label mapping
      lbm.in.source_to_lbm_label_map[0]  = LBM::BoundaryTypes::Inactive;
      lbm.in.source_to_lbm_label_map[1]  = LBM::BoundaryTypes::Wall;   // Left Wall
      lbm.in.source_to_lbm_label_map[2]  = LBM::BoundaryTypes::Wall;   // Right Wall
      lbm.in.source_to_lbm_label_map[3]  = LBM::BoundaryTypes::Wall;   // Left Outlet Wall
      lbm.in.source_to_lbm_label_map[4]  = LBM::BoundaryTypes::Wall;   // Right Outlet Wall
      lbm.in.source_to_lbm_label_map[5]  = LBM::BoundaryTypes::Wall;   // Wall Behind Split
      lbm.in.source_to_lbm_label_map[11] = LBM::BoundaryTypes::Inlet;  // Left Inlet
      lbm.in.source_to_lbm_label_map[12] = LBM::BoundaryTypes::Inlet;  // Right Inlet
      lbm.in.source_to_lbm_label_map[20] = LBM::BoundaryTypes::Outlet;
      if (!LBM::ITK::Compute(cfg, lbm))
      {
        lbm.Error("Unable to run LBM on file : " + run.input);
        return 1;
      }
#ifdef VTK_EXTENSION
      lbm.Info("Writing out VTR file " + run.output_base_path +".vtr");
      LBM::VTK::WriteVTR(run.output_base_path +".vtr", lbm);
#endif
    }
  }
  catch (std::exception ex)
  {
    log.Error(ex.what());
    return 1;
  }

  return 0;
}
