/*
 * Copyright (C) 2012 Open Source Robotics Foundation
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */
/* Desc: A world state
 * Author: Nate Koenig
 */
#include <boost/algorithm/string.hpp>

#include "gazebo/common/Console.hh"
#include "gazebo/common/Exception.hh"
#include "gazebo/physics/World.hh"
#include "gazebo/physics/Model.hh"
#include "gazebo/physics/Light.hh"
#include "gazebo/physics/WorldState.hh"

using namespace gazebo;
using namespace physics;

// TODO added here for ABI compatibility
// move to class when merging forward
static std::string worldStateFilter;

/////////////////////////////////////////////////
WorldState::WorldState()
  : State()
{
  worldStateFilter = "";
}

/////////////////////////////////////////////////
WorldState::WorldState(const WorldPtr _world)
  : State(_world->Name(), _world->RealTime(), _world->SimTime(),
      _world->Iterations())
{
  this->world = _world;

  // Add a state for all the models
  Model_V models = _world->Models();
  for (Model_V::const_iterator iter = models.begin();
       iter != models.end(); ++iter)
  {
    this->modelStates.insert(std::make_pair((*iter)->GetName(),
          ModelState(*iter, this->realTime, this->simTime, this->iterations)));
  }

  // Add states for all the lights
  for (const auto &light : _world->Lights())
  {
    this->lightStates.insert(std::make_pair(light->GetName(),
          LightState(light, this->realTime, this->simTime, this->iterations)));
  }
}

/////////////////////////////////////////////////
WorldState::WorldState(const sdf::ElementPtr _sdf)
  : State()
{
  this->Load(_sdf);
}

/////////////////////////////////////////////////
WorldState::~WorldState()
{
  this->world.reset();
  this->modelStates.clear();
  this->lightStates.clear();
}

/////////////////////////////////////////////////
void WorldState::LoadWithFilter(const WorldPtr _world,
                                const std::string &_filter)
{
  worldStateFilter = _filter;
  this->Load(_world);
}

/////////////////////////////////////////////////
void WorldState::Load(const WorldPtr _world)
{
  this->world = _world;
  this->name = _world->Name();
  this->wallTime = common::Time::GetWallTime();
  this->simTime = _world->SimTime();
  this->realTime = _world->RealTime();
  this->iterations = _world->Iterations();
  this->insertions.clear();
  this->deletions.clear();

  std::string filter = worldStateFilter;
  std::list<std::string> mainParts, parts;
  boost::split(mainParts, filter, boost::is_any_of("/"));

  // Create the model filter
  if (!mainParts.empty())
  {
    boost::split(parts, mainParts.front(), boost::is_any_of("."));
    if (parts.empty() && !mainParts.front().empty())
      parts.push_back(mainParts.front());
  }
  std::list<std::string>::iterator partIter = parts.begin();

  // Add a state for all the models that match the filter
  Model_V models = _world->Models();
  for (Model_V::const_iterator iter = models.begin();
       iter != models.end(); ++iter)
  {
    bool add = true;

    // The first element in the filter must be a model name or a star.
    if (partIter != parts.end() && !parts.empty() &&
        !(*partIter).empty() && (*partIter) != "*")
    {
      std::string regexStr = *partIter;
      boost::replace_all(regexStr, "*", ".*");
      boost::regex regex(regexStr);
      add = boost::regex_match((*iter)->GetName(), regex);
    }

    if (add)
    {
      this->modelStates[(*iter)->GetName()].Load(*iter, this->realTime,
          this->simTime, this->iterations);
    }
  }

  // Remove models that no longer exist. We determine this by check the time
  // stamp on each model.
  for (ModelState_M::iterator iter = this->modelStates.begin();
       iter != this->modelStates.end();)
  {
    if (iter->second.GetRealTime() != this->realTime)
      this->modelStates.erase(iter++);
    else
      ++iter;
  }

  // Add states for all the lights
  this->lightStates.clear();
  Light_V lights = _world->Lights();
  for (const auto &light : lights)
  {
    this->lightStates[light->GetName()].Load(light, this->realTime,
        this->simTime, this->iterations);
  }
}

/////////////////////////////////////////////////
void WorldState::Load(const sdf::ElementPtr _elem)
{
  // Copy the name and time information
  this->name = _elem->Get<std::string>("world_name");
  this->simTime = _elem->Get<common::Time>("sim_time");
  this->wallTime = _elem->Get<common::Time>("wall_time");
  this->realTime = _elem->Get<common::Time>("real_time");
  this->iterations = _elem->Get<uint64_t>("iterations");

  // Add the model states
  this->modelStates.clear();
  if (_elem->HasElement("model"))
  {
    sdf::ElementPtr childElem = _elem->GetElement("model");

    while (childElem)
    {
      std::string modelName = childElem->Get<std::string>("name");
      this->modelStates.insert(std::make_pair(
            modelName, ModelState(childElem)));
      this->modelStates[modelName].SetSimTime(this->simTime);
      this->modelStates[modelName].SetWallTime(this->wallTime);
      this->modelStates[modelName].SetRealTime(this->realTime);
      this->modelStates[modelName].SetIterations(this->iterations);
      childElem = childElem->GetNextElement("model");
    }
  }

  // Add the light states
  this->lightStates.clear();
  if (_elem->HasElement("light"))
  {
    sdf::ElementPtr childElem = _elem->GetElement("light");

    while (childElem)
    {
      LightState lightState(childElem);
      lightState.SetSimTime(this->simTime);
      lightState.SetWallTime(this->wallTime);
      lightState.SetRealTime(this->realTime);
      lightState.SetIterations(this->iterations);

      std::string lightName = childElem->Get<std::string>("name");
      this->lightStates.insert(std::make_pair(
            lightName, lightState));

      childElem = childElem->GetNextElement("light");
    }
  }

  // Add insertions
  this->insertions.clear();
  if (_elem->HasElement("insertions"))
  {
    sdf::ElementPtr insertionsElem = _elem->GetElement("insertions");

    // Models
    if (insertionsElem->HasElement("model"))
    {
      sdf::ElementPtr modelElem = insertionsElem->GetElement("model");

      while (modelElem)
      {
        this->insertions.push_back(modelElem->ToString(""));
        modelElem = modelElem->GetNextElement("model");
      }
    }

    // Lights
    if (insertionsElem->HasElement("light"))
    {
      sdf::ElementPtr lightElem = insertionsElem->GetElement("light");

      while (lightElem)
      {
        this->insertions.push_back(lightElem->ToString(""));
        lightElem = lightElem->GetNextElement("light");
      }
    }
  }

  // Add deletions
  this->deletions.clear();
  if (_elem->HasElement("deletions"))
  {
    sdf::ElementPtr deletionsElem = _elem->GetElement("deletions");

    if (deletionsElem->HasElement("name"))
    {
      sdf::ElementPtr nameElem = deletionsElem->GetElement("name");

      while (nameElem)
      {
        this->deletions.push_back(nameElem->Get<std::string>());
        nameElem = nameElem->GetNextElement("name");
      }
    }
  }
}

/////////////////////////////////////////////////
void WorldState::SetWorld(const WorldPtr _world)
{
  this->world = _world;
}

/////////////////////////////////////////////////
const ModelState_M &WorldState::GetModelStates() const
{
  return this->modelStates;
}

/////////////////////////////////////////////////
const LightState_M &WorldState::LightStates() const
{
  return this->lightStates;
}

/////////////////////////////////////////////////
ModelState_M WorldState::GetModelStates(const boost::regex &_regex) const
{
  ModelState_M result;

  // Search for matching link names
  for (ModelState_M::const_iterator iter = this->modelStates.begin();
       iter != this->modelStates.end(); ++iter)
  {
    if (boost::regex_match(iter->first, _regex))
      result.insert(std::make_pair(iter->first, iter->second));
  }

  return result;
}

/////////////////////////////////////////////////
unsigned int WorldState::GetModelStateCount() const
{
  return this->modelStates.size();
}

/////////////////////////////////////////////////
unsigned int WorldState::LightStateCount() const
{
  return this->lightStates.size();
}

/////////////////////////////////////////////////
ModelState WorldState::GetModelState(const std::string &_modelName) const
{
  // Search for the model name
  ModelState_M::const_iterator iter = this->modelStates.find(_modelName);
  if (iter != this->modelStates.end())
    return iter->second;

  // Error if the model name doesn't exist.
  gzerr << "Invalid model name[" + _modelName + "]." << std::endl;
  return ModelState();
}

/////////////////////////////////////////////////
LightState WorldState::GetLightState(const std::string &_lightName) const
{
  // Search for the light name
  auto iter = this->lightStates.find(_lightName);
  if (iter != this->lightStates.end())
    return iter->second;

  // Error if the light name doesn't exist.
  gzerr << "Invalid light name[" + _lightName + "]." << std::endl;
  return LightState();
}

/////////////////////////////////////////////////
bool WorldState::HasModelState(const std::string &_modelName) const
{
  return this->modelStates.find(_modelName) != this->modelStates.end();
}

/////////////////////////////////////////////////
bool WorldState::HasLightState(const std::string &_lightName) const
{
  return this->lightStates.find(_lightName) != this->lightStates.end();
}

/////////////////////////////////////////////////
const std::vector<std::string> &WorldState::Insertions() const
{
  return this->insertions;
}

/////////////////////////////////////////////////
void WorldState::SetInsertions(const std::vector<std::string> &_insertions)
{
  static sdf::SDFPtr rootSDF = nullptr;
  if (rootSDF == nullptr)
  {
    rootSDF.reset(new sdf::SDF);
    sdf::initFile("root.sdf", rootSDF);
  }

  // Unwrap insertions from <sdf>
  for (const auto &insertion : _insertions)
  {
    rootSDF->Root()->ClearElements();
    // <sdf>
    if (sdf::readString(insertion, rootSDF))
    {
      // <model>
      if (rootSDF->Root()->HasElement("model"))
      {
        this->insertions.push_back(
            rootSDF->Root()->GetElement("model")->ToString(""));
      }
      // <light>
      else if (rootSDF->Root()->HasElement("light"))
      {
        this->insertions.push_back(
            rootSDF->Root()->GetElement("light")->ToString(""));
      }
      // <actor>
      else if (rootSDF->Root()->HasElement("actor"))
      {
        this->insertions.push_back(
            rootSDF->Root()->GetElement("actor")->ToString(""));
      }
      else
      {
        gzwarn << "Unsupported insertion [" << insertion << "]" << std::endl;
        continue;
      }
    }
    // Otherwise copy as-is without validating
    else
    {
      this->insertions.push_back(insertion);
    }
  }
}

/////////////////////////////////////////////////
const std::vector<std::string> &WorldState::Deletions() const
{
  return this->deletions;
}

/////////////////////////////////////////////////
void WorldState::SetDeletions(const std::vector<std::string> &_deletions)
{
  this->deletions = _deletions;
}

/////////////////////////////////////////////////
bool WorldState::IsZero() const
{
  bool result = this->insertions.size() == 0 && this->deletions.size() == 0;

  // Model
  for (ModelState_M::const_iterator iter = this->modelStates.begin();
       iter != this->modelStates.end() && result; ++iter)
  {
    result = result && iter->second.IsZero();
  }

  // Light
  for (const auto &light : this->lightStates)
  {
    result = result && light.second.IsZero();
  }

  return result;
}

/////////////////////////////////////////////////
WorldState &WorldState::operator=(const WorldState &_state)
{
  State::operator=(_state);

  // Clear the states
  this->modelStates.clear();
  this->lightStates.clear();

  this->insertions.clear();
  this->deletions.clear();

  // Copy the model states.
  for (ModelState_M::const_iterator iter =
       _state.modelStates.begin(); iter != _state.modelStates.end(); ++iter)
  {
    this->modelStates.insert(std::make_pair(iter->first,
          ModelState(iter->second)));
  }

  // Copy the light states.
  for (const auto &light : _state.lightStates)
  {
    this->lightStates.insert(std::make_pair(light.first,
          LightState(light.second)));
  }

  // Copy the insertions
  std::copy(_state.insertions.begin(),
            _state.insertions.end(), this->insertions.begin());

  // Copy the deletions
  std::copy(_state.deletions.begin(),
            _state.deletions.end(), this->deletions.begin());

  return *this;
}

/////////////////////////////////////////////////
WorldState WorldState::operator-(const WorldState &_state) const
{
  WorldState result;

  result.name = this->name;
  result.simTime = this->simTime;
  result.realTime = this->realTime;
  result.wallTime = this->wallTime;
  result.iterations = this->iterations;

  // Subtract the model states.
  for (ModelState_M::const_iterator iter =
       _state.modelStates.begin(); iter != _state.modelStates.end(); ++iter)
  {
    if (this->HasModelState(iter->second.GetName()))
    {
      ModelState state = this->GetModelState(iter->second.GetName()) -
        iter->second;

      if (!state.IsZero())
      {
        result.modelStates.insert(std::make_pair(state.GetName(), state));
      }
    }
    else
    {
      result.deletions.push_back(iter->second.GetName());
    }
  }

  // Subtract the light states.
  for (const auto &light : _state.lightStates)
  {
    if (this->HasLightState(light.second.GetName()))
    {
      LightState state = this->GetLightState(light.second.GetName()) -
          light.second;

      if (!state.IsZero())
      {
        result.lightStates.insert(std::make_pair(state.GetName(), state));
      }
    }
    else
    {
      result.deletions.push_back(light.second.GetName());
    }
  }

  // Add in the new model states
  for (const auto &modelState : this->modelStates)
  {
    if (!_state.HasModelState(modelState.second.GetName()) && this->world)
    {
      ModelPtr model = this->world->ModelByName(modelState.second.GetName());
      if (model)
        result.insertions.push_back(model->UnscaledSDF()->ToString(""));
    }
  }

  // Add in the new light states
  for (const auto &light : this->lightStates)
  {
    if (!_state.HasLightState(light.second.GetName()) && this->world)
    {
      LightPtr lightPtr = this->world->LightByName(light.second.GetName());
      result.insertions.push_back(lightPtr->GetSDF()->ToString(""));
    }
  }

  return result;
}

/////////////////////////////////////////////////
WorldState WorldState::operator+(const WorldState &_state) const
{
  WorldState result;

  result.name = this->name;
  result.simTime = this->simTime;
  result.realTime = this->realTime;
  result.wallTime = this->wallTime;
  result.iterations = this->iterations;

  // Add the model states.
  for (ModelState_M::const_iterator iter =
       _state.modelStates.begin(); iter != _state.modelStates.end(); ++iter)
  {
    ModelState state = this->GetModelState(iter->second.GetName()) +
      iter->second;
    result.modelStates.insert(std::make_pair(state.GetName(), state));
  }

  // Add the light states.
  for (const auto &light : _state.lightStates)
  {
    LightState state = this->GetLightState(light.second.GetName()) +
        light.second;
    result.lightStates.insert(std::make_pair(state.GetName(), state));
  }

  return result;
}

/////////////////////////////////////////////////
void WorldState::FillSDF(sdf::ElementPtr _sdf)
{
  _sdf->ClearElements();

  _sdf->GetAttribute("world_name")->Set(this->name);
  _sdf->GetElement("sim_time")->Set(this->simTime);
  _sdf->GetElement("real_time")->Set(this->realTime);
  _sdf->GetElement("wall_time")->Set(this->wallTime);
  _sdf->GetElement("iterations")->Set(this->iterations);

  // Models
  for (ModelState_M::iterator iter =
       this->modelStates.begin(); iter != this->modelStates.end(); ++iter)
  {
    sdf::ElementPtr elem = _sdf->AddElement("model");
    iter->second.FillSDF(elem);
  }

  // Lights
  for (auto &light : this->lightStates)
  {
    sdf::ElementPtr elem = _sdf->AddElement("light");
    light.second.FillSDF(elem);
  }
}

/////////////////////////////////////////////////
void WorldState::SetWallTime(const common::Time &_time)
{
  State::SetWallTime(_time);

  // Models
  for (ModelState_M::iterator iter = this->modelStates.begin();
       iter != this->modelStates.end(); ++iter)
  {
    iter->second.SetWallTime(_time);
  }

  // Lights
  for (auto &light : this->lightStates)
  {
    light.second.SetWallTime(_time);
  }
}

/////////////////////////////////////////////////
void WorldState::SetRealTime(const common::Time &_time)
{
  State::SetRealTime(_time);

  // Models
  for (ModelState_M::iterator iter = this->modelStates.begin();
           iter != this->modelStates.end(); ++iter)
  {
    iter->second.SetRealTime(_time);
  }

  // Lights
  for (auto &light : this->lightStates)
  {
    light.second.SetRealTime(_time);
  }
}

/////////////////////////////////////////////////
void WorldState::SetSimTime(const common::Time &_time)
{
  State::SetSimTime(_time);

  // Models
  for (ModelState_M::iterator iter = this->modelStates.begin();
       iter != this->modelStates.end(); ++iter)
  {
    iter->second.SetSimTime(_time);
  }

  // Lights
  for (auto &light : this->lightStates)
  {
    light.second.SetSimTime(_time);
  }
}

/////////////////////////////////////////////////
void WorldState::SetIterations(const uint64_t _iterations)
{
  State::SetIterations(_iterations);

  // Models
  for (auto &modelState : this->modelStates)
    modelState.second.SetIterations(_iterations);

  // Lights
  for (auto &lightState : this->lightStates)
    lightState.second.SetIterations(_iterations);
}
