/* Copyright (C) 2022-2023 GSI Helmholtzzentrum fuer Schwerionenforschung, Darmstadt
   SPDX-License-Identifier: GPL-3.0-only
   Authors: Sergey Gorbunov, Sergei Zharko [committer] */

/// \file   CaStationInitializer.cxx
/// \brief
/// \author S.Zharko <s.zharko@gsi.de>
/// \since  18.01.2022
///

#include "CaStationInitializer.h"

#include "AlgoFairloggerCompat.h"
#include "CaDefs.h"

#include <iomanip>
#include <sstream>
#include <utility>

using cbm::algo::ca::EDetectorID;
using cbm::algo::ca::fvec;
using cbm::algo::ca::MaterialMap;
using cbm::algo::ca::Station;
using cbm::algo::ca::StationInitializer;

// ---------------------------------------------------------------------------------------------------------------------
//
StationInitializer::StationInitializer(EDetectorID detectorID, int stationID) noexcept
  : fDetectorID(detectorID)
  , fStationID(stationID)
{
  fInitController.SetFlag(EInitKey::kDetectorID);
  fInitController.SetFlag(EInitKey::kStationID);
}

// ---------------------------------------------------------------------------------------------------------------------
//
void StationInitializer::Reset()
{
  StationInitializer other;
  this->Swap(other);
}

// ---------------------------------------------------------------------------------------------------------------------
//
const Station<fvec>& StationInitializer::GetStation() const
{
  if (!fInitController.IsFinalized()) {
    std::stringstream msg;
    msg << "StationInitializer::GetStation: attempt to get a Station object from uninitialized L1BaseStation with "
        << "stationID = " << fStationID << " and detectorID = " << static_cast<int>(fDetectorID);
    LOG(fatal) << msg.str();
  }
  return fStation;
}

// ---------------------------------------------------------------------------------------------------------------------
//
const MaterialMap& StationInitializer::GetMaterialMap() const
{
  if (!fInitController.IsFinalized()) {
    LOG(fatal) << "StationInitializer::GetMaterialMap: attempt of getting the material map object from uninitialized "
                  "L1BaseStation class (detectorID = "
               << static_cast<int>(fDetectorID) << ", stationID = " << fStationID << ")";
  }

  if (fManagementFlags[static_cast<int>(EManagementFlag::kThicknessMapMoved)]) {
    LOG(fatal) << "StationInitializer::GetMaterialMap: attempt of getting the material map, which has been moved. The "
                  "thickness map instance have been "
               << "already took from the StationInitializer object (detectorID = " << static_cast<int>(fDetectorID)
               << ", stationID = " << fStationID << ")";
  }

  return fThicknessMap;
}

// ---------------------------------------------------------------------------------------------------------------------
//
void StationInitializer::SetDetectorID(EDetectorID inID)
{
  if (!fInitController.GetFlag(EInitKey::kDetectorID)) {
    fDetectorID = inID;
    fInitController.SetFlag(EInitKey::kDetectorID);
  }
  else {
    LOG(warn) << "StationInitializer::SetDetectorID: Attempt of detector ID redifinition";
  }
}


// ---------------------------------------------------------------------------------------------------------------------
//
void StationInitializer::SetFieldFunction(
  const std::function<void(const double (&xyz)[3], double (&B)[3])>& getFieldValue)
{
  if (fInitController.GetFlag(EInitKey::kFieldSlice)) {
    LOG(warn) << "StationInitializer::SetFieldSlice: Attempt to redifine field slice for station with detectorID = "
              << static_cast<int>(fDetectorID) << " and stationID = " << fStationID << ". Redifinition ignored";
    return;
  }

  if (!fInitController.GetFlag(EInitKey::kZref)) {
    LOG(fatal) << "Attempt to set magnetic field slice before setting z position of the station";
  }
  if (!fInitController.GetFlag(EInitKey::kXmax)) {
    LOG(fatal) << "Attempt to set magnetic field slice before Xmax size of the station";
  }
  if (!fInitController.GetFlag(EInitKey::kYmax)) {
    LOG(fatal) << "Attempt to set magnetic field slice before Ymax size of the station";
  }
  // TODO: Change names of variables according to convention (S.Zh.)
  constexpr int M = constants::size::MaxFieldApproxPolynomialOrder;
  constexpr int N = constants::size::MaxNFieldApproxCoefficients;
  constexpr int D = 3;  ///> number of dimensions

  // SLE initialization
  double A[N][N + D] = {};  // augmented matrix
  double dx          = (fXmax / N / 2 < 1.) ? fXmax / N / 4. : 1.;
  double dy          = (fYmax / N / 2 < 1.) ? fYmax / N / 4. : 1.;

  for (double x = -fXmax; x <= fXmax; x += dx) {
    for (double y = -fYmax; y <= fYmax; y += dy) {
      double r = sqrt(fabs(x * x / fXmax / fXmax + y / fYmax * y / fYmax));
      if (r > 1.) {
        continue;
      }
      double p[D] = {x, y, fZref};
      double B[D] = {};
      getFieldValue(p, B);

      double m[N] = {1};
      m[0]        = 1;
      for (int i = 1; i <= M; ++i) {
        int k = (i - 1) * i / 2;
        int l = i * (i + 1) / 2;
        for (int j = 0; j < i; ++j) {
          m[l + j] = x * m[k + j];
        }
        m[l + i] = y * m[k + i - 1];
      }

      double w = 1. / (r * r + 1);
      for (int i = 0; i < N; ++i) {
        // fill the left part of the matrix
        for (int j = 0; j < N; ++j) {
          A[i][j] += w * m[i] * m[j];
        }
        // fill the right part of the matrix
        for (int j = 0; j < D; ++j) {
          A[i][N + j] += w * B[j] * m[i];
        }
      }
    }
  }

  // SLE solution (Gaussian elimination)
  //
  for (int kCol = 0; kCol < N - 1; ++kCol) {
    for (int jRow = kCol + 1; jRow < N; ++jRow) {
      double factor = A[jRow][kCol] / A[kCol][kCol];
      for (int iCol = kCol; iCol < N + D; ++iCol) {
        A[jRow][iCol] -= factor * A[kCol][iCol];
      }
    }
  }
  for (int kCol = N - 1; kCol > 0; --kCol) {
    for (int jRow = kCol - 1; jRow >= 0; --jRow) {
      double factor = A[jRow][kCol] / A[kCol][kCol];
      for (int iCol = kCol; iCol < N + D; ++iCol) {
        A[jRow][iCol] -= factor * A[kCol][iCol];
      }
    }
  }

  for (int j = 0; j < N; ++j) {
    fStation.fieldSlice.cx[j] = A[j][N + 0] / A[j][j];
    fStation.fieldSlice.cy[j] = A[j][N + 1] / A[j][j];
    fStation.fieldSlice.cz[j] = A[j][N + 2] / A[j][j];
  }
  fStation.fieldSlice.z = fZref;

  fInitController.SetFlag(EInitKey::kFieldSlice);
}

// ---------------------------------------------------------------------------------------------------------------------
//
void StationInitializer::SetFieldStatus(int fieldStatus)
{
  fStation.fieldStatus = fieldStatus;
  fInitController.SetFlag(EInitKey::kFieldStatus);
}

// ---------------------------------------------------------------------------------------------------------------------
//
void StationInitializer::SetMaterialMap(const MaterialMap& thicknessMap)
{
  if (!fInitController.GetFlag(EInitKey::kThicknessMap)) {
    fThicknessMap = thicknessMap;
    fInitController.SetFlag(EInitKey::kThicknessMap);
  }
  else {
    LOG(warn) << "StationInitializer::SetMaterialMap: attempt to reinitialize the material map";
  }
}

// ---------------------------------------------------------------------------------------------------------------------
//
void StationInitializer::SetMaterialMap(MaterialMap&& thicknessMap) noexcept
{
  if (!fInitController.GetFlag(EInitKey::kThicknessMap)) {
    fThicknessMap = std::move(thicknessMap);
    fInitController.SetFlag(EInitKey::kThicknessMap);
  }
  else {
    LOG(warn) << "StationInitializer::SetMaterialMap: attempt to reinitialize the material map";
  }
}

// ---------------------------------------------------------------------------------------------------------------------
//
void StationInitializer::SetStationID(int inID)
{
  if (!fInitController.GetFlag(EInitKey::kStationID)) {
    fStationID = inID;
    fInitController.SetFlag(EInitKey::kStationID);
  }
  else {
    LOG(warn) << "StationInitializer::SetStationID: Attempt of station ID redifinition";
  }
}

// ---------------------------------------------------------------------------------------------------------------------
//
void StationInitializer::SetStationType(int inType)
{
  if (!fInitController.GetFlag(EInitKey::kType)) {
    fStation.type = inType;
    fInitController.SetFlag(EInitKey::kType);
  }
  else {
    LOG(warn) << "StationInitializer::SetStationType: Attempt of station type redifinition";
  }
}

// ---------------------------------------------------------------------------------------------------------------------
//
void StationInitializer::SetXmax(double aSize)
{
  fXmax         = aSize;
  fStation.Xmax = aSize;
  fInitController.SetFlag(EInitKey::kXmax);
}

// ---------------------------------------------------------------------------------------------------------------------
//
void StationInitializer::SetYmax(double aSize)
{
  fYmax         = aSize;
  fStation.Ymax = aSize;
  fInitController.SetFlag(EInitKey::kYmax);
}

// ---------------------------------------------------------------------------------------------------------------------
//
void StationInitializer::SetTimeInfo(int inTimeInfo)
{
  fStation.timeInfo = inTimeInfo;
  fInitController.SetFlag(EInitKey::kTimeInfo);
}

// ---------------------------------------------------------------------------------------------------------------------
//
void StationInitializer::SetTrackingStatus(bool flag)
{
  fTrackingStatus = flag;
  fInitController.SetFlag(EInitKey::kTrackingStatus);
  if (!fTrackingStatus && !fInitController.GetFlag(EInitKey::kThicknessMap)) {
    // Set initialized status for the material map, if the station is not active, since its material is accounted in the
    // previous tracking station.
    fInitController.SetFlag(EInitKey::kThicknessMap);
  }
}

// ---------------------------------------------------------------------------------------------------------------------
//
void StationInitializer::SetZref(double inZ)
{
  fStation.fZ = inZ;  // setting simd vector of single-precision floats, which is passed to high performanced L1Algo
  fZref       = inZ;  // setting precised value to use in field approximation etc
  fInitController.SetFlag(EInitKey::kZref);
}

// ---------------------------------------------------------------------------------------------------------------------
//
void StationInitializer::SetZmin(double inZ)
{
  fZmin = inZ;  // setting precised value to use in field approximation etc
  fInitController.SetFlag(EInitKey::kZmin);
}

// ---------------------------------------------------------------------------------------------------------------------
//
void StationInitializer::SetZmax(double inZ)
{
  fZmax = inZ;  // setting precised value to use in field approximation etc
  fInitController.SetFlag(EInitKey::kZmax);
}

// ---------------------------------------------------------------------------------------------------------------------
//
void StationInitializer::Swap(StationInitializer& other) noexcept
{
  std::swap(fDetectorID, other.fDetectorID);
  std::swap(fStationID, other.fStationID);
  std::swap(fTrackingStatus, other.fTrackingStatus);
  std::swap(fXmax, other.fXmax);
  std::swap(fYmax, other.fYmax);
  std::swap(fZref, other.fZref);
  std::swap(fZmin, other.fZmin);
  std::swap(fZmax, other.fZmax);
  std::swap(fStation, other.fStation);
  std::swap(fThicknessMap, other.fThicknessMap);
  std::swap(fInitController, other.fInitController);
}

// ---------------------------------------------------------------------------------------------------------------------
//
MaterialMap&& StationInitializer::TakeMaterialMap()
{
  if (!fInitController.IsFinalized()) {
    LOG(fatal) << "StationInitializer::GetMaterialMap: attempt of getting the material map object from uninitialized "
                  "L1BaseStation class (detectorID = "
               << static_cast<int>(fDetectorID) << ", stationID = " << fStationID << ")";
  }

  if (fManagementFlags[static_cast<int>(EManagementFlag::kThicknessMapMoved)]) {
    LOG(fatal) << "StationInitializer::GetMaterialMap: attempt of taking the material map, which has been  moved. The "
                  "thickness map instance have been "
               << "already took from the StationInitializer object (detectorID = " << static_cast<int>(fDetectorID)
               << ", stationID = " << fStationID << ")";
  }
  else {
    fManagementFlags[static_cast<int>(EManagementFlag::kThicknessMapMoved)] = true;
  }

  return std::move(fThicknessMap);
}

// ---------------------------------------------------------------------------------------------------------------------
//
std::string StationInitializer::ToString(int verbosityLevel, int indentLevel) const
{
  std::stringstream aStream{};
  constexpr char indentChar = '\t';
  std::string indent(indentLevel, indentChar);

  if (verbosityLevel == 0) {
    aStream << indent << "StationInitializer object: {stationID, detectorID, z, address} = {" << fStationID << ", "
            << static_cast<int>(fDetectorID) << ", " << fZref << ", " << this << '}';
  }
  else if (verbosityLevel > 0) {
    aStream << indent << "StationInitializer object: at " << this << '\n';
    aStream << indent << indentChar << "Station ID:              " << fStationID << '\n';
    aStream << indent << indentChar << "Detector ID:             " << static_cast<int>(fDetectorID) << '\n';
    aStream << indent << indentChar << "ca::Station object:" << '\n';
    aStream << fStation.ToString(verbosityLevel - 1, indentLevel + 1) << '\n';
    aStream << indent << indentChar << "Additional fields:\n";
    aStream << indent << indentChar << indentChar << "Zmin:                    " << fZmin << '\n';
    aStream << indent << indentChar << indentChar << "Zmax:                    " << fZmax << '\n';
    aStream << indent << indentChar << indentChar << "Xmax:                    " << fXmax << '\n';
    aStream << indent << indentChar << indentChar << "Ymax:                    " << fYmax << '\n';
  }
  return aStream.str();
}