//=========================================================================
//  Copyright (c) Kitware, Inc.
//  All rights reserved.
//  See LICENSE.txt for details.
//
//  This software is distributed WITHOUT ANY WARRANTY; without even
//  the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
//  PURPOSE.  See the above copyright notice for more information.
//=========================================================================

#include "smtk/attribute/Attribute.h"
#include "smtk/attribute/ComponentItem.h"
#include "smtk/attribute/DirectoryItem.h"
#include "smtk/attribute/DoubleItem.h"
#include "smtk/attribute/IntItem.h"
#include "smtk/attribute/ModelEntityItem.h"
#include "smtk/attribute/ResourceItem.h"
#include "smtk/attribute/StringItem.h"

#include "smtk/model/SessionIOJSON.h"
#include "smtk/model/json/jsonResource.h"

#include "smtk/session/rgg/Assembly.h"
#include "smtk/session/rgg/Core.h"
#include "smtk/session/rgg/Pin.h"
#include "smtk/session/rgg/Registrar.h"
#include "smtk/session/rgg/Resource.h"
#include "smtk/session/rgg/Session.h"

#include "smtk/session/rgg/operators/CreateModel.h"
#include "smtk/session/rgg/operators/EditAssembly.h"
#include "smtk/session/rgg/operators/EditCore.h"
#include "smtk/session/rgg/operators/EditDuct.h"
#include "smtk/session/rgg/operators/EditPin.h"
#include "smtk/session/rgg/operators/ExportInp.h"
#include "smtk/session/rgg/operators/Read.h"
#include "smtk/session/rgg/operators/Write.h"

#include "smtk/session/rgg/json/jsonAssembly.h"
#include "smtk/session/rgg/json/jsonCore.h"
#include "smtk/session/rgg/json/jsonDuct.h"
#include "smtk/session/rgg/json/jsonPin.h"

#include "smtk/model/EntityRef.h"
#include "smtk/model/Group.h"
#include "smtk/model/Model.h"

using namespace smtk::session::rgg;
using json = nlohmann::json;

//force to use filesystem version 3
#define BOOST_FILESYSTEM_VERSION 3
#include <boost/filesystem.hpp>
using namespace boost::filesystem;

namespace
{
std::string dataRoot = DATA_DIR;
std::string writeRoot = SCRATCH_DIR;

void cleanup(const std::string& file_path)
{
  //first verify the file exists
  ::boost::filesystem::path path(file_path);
  if (::boost::filesystem::is_regular_file(path))
  {
    //remove the file_path if it exists.
    ::boost::filesystem::remove(path);
  }
}

int createTargetPins(smtk::model::Model model)
{
  /// Create four pins
  auto editPinOp = smtk::session::rgg::EditPin::create();
  if (!editPinOp)
  {
    std::cerr << "No \"Edit Pin\" operator\n";
    return 1;
  }
  // Pin1-P1
  {
    Pin pin = Pin("Pin1", "P1");
    pin.setColor({0.258824, 0.572549, 0.776471, 1});
    pin.setCellMaterialIndex(0);
    std::vector<Pin::Piece> pinPieces;
    pinPieces.push_back(Pin::Piece(Pin::PieceType::CYLINDER, 100, 2, 2));
    pin.setPieces(pinPieces);
    std::vector<Pin::LayerMaterial> pinMaterials;
    pinMaterials.push_back(Pin::LayerMaterial(6,1));
    pin.setLayerMaterials(pinMaterials);

    nlohmann::json PinJson = pin;
    std::string targetPinStr = PinJson.dump();

    editPinOp->parameters()->associate(model.component());

    editPinOp->parameters()
      ->findString("pin representation")
      ->setValue(targetPinStr);

    auto createPinResult = editPinOp->operate();
    if (createPinResult->findInt("outcome")->value() !=
      static_cast<int>(smtk::operation::Operation::Outcome::SUCCEEDED))
    {
      std::cerr << "\"Edit Pin\" operator failed to create a pin\n";
      return 1;
    }
  }

  // ControlRod-CR
  {
    editPinOp->parameters()->removeAllAssociations();
    Pin pin = Pin("ControlRod", "CR");
    pin.setColor({0.254902, 0.670588, 0.364706, 1});
    pin.setCellMaterialIndex(0);
    std::vector<Pin::Piece> pinPieces;
    pinPieces.push_back(Pin::Piece(Pin::PieceType::CYLINDER, 100, 1, 1));
    pin.setPieces(pinPieces);
    std::vector<Pin::LayerMaterial> pinMaterials;
    pinMaterials.push_back(Pin::LayerMaterial(1,1));
    pin.setLayerMaterials(pinMaterials);

    nlohmann::json PinJson = pin;
    std::string targetPinStr = PinJson.dump();

    editPinOp->parameters()->associate(model.component());

    editPinOp->parameters()
      ->findString("pin representation")
      ->setValue(targetPinStr);

    auto createPinResult = editPinOp->operate();
    if (createPinResult->findInt("outcome")->value() !=
      static_cast<int>(smtk::operation::Operation::Outcome::SUCCEEDED))
    {
      std::cerr << "\"Edit Pin\" operator failed to create a pin\n";
      return 1;
    }
  }

  // FuelCell_1-F2
  {
    editPinOp->parameters()->removeAllAssociations();
    Pin pin = Pin("FuelCell_1", "F2");
    pin.setColor({0.45098, 0.45098, 0.45098, 1});
    pin.setCellMaterialIndex(0);
    std::vector<Pin::Piece> pinPieces;
    pinPieces.push_back(Pin::Piece(Pin::PieceType::CYLINDER, 100, 0.5, 0.5));
    pin.setPieces(pinPieces);
    std::vector<Pin::LayerMaterial> pinMaterials;
    pinMaterials.push_back(Pin::LayerMaterial(2,1));
    pin.setLayerMaterials(pinMaterials);

    nlohmann::json PinJson = pin;
    std::string targetPinStr = PinJson.dump();

    editPinOp->parameters()->associate(model.component());

    editPinOp->parameters()
      ->findString("pin representation")
      ->setValue(targetPinStr);

    auto createPinResult = editPinOp->operate();
    if (createPinResult->findInt("outcome")->value() !=
      static_cast<int>(smtk::operation::Operation::Outcome::SUCCEEDED))
    {
      std::cerr << "\"Edit Pin\" operator failed to create a pin\n";
      return 1;
    }
  }

  // FuelCell_2-F1
  {
    editPinOp->parameters()->removeAllAssociations();
    Pin pin = Pin("FuelCell_2", "F1");
    pin.setColor({0.45098, 0.45098, 0.45098, 1});
    pin.setCellMaterialIndex(0);
    std::vector<Pin::Piece> pinPieces;
    pinPieces.push_back(Pin::Piece(Pin::PieceType::CYLINDER, 100, 0.6, 0.6));
    pin.setPieces(pinPieces);
    std::vector<Pin::LayerMaterial> pinMaterials;
    pinMaterials.push_back(Pin::LayerMaterial(2,0.66666));
    pinMaterials.push_back(Pin::LayerMaterial(5,1));
    pin.setLayerMaterials(pinMaterials);

    nlohmann::json PinJson = pin;
    std::string targetPinStr = PinJson.dump();

    editPinOp->parameters()->associate(model.component());

    editPinOp->parameters()
      ->findString("pin representation")
      ->setValue(targetPinStr);

    auto createPinResult = editPinOp->operate();
    if (createPinResult->findInt("outcome")->value() !=
      static_cast<int>(smtk::operation::Operation::Outcome::SUCCEEDED))
    {
      std::cerr << "\"Edit Pin\" operator failed to create a pin\n";
      return 1;
    }
  }

  return 0;
}


int createTargetDucts(smtk::model::Model model)
{
  auto editDuctOp = smtk::session::rgg::EditDuct::create();
  if(!editDuctOp)
  {
    std::cerr << "No \"Edit Duct\" operator\n";
    return 1;
  }

  // Create Duct C1_Duct
  {
    Duct targetDuct = Duct("C1_Duct", false);
    Duct::Segment initSeg;
    initSeg.baseZ = 0.0;
    initSeg.height = 100.0;
    initSeg.layers.push_back(std::make_tuple(7, 1.0, 1.0));
    targetDuct.setSegments({initSeg});
    nlohmann::json targetDuctJson = targetDuct;
    std::string targetDuctStr = targetDuctJson.dump();

    editDuctOp->parameters()->associate(model.component());
    editDuctOp->parameters()->findString("duct representation")->setValue(targetDuctStr);

    auto createDuctResult = editDuctOp->operate();

    if (createDuctResult->findInt("outcome")->value() !=
      static_cast<int>(smtk::operation::Operation::Outcome::SUCCEEDED))
    {
      std::cerr << "\"Edit Duct\" operator failed to create a duct\n";
      return 1;
    }
  }

  // Create Duct FF_Duct
  {
    editDuctOp->parameters()->removeAllAssociations();
    Duct targetDuct = Duct("FF_Duct", false);
    Duct::Segment initSeg;
    initSeg.baseZ = 0.0;
    initSeg.height = 100.0;
    initSeg.layers.push_back(std::make_tuple(3, 0.9, 0.9));
    initSeg.layers.push_back(std::make_tuple(4, 1.0, 1.0));
    targetDuct.setSegments({initSeg});
    nlohmann::json targetDuctJson = targetDuct;
    std::string targetDuctStr = targetDuctJson.dump();

    editDuctOp->parameters()->associate(model.component());
    editDuctOp->parameters()->findString("duct representation")->setValue(targetDuctStr);

    auto createDuctResult = editDuctOp->operate();

    if (createDuctResult->findInt("outcome")->value() !=
      static_cast<int>(smtk::operation::Operation::Outcome::SUCCEEDED))
    {
      std::cerr << "\"Edit Duct\" operator failed to create a duct\n";
      return 1;
    }
  }

  return 0;
}

int createTargetAssemblies(smtk::model::Model model)
{
  auto editAssyOp = smtk::session::rgg::EditAssembly::create();
  if (!editAssyOp)
  {
    std::cerr << "NO \"Edit Assembly\" operator\n";
    return 1;
  }

  smtk::model::EntityRef pinP1 = model.resource()->findEntitiesByProperty("label", "P1")[0];
  smtk::model::EntityRef pinCR = model.resource()->findEntitiesByProperty("label", "CR")[0];
  smtk::model::EntityRef pinF2 = model.resource()->findEntitiesByProperty("label", "F2")[0];
  smtk::model::EntityRef pinF1 = model.resource()->findEntitiesByProperty("label", "F1")[0];

  smtk::model::EntityRef ductC1 = model.resource()->findEntitiesByProperty("name", "C1_Duct")[0];
  smtk::model::EntityRef ductFF = model.resource()->findEntitiesByProperty("name", "FF_Duct")[0];

  // Create assembly control-C1
  {
    Assembly assy = Assembly("control", "C1");
    assy.setAssociatedDuct(ductC1.entity());

    // modify the export parameters
    AssyExportParameters aep;
    aep.MeshType = "Hex"; // hex or Tet
    aep.GeometryType = "Hexagonal"; // Hexagonal
    aep.Geometry = "Volume";
    aep.MergeTolerance = 1e-3;
    aep.RadialMeshSize = 0.3;
    aep.AxialMeshSize = 20.0;
    aep.MeshScheme = "hole";
    aep.EdgeInterval = 14;
    aep.CreateSideset = "NO";
    aep.Rotate = "Z 30.0";
    aep.MaterialSet_StartId = 100;
    aep.NeumannSet_StartId = 100;
    assy.setExportParams(aep);

    assy.setColor({0.258824, 0.572549, 0.776471, 1});
    assy.setCenterPin(false);
    assy.setPitch(10,10);
    Assembly::UuidToSchema layout;
    layout[pinP1.entity().toString()] = {std::make_pair(0,0)};
    assy.setLayout(layout);
    assy.setLatticeSize(1);

    nlohmann::json assyJson = assy;
    std::string assyStr = assyJson.dump();

    editAssyOp->parameters()->associate(model.component());

    editAssyOp->parameters()->findString("assembly representation")->setValue(assyStr);

    editAssyOp->parameters()->findComponent("associated duct")->setValue(ductC1.component());

    auto createAssyOpResult = editAssyOp->operate();
    if (createAssyOpResult->findInt("outcome")->value() !=
        static_cast<int>(smtk::operation::Operation::Outcome::SUCCEEDED))
    {
      std::cerr << "\"Create Assembly\" operator failed\n";
      return 1;
    }
  }

  // Create assembly fuel-FF
  {
    editAssyOp->parameters()->removeAllAssociations();
    Assembly assy = Assembly("fuel", "FF");
    assy.setAssociatedDuct(ductFF.entity());

    // modify the export parameters
    AssyExportParameters aep;
    aep.MeshType = "Hex"; // hex or Tet
    aep.GeometryType = "Hexagonal"; // Hexagonal
    aep.Geometry = "Volume";
    aep.MergeTolerance = 1e-3;
    aep.RadialMeshSize = 0.3;
    aep.AxialMeshSize = 20.0;
    aep.EdgeInterval = 14;
    aep.CreateSideset = "NO";
    aep.Rotate = "Z -30.0";
    aep.MaterialSet_StartId = -1;
    aep.NeumannSet_StartId = -1;
    assy.setExportParams(aep);

    assy.setColor({0.945098, 0.411765, 0.0745098, 1});
    assy.setCenterPin(false);
    assy.setPitch(1.8796,1.8796);
    Assembly::UuidToSchema layout;
    layout[pinCR.entity().toString()] = {std::make_pair(0,0)};
    layout[pinF1.entity().toString()] = {std::make_pair(1,0), std::make_pair(1,1), std::make_pair(1, 2),
        std::make_pair(1,3), std::make_pair(1,4), std::make_pair(1, 5)};
    layout[pinF2.entity().toString()] = {std::make_pair(2,0), std::make_pair(2, 1), std::make_pair(2,2),
        std::make_pair(2,3), std::make_pair(2,4), std::make_pair(2, 5),
        std::make_pair(2,6), std::make_pair(2,7), std::make_pair(2, 8),
        std::make_pair(2,9), std::make_pair(2,10), std::make_pair(2, 11)};
    assy.setLayout(layout);
    assy.setLatticeSize(3);

    nlohmann::json assyJson = assy;
    std::string assyStr = assyJson.dump();

    editAssyOp->parameters()->associate(model.component());

    editAssyOp->parameters()->findString("assembly representation")->setValue(assyStr);

    editAssyOp->parameters()->findComponent("associated duct")->setValue(ductC1.component());

    auto createAssyOpResult = editAssyOp->operate();
    if (createAssyOpResult->findInt("outcome")->value() !=
        static_cast<int>(smtk::operation::Operation::Outcome::SUCCEEDED))
    {
      std::cerr << "\"Create Assembly\" operator failed\n";
      return 1;
    }
  }
  return 0;
}

int editTargetCore(smtk::model::Model model)
{
  auto editCoreOp = smtk::session::rgg::EditCore::create();
  if (!editCoreOp)
  {
    std::cerr << "NO \"Edit Assembly\" operator\n";
    return 1;
  }

  smtk::model::EntityRef coreG = model.resource()->findEntitiesByProperty("rggType", Core::typeDescription)[0];
  smtk::model::EntityRef assyGC1 = model.resource()->findEntitiesByProperty("label", "C1")[0];
  smtk::model::EntityRef assyGFF = model.resource()->findEntitiesByProperty("label", "FF")[0];

  Core core = nlohmann::json::parse(coreG.stringProperty(Core::propDescription)[0]);
  core.setLatticeSize(3);

  CoreExportParameters cep;
  cep.Geometry = "Volume";
  cep.Symmetry = "1";
  cep.GeometryType = "hexflat";
  cep.NeumannSet.push_back("Top 97");
  cep.NeumannSet.push_back("Bot 98");
  cep.NeumannSet.push_back("SideAll 99");
  cep.BackgroundInfo.first = "outer_cylinder.cub";
  cep.BackgroundInfo.second = dataRoot + "/outer_cylinder.cub";
  cep.OutputFileName = "shfc.h5m";
  core.setExportParams(cep);

  Core::UuidToSchema layout;
  layout[assyGFF.entity().toString()] = {std::make_pair(1,0), std::make_pair(1,1), std::make_pair(1, 2),
      std::make_pair(1,3), std::make_pair(1,4), std::make_pair(1, 5)};
  layout[assyGC1.entity().toString()] = {std::make_pair(0,0), std::make_pair(2,0), std::make_pair(2, 1), std::make_pair(2,2),
      std::make_pair(2,3), std::make_pair(2,4), std::make_pair(2, 5),
      std::make_pair(2,6), std::make_pair(2,7), std::make_pair(2, 8),
      std::make_pair(2,9), std::make_pair(2,10), std::make_pair(2, 11)};
  core.setLayout(layout);

  editCoreOp->parameters()->associate(coreG.component());
  nlohmann::json coreJson = core;
  std::string coreJsonStr = coreJson.dump();
  editCoreOp->parameters()->findString("core representation")->setValue(coreJsonStr);

  auto editCoreOpResult = editCoreOp->operate();
  if (editCoreOpResult->findInt("outcome")->value() !=
      static_cast<int>(smtk::operation::Operation::Outcome::SUCCEEDED))
  {
    std::cerr << "\"Edit core\" operator failed\n";
    return 1;
  }

  return 0;
}


int writeAndReadRGGModel(smtk::model::Model model)
{
  auto writeOp = smtk::session::rgg::Write::create();
  if (!writeOp)
  {
    std::cerr << "NO \"Write \" operator\n";
    return 1;
  }
  writeOp->parameters()->associate(model.resource());

  std::string testDirStr = writeRoot + "/" + smtk::common::UUID::random().toString();
  ::boost::filesystem::path testDir(testDirStr);
  bool created = boost::filesystem::create_directory(testDir);
  if (created == false)
  {
    std::cerr << "Could not create temporary directory \"" << testDirStr << "\"\n";
    return 1;
  }
  std::string fileName = testDirStr +"/rggModel.smtk";
  model.resource()->setLocation(fileName);
  std::cout << "testDirStr=" << testDirStr << std::endl;

  auto writeOpResult = writeOp->operate();


  if (writeOpResult->findInt("outcome")->value() !=
      static_cast<int>(smtk::operation::Operation::Outcome::SUCCEEDED))
  {
    std::cerr << "\"Write \" operator failed\n";
    return 1;
  }

  smtk::model::Resource::Ptr resourceRead = std::dynamic_pointer_cast<smtk::model::Resource>
      (smtk::session::rgg::read(fileName));
  smtk::model::SessionIOJSON::json resourceReadJson = resourceRead;
  smtk::model::SessionIOJSON::json resourceOriginalJson = model.resource();

  // Nlohmann json serialization does not gurarantee order. Comment it out for now.
  //assert (resourceReadJson == resourceOriginalJson);

  ::boost::filesystem::remove_all(testDir);

  return 0;
}
}

int TestRGGWriteRead(int argc, char* argv[])
{
  // Create a resource manager
  smtk::resource::Manager::Ptr resourceManager = smtk::resource::Manager::create();

  // Register rgg resources to the resource manager
  {
    smtk::session::rgg::Registrar::registerTo(resourceManager);
  }

  // Create an operation manager
  smtk::operation::Manager::Ptr operationManager = smtk::operation::Manager::create();

  // Register rgg operators to the operation manager
  {
    smtk::session::rgg::Registrar::registerTo(operationManager);
  }

  // Register the resource manager to the operation manager (newly created
  // resources will be automatically registered to the resource manager).
  operationManager->registerResourceManager(resourceManager);

  smtk::model::Model model;
  smtk::resource::ResourcePtr resource;
  auto createModelOp = smtk::session::rgg::CreateModel::create();
  if (!createModelOp)
  {
    std::cerr << "No create model operator\n";
    return 1;
  }

  // Create a hex core
  createModelOp->parameters()->findString("name")->setValue("simpleHexCore");
  createModelOp->parameters()->findString("geometry type")->setDiscreteIndex(0);
  createModelOp->parameters()->findDouble("z origin")->setValue(0);
  createModelOp->parameters()->findDouble("height")->setValue(100);
  createModelOp->parameters()->findDouble("duct thickness")->setValue(10);
  createModelOp->parameters()->findInt("hex lattice size")->setValue(3);

  auto createModelOpResult = createModelOp->operate();
  if (createModelOpResult->findInt("outcome")->value() !=
    static_cast<int>(smtk::operation::Operation::Outcome::SUCCEEDED))
  {
    std::cerr << "create model operator failed\n";
    return 1;
  }

  // Use a shared pointer to track the created resource so that it's valid out side of this scope
  resource =
    createModelOpResult->findResource("resource")->value();

  model =
    createModelOpResult->findComponent("created")->valueAs<smtk::model::Entity>();
  if (!model.isValid())
  {
    std::cerr << "create model operator constructed an invalid model\n";
    return 1;
  }

  // Set material and its color
  std::vector<std::string> materials{"NoCellMaterial", "CR", "Fuel", "Dt1", "Dt2", "MatH",
                                    "Mat_Control", "Mat_Coolant"};
  model.setStringProperty("materials", materials);
  model.setFloatProperty("NoCellMaterial",{ 1.0, 1.0, 1.0, 1.0 });
  model.setFloatProperty("CR",{ 0.3, 0.5, 1.0, .5 });
  model.setFloatProperty("Fuel",{ 0.3, 0.3, 1.0, .5 });
  model.setFloatProperty("Dt1",{ 0.75, 0.2, 0.75, 1.0});
  model.setFloatProperty("Dt2",{ 1.0, 0.1, 0.1, 1.0 });
  model.setFloatProperty("MatH",{ 0.0, 0.0, 0.0, 0.0 });
  model.setFloatProperty("Mat_Control",{ 0.3, 1.0, 0.5, 1.0 });
  model.setFloatProperty("Mat_Coolant",{ .4, .4, .4, 1.0 });

  smtk::model::Group coreGroup = createModelOpResult->findComponent("created")->
      valueAs<smtk::model::Entity>(1);
  if (!coreGroup.isValid())
  {
    std::cerr << "create model operator constructed an invliad core group\n";
    return 1;
  }
  if (!coreGroup.hasStringProperty(Core::propDescription))
  {
    std::cerr << "The created core does not have json representation string property\n";
    return 1;
  }
  Core core = nlohmann::json::parse(coreGroup.stringProperty(Core::propDescription)[0]);

  if (createTargetPins(model))
  {
    std::cerr << "Fail to create the target pins" <<std::endl;
    return 1;
  }

  assert(model.resource()->findEntitiesByProperty("label", "P1").size() > 0);
  smtk::model::EntityRef pinP1 = model.resource()->findEntitiesByProperty("label", "P1")[0];
  assert(model.resource()->findEntitiesByProperty("label", "CR").size() > 0);
  smtk::model::EntityRef pinCR = model.resource()->findEntitiesByProperty("label", "CR")[0];
  assert(model.resource()->findEntitiesByProperty("label", "F2").size() > 0);
  smtk::model::EntityRef pinF2 = model.resource()->findEntitiesByProperty("label", "F2")[0];
  assert(model.resource()->findEntitiesByProperty("label", "F1").size() > 0);
  smtk::model::EntityRef pinF1 = model.resource()->findEntitiesByProperty("label", "F1")[0];

  if(createTargetDucts(model))
  {
    std::cerr << "Fail to create the target ducts" <<std::endl;
    return 1;
  }
  assert(model.resource()->findEntitiesByProperty("name", "C1_Duct").size() > 0);
  smtk::model::EntityRef ductC1 = model.resource()->findEntitiesByProperty("name", "C1_Duct")[0];
  assert(model.resource()->findEntitiesByProperty("name", "FF_Duct").size() > 0);
  smtk::model::EntityRef ductFF = model.resource()->findEntitiesByProperty("name", "FF_Duct")[0];

  if (createTargetAssemblies(model))
  {
    std::cerr << "Fail to create the target assemblies" <<std::endl;
    return 1;
  }
  assert(model.resource()->findEntitiesByProperty("rggType", Assembly::typeDescription).size() == 2);

  if (editTargetCore(model))
  {
    std::cerr << "Fail to edit the target core" <<std::endl;
    return 1;
  }

  if (writeAndReadRGGModel(model))
  {
    std::cerr << "Fail to write and read rgg model" <<std::endl;
    return 1;
  }

  return 0;
}
