/* +------------------------------------------------------------------------+
   |                     Mobile Robot Programming Toolkit (MRPT)            |
   |                          https://www.mrpt.org/                         |
   |                                                                        |
   | Copyright (c) 2005-2024, Individual contributors, see AUTHORS file     |
   | See: https://www.mrpt.org/Authors - All rights reserved.               |
   | Released under BSD License. See: https://www.mrpt.org/License          |
   +------------------------------------------------------------------------+ */

#include "obs-precomp.h"  // Precompiled headers
//
#include <mrpt/io/CFileGZInputStream.h>
#include <mrpt/io/CFileGZOutputStream.h>
#include <mrpt/io/CFileInputStream.h>
#include <mrpt/obs/CRawlog.h>
#include <mrpt/serialization/CArchive.h>
#include <mrpt/system/filesystem.h>

#include <iostream>

using namespace mrpt;
using namespace mrpt::io;
using namespace mrpt::obs;
using namespace mrpt::poses;
using namespace mrpt::system;
using namespace mrpt::serialization;

IMPLEMENTS_SERIALIZABLE(CRawlog, CSerializable, mrpt::obs)

void CRawlog::clear()
{
  m_seqOfActObs.clear();
  m_commentTexts.text.clear();
}

void CRawlog::insert(CSensoryFrame& observations)
{
  m_seqOfActObs.push_back(
      std::dynamic_pointer_cast<CSerializable>(observations.duplicateGetSmartPtr()));
}

void CRawlog::insert(CActionCollection& actions)
{
  m_seqOfActObs.push_back(std::dynamic_pointer_cast<CSerializable>(actions.duplicateGetSmartPtr()));
}
void CRawlog::insert(const CSerializable::Ptr& obj)
{
  if (IS_CLASS(*obj, CObservationComment))
  {
    CObservationComment::Ptr o = std::dynamic_pointer_cast<CObservationComment>(obj);
    m_commentTexts = *o;
  }
  else
    m_seqOfActObs.push_back(obj);
}

void CRawlog::insert(CAction& action)
{
  CActionCollection::Ptr temp = std::make_shared<CActionCollection>();
  temp->insert(action);
  m_seqOfActObs.push_back(temp);
}

size_t CRawlog::size() const { return m_seqOfActObs.size(); }

bool CRawlog::empty() const { return m_seqOfActObs.empty(); }

CActionCollection::Ptr CRawlog::getAsAction(size_t index) const
{
  MRPT_START

  if (index >= m_seqOfActObs.size()) THROW_EXCEPTION("Index out of bounds");

  CSerializable::Ptr obj = m_seqOfActObs[index];

  if (obj->GetRuntimeClass() == CLASS_ID(CActionCollection))
    return std::dynamic_pointer_cast<CActionCollection>(obj);
  else
    THROW_EXCEPTION_FMT("Element at index %i is not a CActionCollection", (int)index);
  MRPT_END
}

CObservation::Ptr CRawlog::getAsObservation(size_t index) const
{
  MRPT_START

  if (index >= m_seqOfActObs.size()) THROW_EXCEPTION("Index out of bounds");

  CSerializable::Ptr obj = m_seqOfActObs[index];

  if (obj->GetRuntimeClass()->derivedFrom(CLASS_ID(CObservation)))
    return std::dynamic_pointer_cast<CObservation>(obj);
  else
    THROW_EXCEPTION_FMT("Element at index %i is not a CObservation", (int)index);
  MRPT_END
}

CSerializable::Ptr CRawlog::getAsGeneric(size_t index) const
{
  MRPT_START
  if (index >= m_seqOfActObs.size()) THROW_EXCEPTION("Index out of bounds");

  return m_seqOfActObs[index];
  MRPT_END
}

CRawlog::TEntryType CRawlog::getType(size_t index) const
{
  MRPT_START
  if (index >= m_seqOfActObs.size()) THROW_EXCEPTION("Index out of bounds");

  const CSerializable::Ptr& obj = m_seqOfActObs[index];

  if (obj->GetRuntimeClass()->derivedFrom(CLASS_ID(CObservation)))
    return etObservation;
  else if (obj->GetRuntimeClass() == CLASS_ID(CActionCollection))
    return etActionCollection;
  else if (obj->GetRuntimeClass() == CLASS_ID(CSensoryFrame))
    return etSensoryFrame;
  else
    return etOther;

  MRPT_END
}

CSensoryFrame::Ptr CRawlog::getAsObservations(size_t index) const
{
  MRPT_START
  if (index >= m_seqOfActObs.size()) THROW_EXCEPTION("Index out of bounds");

  CSerializable::Ptr obj = m_seqOfActObs[index];

  if (obj->GetRuntimeClass()->derivedFrom(CLASS_ID(CSensoryFrame)))
    return std::dynamic_pointer_cast<CSensoryFrame>(obj);
  else
    THROW_EXCEPTION_FMT("Element at index %i is not a CSensoryFrame", (int)index);
  MRPT_END
}

uint8_t CRawlog::serializeGetVersion() const { return 1; }
void CRawlog::serializeTo(mrpt::serialization::CArchive& out) const
{
  out.WriteAs<uint32_t>(m_seqOfActObs.size());
  for (const auto& a : m_seqOfActObs) out << a;
  out << m_commentTexts;
}

void CRawlog::serializeFrom(mrpt::serialization::CArchive& in, uint8_t version)
{
  switch (version)
  {
    case 0:
    case 1:
    {
      clear();
      m_seqOfActObs.resize(in.ReadAs<uint32_t>());
      for (auto& a : m_seqOfActObs) a = in.ReadObject();
      in >> m_commentTexts;
    }
    break;
    default:
      MRPT_THROW_UNKNOWN_SERIALIZATION_VERSION(version);
  };
}

bool CRawlog::loadFromRawLogFile(const std::string& fileName, bool non_obs_objects_are_legal)
{
  // Open for read.
  CFileGZInputStream fi;
  if (!fi.open(fileName)) return false;
  auto fs = archiveFrom(fi);

  clear();  // Clear first

  // OK: read objects:
  bool keepReading = true;
  while (keepReading)
  {
    CSerializable::Ptr newObj;
    try
    {
      fs >> newObj;
      bool add_obj = false;
      // Check type:
      if (newObj->GetRuntimeClass() == CLASS_ID(CRawlog))
      {
        // It is an entire object: Copy and finish:
        CRawlog::Ptr ao = std::dynamic_pointer_cast<CRawlog>(newObj);
        this->swap(*ao);
        return true;
      }
      else if (newObj->GetRuntimeClass()->derivedFrom(CLASS_ID(CObservation)))
      {
        if (IS_CLASS(*newObj, CObservationComment))
        {
          CObservationComment::Ptr o = std::dynamic_pointer_cast<CObservationComment>(newObj);
          m_commentTexts = *o;
        }
        else
        {
          add_obj = true;
        }
      }
      else if (newObj->GetRuntimeClass() == CLASS_ID(CSensoryFrame))
      {
        add_obj = true;
      }
      else if (newObj->GetRuntimeClass() == CLASS_ID(CActionCollection))
      {
        add_obj = true;
      }
      else
      {
        // Other classes:
        if (non_obs_objects_are_legal)
        {
          add_obj = true;
        }
        else
        {
          keepReading = false;
        }
      }
      if (add_obj) m_seqOfActObs.push_back(newObj);
    }
    catch (CExceptionEOF&)
    {  // EOF, just finish the loop
      keepReading = false;
    }
    catch (const std::exception& e)
    {
      std::cerr << mrpt::exception_to_str(e) << std::endl;
      keepReading = false;
    }
    catch (...)
    {
      keepReading = false;
    }
  }
  return true;
}

void CRawlog::remove(size_t index)
{
  MRPT_START
  if (index >= m_seqOfActObs.size()) THROW_EXCEPTION("Index out of bounds");
  m_seqOfActObs.erase(m_seqOfActObs.begin() + index);
  MRPT_END
}

void CRawlog::remove(size_t first_index, size_t last_index)
{
  MRPT_START
  if (first_index >= m_seqOfActObs.size() || last_index >= m_seqOfActObs.size())
    THROW_EXCEPTION("Index out of bounds");
  m_seqOfActObs.erase(m_seqOfActObs.begin() + first_index, m_seqOfActObs.begin() + last_index + 1);
  MRPT_END
}

bool CRawlog::saveToRawLogFile(const std::string& fileName) const
{
  try
  {
    CFileGZOutputStream fo(fileName);
    auto f = archiveFrom(fo);
    if (!m_commentTexts.text.empty()) f << m_commentTexts;
    for (const auto& m_seqOfActOb : m_seqOfActObs) f << *m_seqOfActOb;
    return true;
  }
  catch (const std::exception& e)
  {
    std::cerr << mrpt::exception_to_str(e) << std::endl;
    return false;
  }
}

void CRawlog::swap(CRawlog& obj)
{
  if (this == &obj) return;
  m_seqOfActObs.swap(obj.m_seqOfActObs);
  std::swap(m_commentTexts, obj.m_commentTexts);
}

bool CRawlog::readActionObservationPair(
    CArchive& inStream,
    CActionCollection::Ptr& action,
    CSensoryFrame::Ptr& observations,
    size_t& rawlogEntry)
{
  try
  {
    // Load pose change from the rawlog:
    action.reset();
    while (!action)
    {
      CSerializable::Ptr obj;
      inStream >> obj;
      if (obj->GetRuntimeClass() == CLASS_ID(CActionCollection))
      {
        action = std::dynamic_pointer_cast<CActionCollection>(obj);
      }
      else
      {
        obj.reset();
      }
      rawlogEntry++;
    };

    // Load sensory frame from the rawlog:
    observations.reset();
    while (!observations)
    {
      CSerializable::Ptr obj;
      inStream >> obj;
      if (obj->GetRuntimeClass() == CLASS_ID(CSensoryFrame))
      {
        observations = std::dynamic_pointer_cast<CSensoryFrame>(obj);
      }
      else
      {
        obj.reset();
      }
      rawlogEntry++;
    }
    return true;
  }
  catch (CExceptionEOF&)
  {
    return false;
  }
  catch (const std::exception& e)
  {
    std::cerr << "[CRawlog::readActionObservationPair] Found exception:" << std::endl
              << mrpt::exception_to_str(e) << std::endl;
    return false;
  }
  catch (...)
  {
    std::cerr << "Untyped exception reading rawlog file!!" << std::endl;
    return false;
  }
}

bool CRawlog::getActionObservationPairOrObservation(
    CArchive& inStream,
    CActionCollection::Ptr& action,
    CSensoryFrame::Ptr& observations,
    CObservation::Ptr& observation,
    size_t& rawlogEntry)
{
  try
  {
    // Load pose change from the rawlog:
    observations.reset();
    observation.reset();
    action.reset();
    while (!action)
    {
      CSerializable::Ptr obj;
      inStream >> obj;
      if (IS_CLASS(*obj, CActionCollection))
      {
        action = std::dynamic_pointer_cast<CActionCollection>(obj);
      }
      else if (IS_DERIVED(*obj, CObservation))
      {
        observation = std::dynamic_pointer_cast<CObservation>(obj);
        rawlogEntry++;
        return true;
      }
      else
        obj.reset();
      rawlogEntry++;
    };

    // Load sensory frame from the rawlog:
    observations.reset();
    while (!observations)
    {
      CSerializable::Ptr obj;
      inStream >> obj;
      if (obj->GetRuntimeClass() == CLASS_ID(CSensoryFrame))
      {
        observations = std::dynamic_pointer_cast<CSensoryFrame>(obj);
      }
      rawlogEntry++;
    }
    return true;
  }
  catch (CExceptionEOF&)
  {
    return false;
  }
  catch (const std::exception& e)
  {
    std::cerr << "[CRawlog::readActionObservationPair] Found exception:" << std::endl
              << mrpt::exception_to_str(e) << std::endl;
    return false;
  }
  catch (...)
  {
    std::cerr << "Untyped exception reading rawlog file!!" << std::endl;
    return false;
  }
}

void CRawlog::findObservationsByClassInRange(
    mrpt::system::TTimeStamp time_start,
    mrpt::system::TTimeStamp time_end,
    const mrpt::rtti::TRuntimeClassId* class_type,
    TListTimeAndObservations& out_found,
    [[maybe_unused]] size_t guess_start_position) const
{
  MRPT_START

  out_found.clear();

  if (m_seqOfActObs.empty()) return;

  // Find the first appearance of time_start:
  // ---------------------------------------------------
  auto first = m_seqOfActObs.begin();
  const auto last = m_seqOfActObs.end();
  {
    // The following is based on lower_bound:
    size_t count, step;
    count = std::distance(first, last);
    while (count > 0)
    {
      auto it = first;
      step = count / 2;
      std::advance(it, step);

      // The comparison function:
      TTimeStamp this_timestamp;
      if ((*it)->GetRuntimeClass()->derivedFrom(CLASS_ID(CObservation)))
      {
        CObservation::Ptr o = std::dynamic_pointer_cast<CObservation>(*it);
        this_timestamp = o->timestamp;
        ASSERT_(this_timestamp != INVALID_TIMESTAMP);
      }
      else
        THROW_EXCEPTION("Element found which is not derived from CObservation");

      if (this_timestamp < time_start)  // *it < time_start
      {
        first = ++it;
        count -= step + 1;
      }
      else
        count = step;
    }
    // "first" is our guy
  }

  // Iterate until we get out of the time window:
  while (first != last)
  {
    TTimeStamp this_timestamp;
    if ((*first)->GetRuntimeClass()->derivedFrom(CLASS_ID(CObservation)))
    {
      CObservation::Ptr o = std::dynamic_pointer_cast<CObservation>(*first);
      this_timestamp = o->timestamp;
      ASSERT_(this_timestamp != INVALID_TIMESTAMP);

      if (this_timestamp < time_end)
      {
        if (o->GetRuntimeClass()->derivedFrom(class_type))
          out_found.insert(TTimeObservationPair(this_timestamp, o));
      }
      else
      {
        break;  // end of time window!
      }
    }
    else
      THROW_EXCEPTION("Element found which is not derived from CObservation");

    first++;
  }

  MRPT_END
}

bool CRawlog::getActionObservationPair(
    CActionCollection::Ptr& action, CSensoryFrame::Ptr& observations, size_t& rawlogEntry) const
{
  try
  {
    while (getType(rawlogEntry) != CRawlog::etActionCollection)
    {
      rawlogEntry++;
    }
    action = getAsAction(rawlogEntry++);

    while (getType(rawlogEntry) != CRawlog::etSensoryFrame)
    {
      rawlogEntry++;
    }
    observations = getAsObservations(rawlogEntry++);

    return true;
  }
  catch (const std::exception& e)
  {
    std::cerr << mrpt::exception_to_str(e) << std::endl;
    return false;
  }
  catch (...)
  {
    std::cerr << "Untyped exception getting act-obs pair from rawlog!!" << std::endl;
    return false;
  }
}

void CRawlog::getCommentText(std::string& t) const { t = m_commentTexts.text; }
std::string CRawlog::getCommentText() const { return m_commentTexts.text; }
void CRawlog::getCommentTextAsConfigFile(mrpt::config::CConfigFileMemory& memCfg) const
{
  memCfg.setContent(m_commentTexts.text);
}

void CRawlog::setCommentText(const std::string& t) { m_commentTexts.text = t; }

std::string CRawlog::detectImagesDirectory(const std::string& str)
{
  using namespace std::string_literals;
  using mrpt::system::pathJoin;

  const std::string rawlogPath = extractFileDirectory(str);

  std::vector<std::string> temptativeImgPaths = {
      pathJoin({rawlogPath, extractFileName(str) + "_Images"s}),
      pathJoin({rawlogPath, extractFileName(str) + "_images"s}),
      pathJoin({rawlogPath, extractFileName(str) + "_IMAGES"s}), pathJoin({rawlogPath, "Images"})};

  for (const auto& p : temptativeImgPaths)
    if (mrpt::system::directoryExists(p)) return p;

  // fallback:
  return temptativeImgPaths.back();
}
