/*
 * \class L2TauTagProducer
 *
 * L2Tau identification using Convolutional NN.
 *
 * \author Valeria D'Amante, Università di Siena and INFN Pisa
 *         Konstantin Androsov, EPFL and ETHZ
*/
#include <memory>
#include <boost/property_tree/json_parser.hpp>
#include <cmath>
#include "FWCore/Framework/interface/stream/EDProducer.h"
#include "FWCore/Framework/interface/ESHandle.h"
#include "FWCore/Framework/interface/Event.h"
#include "FWCore/Framework/interface/EventSetup.h"
#include "FWCore/Framework/interface/Frameworkfwd.h"
#include "FWCore/MessageLogger/interface/MessageLogger.h"
#include "DataFormats/Math/interface/deltaR.h"
#include "DataFormats/Common/interface/Handle.h"
#include "FWCore/Utilities/interface/InputTag.h"
#include "FWCore/Utilities/interface/isFinite.h"
#include "FWCore/ParameterSet/interface/ParameterSet.h"
#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
#include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
#include "Geometry/CaloGeometry/interface/CaloCellGeometry.h"
#include "Geometry/CaloGeometry/interface/CaloGeometry.h"
#include "Geometry/CaloTopology/interface/HcalTopology.h"
#include "Geometry/Records/interface/CaloGeometryRecord.h"
#include "DataFormats/CaloRecHit/interface/CaloRecHit.h"
#include "DataFormats/EcalRecHit/interface/EcalRecHit.h"
#include "DataFormats/EcalRecHit/interface/EcalRecHitCollections.h"
#include "DataFormats/EcalDetId/interface/EcalDetIdCollections.h"
#include "DataFormats/HcalDetId/interface/HcalDetId.h"
#include "DataFormats/HcalRecHit/interface/HBHERecHit.h"
#include "DataFormats/HcalRecHit/interface/HcalRecHitDefs.h"
#include "DataFormats/HcalRecHit/interface/HFRecHit.h"
#include "DataFormats/HcalRecHit/interface/HORecHit.h"
#include "DataFormats/HLTReco/interface/TriggerTypeDefs.h"
#include "DataFormats/HLTReco/interface/TriggerFilterObjectWithRefs.h"
#include "TrackingTools/TrajectoryParametrization/interface/CurvilinearTrajectoryError.h"
#include "RecoPixelVertexing/PixelTrackFitting/interface/FitUtils.h"
#include "TrackingTools/TrajectoryParametrization/interface/GlobalTrajectoryParameters.h"
#include "DataFormats/TrackReco/interface/HitPattern.h"
#include "TrackingTools/AnalyticalJacobians/interface/JacobianLocalToCurvilinear.h"
#include "DataFormats/TrajectoryState/interface/LocalTrajectoryParameters.h"
#include "DataFormats/GeometrySurface/interface/Plane.h"
#include "DataFormats/BeamSpot/interface/BeamSpot.h"
#include "CUDADataFormats/Track/interface/PixelTrackHeterogeneous.h"
#include "MagneticField/Records/interface/IdealMagneticFieldRecord.h"
#include "CUDADataFormats/SiPixelCluster/interface/gpuClusteringConstants.h"
#include "CUDADataFormats/Track/interface/TrackSoAHeterogeneousT.h"
#include "CUDADataFormats/Vertex/interface/ZVertexSoA.h"
#include "CUDADataFormats/Vertex/interface/ZVertexHeterogeneous.h"

namespace L2TauTagNNv1 {
  constexpr int nCellEta = 5;
  constexpr int nCellPhi = 5;
  constexpr int nVars = 31;
  constexpr float dR_max = 0.5;
  enum class NNInputs {
    nVertices = 0,
    l1Tau_pt,
    l1Tau_eta,
    l1Tau_hwIso,
    EcalEnergySum,
    EcalSize,
    EcalEnergyStdDev,
    EcalDeltaEta,
    EcalDeltaPhi,
    EcalChi2,
    EcalEnergySumForPositiveChi2,
    EcalSizeForPositiveChi2,
    HcalEnergySum,
    HcalSize,
    HcalEnergyStdDev,
    HcalDeltaEta,
    HcalDeltaPhi,
    HcalChi2,
    HcalEnergySumForPositiveChi2,
    HcalSizeForPositiveChi2,
    PatatrackPtSum,
    PatatrackSize,
    PatatrackSizeWithVertex,
    PatatrackPtSumWithVertex,
    PatatrackChargeSum,
    PatatrackDeltaEta,
    PatatrackDeltaPhi,
    PatatrackChi2OverNdof,
    PatatrackNdof,
    PatatrackDxy,
    PatatrackDz
  };

  const std::map<NNInputs, std::string> varNameMap = {
      {NNInputs::nVertices, "nVertices"},
      {NNInputs::l1Tau_pt, "l1Tau_pt"},
      {NNInputs::l1Tau_eta, "l1Tau_eta"},
      {NNInputs::l1Tau_hwIso, "l1Tau_hwIso"},
      {NNInputs::EcalEnergySum, "EcalEnergySum"},
      {NNInputs::EcalSize, "EcalSize"},
      {NNInputs::EcalEnergyStdDev, "EcalEnergyStdDev"},
      {NNInputs::EcalDeltaEta, "EcalDeltaEta"},
      {NNInputs::EcalDeltaPhi, "EcalDeltaPhi"},
      {NNInputs::EcalChi2, "EcalChi2"},
      {NNInputs::EcalEnergySumForPositiveChi2, "EcalEnergySumForPositiveChi2"},
      {NNInputs::EcalSizeForPositiveChi2, "EcalSizeForPositiveChi2"},
      {NNInputs::HcalEnergySum, "HcalEnergySum"},
      {NNInputs::HcalSize, "HcalSize"},
      {NNInputs::HcalEnergyStdDev, "HcalEnergyStdDev"},
      {NNInputs::HcalDeltaEta, "HcalDeltaEta"},
      {NNInputs::HcalDeltaPhi, "HcalDeltaPhi"},
      {NNInputs::HcalChi2, "HcalChi2"},
      {NNInputs::HcalEnergySumForPositiveChi2, "HcalEnergySumForPositiveChi2"},
      {NNInputs::HcalSizeForPositiveChi2, "HcalSizeForPositiveChi2"},
      {NNInputs::PatatrackPtSum, "PatatrackPtSum"},
      {NNInputs::PatatrackSize, "PatatrackSize"},
      {NNInputs::PatatrackSizeWithVertex, "PatatrackSizeWithVertex"},
      {NNInputs::PatatrackPtSumWithVertex, "PatatrackPtSumWithVertex"},
      {NNInputs::PatatrackChargeSum, "PatatrackChargeSum"},
      {NNInputs::PatatrackDeltaEta, "PatatrackDeltaEta"},
      {NNInputs::PatatrackDeltaPhi, "PatatrackDeltaPhi"},
      {NNInputs::PatatrackChi2OverNdof, "PatatrackChi2OverNdof"},
      {NNInputs::PatatrackNdof, "PatatrackNdof"},
      {NNInputs::PatatrackDxy, "PatatrackDxy"},
      {NNInputs::PatatrackDz, "PatatrackDz"}};
}  // namespace L2TauTagNNv1
namespace {
  inline float& getCellImpl(
      tensorflow::Tensor& cellGridMatrix, int tau_idx, int phi_idx, int eta_idx, L2TauTagNNv1::NNInputs NNInput_idx) {
    return cellGridMatrix.tensor<float, 4>()(tau_idx, phi_idx, eta_idx, static_cast<int>(NNInput_idx));
  }
}  // namespace
struct normDictElement {
  float mean;
  float std;
  float min;
  float max;
};

struct L2TauNNProducerCacheData {
  L2TauNNProducerCacheData() : graphDef(nullptr), session(nullptr) {}
  tensorflow::GraphDef* graphDef;
  tensorflow::Session* session;
  std::vector<normDictElement> normVec;
};

class L2TauNNProducer : public edm::stream::EDProducer<edm::GlobalCache<L2TauNNProducerCacheData>> {
public:
  struct caloRecHitCollections {
    const HBHERecHitCollection* hbhe;
    const HORecHitCollection* ho;
    const EcalRecHitCollection* eb;
    const EcalRecHitCollection* ee;
    const CaloGeometry* geometry;
  };

  struct InputDescTau {
    std::string CollectionName;
    edm::EDGetTokenT<trigger::TriggerFilterObjectWithRefs> inputToken_;
  };

  static constexpr float dR2_max = L2TauTagNNv1::dR_max * L2TauTagNNv1::dR_max;
  static constexpr float dEta_width = 2 * L2TauTagNNv1::dR_max / static_cast<float>(L2TauTagNNv1::nCellEta);
  static constexpr float dPhi_width = 2 * L2TauTagNNv1::dR_max / static_cast<float>(L2TauTagNNv1::nCellPhi);

  explicit L2TauNNProducer(const edm::ParameterSet&, const L2TauNNProducerCacheData*);
  static void fillDescriptions(edm::ConfigurationDescriptions&);
  static std::unique_ptr<L2TauNNProducerCacheData> initializeGlobalCache(const edm::ParameterSet&);
  static void globalEndJob(L2TauNNProducerCacheData*);

private:
  void checknan(tensorflow::Tensor& tensor, int debugLevel);
  void standardizeTensor(tensorflow::Tensor& tensor);
  std::vector<float> getTauScore(const tensorflow::Tensor& cellGridMatrix);
  void produce(edm::Event& event, const edm::EventSetup& eventsetup) override;
  void fillL1TauVars(tensorflow::Tensor& cellGridMatrix, const std::vector<l1t::TauRef>& allTaus);
  void fillCaloRecHits(tensorflow::Tensor& cellGridMatrix,
                       const std::vector<l1t::TauRef>& allTaus,
                       const caloRecHitCollections& caloRecHits);
  void fillPatatracks(tensorflow::Tensor& cellGridMatrix,
                      const std::vector<l1t::TauRef>& allTaus,
                      const pixelTrack::TrackSoA& patatracks_tsoa,
                      const ZVertexSoA& patavtx_soa,
                      const reco::BeamSpot& beamspot,
                      const MagneticField* magfi);
  std::vector<int> selectGoodVertices(const ZVertexSoA& patavtx_soa,
                                      const pixelTrack::TrackSoA& patatracks_tsoa,
                                      const std::vector<int>& TrackGood);
  std::pair<float, float> impactParameter(int it,
                                          const pixelTrack::TrackSoA& patatracks_tsoa,
                                          float patatrackPhi,
                                          const reco::BeamSpot& beamspot,
                                          const MagneticField* magfi);
  template <typename VPos, typename LVec>
  std::tuple<float, float, int, int> getEtaPhiIndices(const VPos& position, const LVec& tau_p4);
  template <typename LVec>
  std::tuple<float, float, int, int> getEtaPhiIndices(float eta, float phi, const LVec& tau_p4);

private:
  const int debugLevel_;
  const edm::EDGetTokenT<trigger::TriggerFilterObjectWithRefs> tauTriggerToken_;
  std::vector<InputDescTau> L1TauDesc_;
  const edm::EDGetTokenT<HBHERecHitCollection> hbheToken_;
  const edm::EDGetTokenT<HORecHitCollection> hoToken_;
  const edm::EDGetTokenT<EcalRecHitCollection> ebToken_;
  const edm::EDGetTokenT<EcalRecHitCollection> eeToken_;
  const edm::ESGetToken<CaloGeometry, CaloGeometryRecord> geometryToken_;
  const edm::ESGetToken<MagneticField, IdealMagneticFieldRecord> bFieldToken_;
  const edm::EDGetTokenT<ZVertexHeterogeneous> pataVerticesToken_;
  const edm::EDGetTokenT<PixelTrackHeterogeneous> pataTracksToken_;
  const edm::EDGetTokenT<reco::BeamSpot> beamSpotToken_;
  const unsigned int maxVtx_;
  const double fractionSumPt2_;
  const double minSumPt2_;
  const double trackPtMin_;
  const double trackPtMax_;
  const double trackChi2Max_;
  std::string inputTensorName_;
  std::string outputTensorName_;
  const L2TauNNProducerCacheData* L2cacheData_;
};

std::unique_ptr<L2TauNNProducerCacheData> L2TauNNProducer::initializeGlobalCache(const edm::ParameterSet& cfg) {
  std::unique_ptr<L2TauNNProducerCacheData> cacheData = std::make_unique<L2TauNNProducerCacheData>();
  cacheData->normVec.reserve(L2TauTagNNv1::nVars);

  auto const graphPath = edm::FileInPath(cfg.getParameter<std::string>("graphPath")).fullPath();

  cacheData->graphDef = tensorflow::loadGraphDef(graphPath);
  cacheData->session = tensorflow::createSession(cacheData->graphDef);

  tensorflow::setLogging("2");

  boost::property_tree::ptree loadPtreeRoot;
  auto const normalizationDict = edm::FileInPath(cfg.getParameter<std::string>("normalizationDict")).fullPath();
  boost::property_tree::read_json(normalizationDict, loadPtreeRoot);
  for (const auto& [key, val] : L2TauTagNNv1::varNameMap) {
    boost::property_tree::ptree var = loadPtreeRoot.get_child(val);
    normDictElement current_element;
    current_element.mean = var.get_child("mean").get_value<float>();
    current_element.std = var.get_child("std").get_value<float>();
    current_element.min = var.get_child("min").get_value<float>();
    current_element.max = var.get_child("max").get_value<float>();
    cacheData->normVec.push_back(current_element);
  }
  return cacheData;
}
void L2TauNNProducer::globalEndJob(L2TauNNProducerCacheData* cacheData) {
  if (cacheData->graphDef != nullptr) {
    delete cacheData->graphDef;
  }
  tensorflow::closeSession(cacheData->session);
}
void L2TauNNProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
  edm::ParameterSetDescription desc;
  desc.add<int>("debugLevel", 0)->setComment("set debug level for printing out info");
  edm::ParameterSetDescription l1TausPset;
  l1TausPset.add<std::string>("L1CollectionName", "DoubleTau")->setComment("Name of collections");
  l1TausPset.add<edm::InputTag>("L1TauTrigger", edm::InputTag("hltL1sDoubleTauBigOR"))
      ->setComment("Which trigger should the L1 Taus collection pass");
  edm::ParameterSet l1TausPSetDefault;
  l1TausPSetDefault.addParameter<std::string>("L1CollectionName", "DoubleTau");
  l1TausPSetDefault.addParameter<edm::InputTag>("L1TauTrigger", edm::InputTag("hltL1sDoubleTauBigOR"));
  desc.addVPSet("L1Taus", l1TausPset, {l1TausPSetDefault});
  desc.add<edm::InputTag>("hbheInput", edm::InputTag("hltHbhereco"))->setComment("HBHE recHit collection");
  desc.add<edm::InputTag>("hoInput", edm::InputTag("hltHoreco"))->setComment("HO recHit Collection");
  desc.add<edm::InputTag>("ebInput", edm::InputTag("hltEcalRecHit:EcalRecHitsEB"))->setComment("EB recHit Collection");
  desc.add<edm::InputTag>("eeInput", edm::InputTag("hltEcalRecHit:EcalRecHitsEE"))->setComment("EE recHit Collection");
  desc.add<edm::InputTag>("pataVertices", edm::InputTag("hltPixelVerticesSoA"))
      ->setComment("patatrack vertices collection");
  desc.add<edm::InputTag>("pataTracks", edm::InputTag("hltPixelTracksSoA"))->setComment("patatrack collection");
  desc.add<edm::InputTag>("BeamSpot", edm::InputTag("hltOnlineBeamSpot"))->setComment("BeamSpot Collection");
  desc.add<uint>("maxVtx", 100)->setComment("max output collection size (number of accepted vertices)");
  desc.add<double>("fractionSumPt2", 0.3)->setComment("threshold on sumPt2 fraction of the leading vertex");
  desc.add<double>("minSumPt2", 0.)->setComment("min sumPt2");
  desc.add<double>("track_pt_min", 1.0)->setComment("min track p_T");
  desc.add<double>("track_pt_max", 10.0)->setComment("max track p_T");
  desc.add<double>("track_chi2_max", 99999.)->setComment("max track chi2");
  desc.add<std::string>("graphPath", "RecoTauTag/TrainingFiles/data/L2TauNNTag/L2TauTag_Run3v1.pb")
      ->setComment("path to the saved CNN");
  desc.add<std::string>("normalizationDict", "RecoTauTag/TrainingFiles/data/L2TauNNTag/NormalizationDict.json")
      ->setComment("path to the dictionary for variable standardization");
  descriptions.addWithDefaultLabel(desc);
}

L2TauNNProducer::L2TauNNProducer(const edm::ParameterSet& cfg, const L2TauNNProducerCacheData* cacheData)
    : debugLevel_(cfg.getParameter<int>("debugLevel")),
      hbheToken_(consumes<HBHERecHitCollection>(cfg.getParameter<edm::InputTag>("hbheInput"))),
      hoToken_(consumes<HORecHitCollection>(cfg.getParameter<edm::InputTag>("hoInput"))),
      ebToken_(consumes<EcalRecHitCollection>(cfg.getParameter<edm::InputTag>("ebInput"))),
      eeToken_(consumes<EcalRecHitCollection>(cfg.getParameter<edm::InputTag>("eeInput"))),
      geometryToken_(esConsumes<CaloGeometry, CaloGeometryRecord>()),
      bFieldToken_(esConsumes<MagneticField, IdealMagneticFieldRecord>()),
      pataVerticesToken_(consumes<ZVertexHeterogeneous>(cfg.getParameter<edm::InputTag>("pataVertices"))),
      pataTracksToken_(consumes<PixelTrackHeterogeneous>(cfg.getParameter<edm::InputTag>("pataTracks"))),
      beamSpotToken_(consumes<reco::BeamSpot>(cfg.getParameter<edm::InputTag>("BeamSpot"))),
      maxVtx_(cfg.getParameter<uint>("maxVtx")),
      fractionSumPt2_(cfg.getParameter<double>("fractionSumPt2")),
      minSumPt2_(cfg.getParameter<double>("minSumPt2")),
      trackPtMin_(cfg.getParameter<double>("track_pt_min")),
      trackPtMax_(cfg.getParameter<double>("track_pt_max")),
      trackChi2Max_(cfg.getParameter<double>("track_chi2_max")) {
  if (cacheData->graphDef == nullptr) {
    throw cms::Exception("InvalidCacheData") << "Invalid Cache Data.";
  }
  inputTensorName_ = cacheData->graphDef->node(0).name();
  outputTensorName_ = cacheData->graphDef->node(cacheData->graphDef->node_size() - 1).name();
  L2cacheData_ = cacheData;
  std::vector<edm::ParameterSet> L1TauCollections = cfg.getParameter<std::vector<edm::ParameterSet>>("L1Taus");
  L1TauDesc_.reserve(L1TauCollections.size());
  for (const auto& l1TauInput : L1TauCollections) {
    InputDescTau toInsert;
    toInsert.CollectionName = l1TauInput.getParameter<std::string>("L1CollectionName");
    toInsert.inputToken_ =
        consumes<trigger::TriggerFilterObjectWithRefs>(l1TauInput.getParameter<edm::InputTag>("L1TauTrigger"));
    L1TauDesc_.push_back(toInsert);
  }
  for (const auto& desc : L1TauDesc_)
    produces<std::vector<float>>(desc.CollectionName);
}

void L2TauNNProducer::checknan(tensorflow::Tensor& tensor, int debugLevel) {
  using NNInputs = L2TauTagNNv1::NNInputs;
  std::vector<int> tensor_shape(tensor.shape().dims());
  for (int d = 0; d < tensor.shape().dims(); d++) {
    tensor_shape.at(d) = tensor.shape().dim_size(d);
  }
  if (tensor_shape.size() != 4) {
    throw cms::Exception("InvalidTensor") << "Tensor shape does not have 4 dimensions!";
  }
  for (int tau_idx = 0; tau_idx < tensor_shape.at(0); tau_idx++) {
    for (int phi_idx = 0; phi_idx < tensor_shape.at(1); phi_idx++) {
      for (int eta_idx = 0; eta_idx < tensor_shape.at(2); eta_idx++) {
        for (int var_idx = 0; var_idx < tensor_shape.at(3); var_idx++) {
          auto getCell = [&](NNInputs input) -> float& {
            return getCellImpl(tensor, tau_idx, phi_idx, eta_idx, input);
          };
          auto nonstd_var = getCell(static_cast<NNInputs>(var_idx));
          if (edm::isNotFinite(nonstd_var)) {
            edm::LogWarning("InputVar") << "var is nan \nvar name= "
                                        << L2TauTagNNv1::varNameMap.at(static_cast<L2TauTagNNv1::NNInputs>(var_idx))
                                        << "\t var_idx = " << var_idx << "\t eta_idx = " << eta_idx
                                        << "\t phi_idx = " << phi_idx << "\t tau_idx = " << tau_idx << std::endl;
            if (debugLevel > 2) {
              edm::LogWarning("InputVar") << "other vars in same cell \n";
              if (var_idx + 1 < tensor_shape.at(3))
                edm::LogWarning("InputVar") << L2TauTagNNv1::varNameMap.at(static_cast<NNInputs>(var_idx + 1))
                                            << "\t = " << getCell(static_cast<NNInputs>(var_idx + 1)) << std::endl;
              if (var_idx + 2 < tensor_shape.at(3))
                edm::LogWarning("InputVar") << L2TauTagNNv1::varNameMap.at(static_cast<NNInputs>(var_idx + 2))
                                            << "\t = " << getCell(static_cast<NNInputs>(var_idx + 2)) << std::endl;
              if (var_idx + 3 < tensor_shape.at(3))
                edm::LogWarning("InputVar") << L2TauTagNNv1::varNameMap.at(static_cast<NNInputs>(var_idx + 3))
                                            << "\t = " << getCell(static_cast<NNInputs>(var_idx + 3)) << std::endl;
              if (var_idx + 4 < tensor_shape.at(3))
                edm::LogWarning("InputVar") << L2TauTagNNv1::varNameMap.at(static_cast<NNInputs>(var_idx + 4))
                                            << "\t = " << getCell(static_cast<NNInputs>(var_idx + 4)) << std::endl;
            }
          }
        }
      }
    }
  }
}

void L2TauNNProducer::standardizeTensor(tensorflow::Tensor& tensor) {
  using NNInputs = L2TauTagNNv1::NNInputs;
  std::vector<int> tensor_shape(tensor.shape().dims());
  for (int d = 0; d < tensor.shape().dims(); d++) {
    tensor_shape.at(d) = tensor.shape().dim_size(d);
  }
  if (tensor_shape.size() != 4) {
    throw cms::Exception("InvalidTensor") << "Tensor shape does not have 4 dimensions!";
  }
  for (int tau_idx = 0; tau_idx < tensor_shape.at(0); tau_idx++) {
    for (int phi_idx = 0; phi_idx < tensor_shape.at(1); phi_idx++) {
      for (int eta_idx = 0; eta_idx < tensor_shape.at(2); eta_idx++) {
        for (int var_idx = 0; var_idx < tensor_shape.at(3); var_idx++) {
          auto getCell = [&](NNInputs input) -> float& {
            return getCellImpl(tensor, tau_idx, phi_idx, eta_idx, input);
          };
          float mean = L2cacheData_->normVec.at(var_idx).mean;
          float std = L2cacheData_->normVec.at(var_idx).std;
          float min = L2cacheData_->normVec.at(var_idx).min;
          float max = L2cacheData_->normVec.at(var_idx).max;
          float nonstd_var = getCell(static_cast<NNInputs>(var_idx));
          float std_var = static_cast<float>((nonstd_var - mean) / std);
          if (std_var > max) {
            std_var = static_cast<float>(max);
          } else if (std_var < min) {
            std_var = static_cast<float>(min);
          }
          getCell(static_cast<NNInputs>(var_idx)) = std_var;
        }
      }
    }
  }
}

void L2TauNNProducer::fillL1TauVars(tensorflow::Tensor& cellGridMatrix, const std::vector<l1t::TauRef>& allTaus) {
  using NNInputs = L2TauTagNNv1::NNInputs;
  const int nTaus = static_cast<int>(allTaus.size());
  for (int tau_idx = 0; tau_idx < nTaus; tau_idx++) {
    for (int eta_idx = 0; eta_idx < L2TauTagNNv1::nCellEta; eta_idx++) {
      for (int phi_idx = 0; phi_idx < L2TauTagNNv1::nCellPhi; phi_idx++) {
        auto getCell = [&](NNInputs input) -> float& {
          return getCellImpl(cellGridMatrix, tau_idx, phi_idx, eta_idx, input);
        };
        getCell(NNInputs::l1Tau_pt) = allTaus[tau_idx]->pt();
        getCell(NNInputs::l1Tau_eta) = allTaus[tau_idx]->eta();
        getCell(NNInputs::l1Tau_hwIso) = allTaus[tau_idx]->hwIso();
      }
    }
  }
}

template <typename LVec>
std::tuple<float, float, int, int> L2TauNNProducer::getEtaPhiIndices(float eta, float phi, const LVec& tau_p4) {
  const float deta = eta - tau_p4.eta();
  const float dphi = reco::deltaPhi(phi, tau_p4.phi());
  const int eta_idx = static_cast<int>(floor((deta + L2TauTagNNv1::dR_max) / dEta_width));
  const int phi_idx = static_cast<int>(floor((dphi + L2TauTagNNv1::dR_max) / dPhi_width));
  return std::make_tuple(deta, dphi, eta_idx, phi_idx);
}

template <typename VPos, typename LVec>
std::tuple<float, float, int, int> L2TauNNProducer::getEtaPhiIndices(const VPos& position, const LVec& tau_p4) {
  return getEtaPhiIndices(position.eta(), position.phi(), tau_p4);
}

void L2TauNNProducer::fillCaloRecHits(tensorflow::Tensor& cellGridMatrix,
                                      const std::vector<l1t::TauRef>& allTaus,
                                      const caloRecHitCollections& caloRecHits) {
  using NNInputs = L2TauTagNNv1::NNInputs;
  const int nTaus = static_cast<int>(allTaus.size());
  float deta, dphi;
  int eta_idx = 0;
  int phi_idx = 0;
  int tau_idx = 0;

  auto getCell = [&](NNInputs input) -> float& {
    return getCellImpl(cellGridMatrix, tau_idx, phi_idx, eta_idx, input);
  };
  for (tau_idx = 0; tau_idx < nTaus; tau_idx++) {
    // calorechit_EE
    for (const auto& caloRecHit_ee : *caloRecHits.ee) {
      if (caloRecHit_ee.energy() <= 0)
        continue;
      const auto& position = caloRecHits.geometry->getGeometry(caloRecHit_ee.id())->getPosition();
      const float eeCalEn = caloRecHit_ee.energy();
      const float eeCalChi2 = caloRecHit_ee.chi2();
      if (reco::deltaR2(position, allTaus[tau_idx]->polarP4()) < dR2_max) {
        std::tie(deta, dphi, eta_idx, phi_idx) = getEtaPhiIndices(position, allTaus[tau_idx]->polarP4());
        getCell(NNInputs::EcalEnergySum) += eeCalEn;
        getCell(NNInputs::EcalSize) += 1.;
        getCell(NNInputs::EcalEnergyStdDev) += eeCalEn * eeCalEn;
        getCell(NNInputs::EcalDeltaEta) += deta * eeCalEn;
        getCell(NNInputs::EcalDeltaPhi) += dphi * eeCalEn;
        if (eeCalChi2 >= 0) {
          getCell(NNInputs::EcalChi2) += eeCalChi2 * eeCalEn;
          getCell(NNInputs::EcalEnergySumForPositiveChi2) += eeCalEn;
          getCell(NNInputs::EcalSizeForPositiveChi2) += 1.;
        }
      }
    }

    // calorechit_EB
    for (const auto& caloRecHit_eb : *caloRecHits.eb) {
      if (caloRecHit_eb.energy() <= 0)
        continue;
      const auto& position = caloRecHits.geometry->getGeometry(caloRecHit_eb.id())->getPosition();
      const float ebCalEn = caloRecHit_eb.energy();
      const float ebCalChi2 = caloRecHit_eb.chi2();
      if (reco::deltaR2(position, allTaus[tau_idx]->polarP4()) < dR2_max) {
        std::tie(deta, dphi, eta_idx, phi_idx) = getEtaPhiIndices(position, allTaus[tau_idx]->polarP4());
        getCell(NNInputs::EcalEnergySum) += ebCalEn;
        getCell(NNInputs::EcalSize) += 1.;
        getCell(NNInputs::EcalEnergyStdDev) += ebCalEn * ebCalEn;
        getCell(NNInputs::EcalDeltaEta) += deta * ebCalEn;
        getCell(NNInputs::EcalDeltaPhi) += dphi * ebCalEn;
        if (ebCalChi2 >= 0) {
          getCell(NNInputs::EcalChi2) += ebCalChi2 * ebCalEn;
          getCell(NNInputs::EcalEnergySumForPositiveChi2) += ebCalEn;
          getCell(NNInputs::EcalSizeForPositiveChi2) += 1.;
        }
      }
    }

    // calorechit_HBHE
    for (const auto& caloRecHit_hbhe : *caloRecHits.hbhe) {
      if (caloRecHit_hbhe.energy() <= 0)
        continue;
      const auto& position = caloRecHits.geometry->getGeometry(caloRecHit_hbhe.id())->getPosition();
      const float hbheCalEn = caloRecHit_hbhe.energy();
      const float hbheCalChi2 = caloRecHit_hbhe.chi2();
      if (reco::deltaR2(position, allTaus[tau_idx]->polarP4()) < dR2_max) {
        std::tie(deta, dphi, eta_idx, phi_idx) = getEtaPhiIndices(position, allTaus[tau_idx]->polarP4());
        getCell(NNInputs::HcalEnergySum) += hbheCalEn;
        getCell(NNInputs::HcalEnergyStdDev) += hbheCalEn * hbheCalEn;
        getCell(NNInputs::HcalSize) += 1.;
        getCell(NNInputs::HcalDeltaEta) += deta * hbheCalEn;
        getCell(NNInputs::HcalDeltaPhi) += dphi * hbheCalEn;
        if (hbheCalChi2 >= 0) {
          getCell(NNInputs::HcalChi2) += hbheCalChi2 * hbheCalEn;
          getCell(NNInputs::HcalEnergySumForPositiveChi2) += hbheCalEn;
          getCell(NNInputs::HcalSizeForPositiveChi2) += 1.;
        }
      }
    }

    // calorechit_HO
    for (const auto& caloRecHit_ho : *caloRecHits.ho) {
      if (caloRecHit_ho.energy() <= 0)
        continue;
      const auto& position = caloRecHits.geometry->getGeometry(caloRecHit_ho.id())->getPosition();
      const float hoCalEn = caloRecHit_ho.energy();
      if (reco::deltaR2(position, allTaus[tau_idx]->polarP4()) < dR2_max) {
        std::tie(deta, dphi, eta_idx, phi_idx) = getEtaPhiIndices(position, allTaus[tau_idx]->polarP4());
        getCell(NNInputs::HcalEnergySum) += hoCalEn;
        getCell(NNInputs::HcalEnergyStdDev) += hoCalEn * hoCalEn;
        getCell(NNInputs::HcalSize) += 1.;
        getCell(NNInputs::HcalDeltaEta) += deta * hoCalEn;
        getCell(NNInputs::HcalDeltaPhi) += dphi * hoCalEn;
      }
    }

    // normalize to sum and define stdDev
    for (eta_idx = 0; eta_idx < L2TauTagNNv1::nCellEta; eta_idx++) {
      for (phi_idx = 0; phi_idx < L2TauTagNNv1::nCellPhi; phi_idx++) {
        /* normalize eCal vars*/
        if (getCell(NNInputs::EcalEnergySum) > 0.) {
          getCell(NNInputs::EcalDeltaEta) /= getCell(NNInputs::EcalEnergySum);
          getCell(NNInputs::EcalDeltaPhi) /= getCell(NNInputs::EcalEnergySum);
        }
        if (getCell(NNInputs::EcalEnergySumForPositiveChi2) > 0.) {
          getCell(NNInputs::EcalChi2) /= getCell(NNInputs::EcalEnergySumForPositiveChi2);
        }
        if (getCell(NNInputs::EcalSize) > 1.) {
          // (stdDev - (enSum*enSum)/size) / (size-1)
          getCell(NNInputs::EcalEnergyStdDev) =
              (getCell(NNInputs::EcalEnergyStdDev) -
               (getCell(NNInputs::EcalEnergySum) * getCell(NNInputs::EcalEnergySum)) / getCell(NNInputs::EcalSize)) /
              (getCell(NNInputs::EcalSize) - 1);
        } else {
          getCell(NNInputs::EcalEnergyStdDev) = 0.;
        }
        /* normalize hCal Vars */
        if (getCell(NNInputs::HcalEnergySum) > 0.) {
          getCell(NNInputs::HcalDeltaEta) /= getCell(NNInputs::HcalEnergySum);
          getCell(NNInputs::HcalDeltaPhi) /= getCell(NNInputs::HcalEnergySum);
        }
        if (getCell(NNInputs::HcalEnergySumForPositiveChi2) > 0.) {
          getCell(NNInputs::HcalChi2) /= getCell(NNInputs::HcalEnergySumForPositiveChi2);
        }
        if (getCell(NNInputs::HcalSize) > 1.) {
          // (stdDev - (enSum*enSum)/size) / (size-1)
          getCell(NNInputs::HcalEnergyStdDev) =
              (getCell(NNInputs::HcalEnergyStdDev) -
               (getCell(NNInputs::HcalEnergySum) * getCell(NNInputs::HcalEnergySum)) / getCell(NNInputs::HcalSize)) /
              (getCell(NNInputs::HcalSize) - 1);
        } else {
          getCell(NNInputs::HcalEnergyStdDev) = 0.;
        }
      }
    }
  }
}

std::vector<int> L2TauNNProducer::selectGoodVertices(const ZVertexSoA& patavtx_soa,
                                                     const pixelTrack::TrackSoA& patatracks_tsoa,
                                                     const std::vector<int>& TrackGood) {
  auto maxTracks = patatracks_tsoa.stride();
  const int nv = patavtx_soa.nvFinal;
  std::vector<int> VtxGood;
  if (nv == 0)
    return VtxGood;
  VtxGood.reserve(nv);

  std::vector<double> maxChi2_;
  std::vector<double> pTSquaredSum(nv);

  for (int j = nv - 1; j >= 0; --j) {
    std::vector<int> trk_ass_to_vtx;
    auto vtx_idx = patavtx_soa.sortInd[j];
    assert(vtx_idx < nv);
    for (int trk_idx = 0; trk_idx < maxTracks; trk_idx++) {
      int vtx_ass_to_track = patavtx_soa.idv[trk_idx];
      if (vtx_ass_to_track == int16_t(vtx_idx))
        trk_ass_to_vtx.push_back(trk_idx);
    }
    auto nt = trk_ass_to_vtx.size();
    if (nt == 0) {
      continue;
    }
    if (nt < 2) {
      trk_ass_to_vtx.clear();
      continue;
    }
    for (const auto& trk_idx : trk_ass_to_vtx) {
      int vtx_ass_to_track = patavtx_soa.idv[trk_idx];
      if (vtx_ass_to_track != vtx_idx)
        continue;
      double patatrackPt = patatracks_tsoa.pt[trk_idx];
      if (patatrackPt < trackPtMin_)
        continue;
      if (patatracks_tsoa.chi2(trk_idx) > trackChi2Max_)
        continue;
      if (patatrackPt > trackPtMax_) {
        patatrackPt = trackPtMax_;
      }
      pTSquaredSum.at(vtx_idx) += patatrackPt * patatrackPt;
    }
  }
  std::vector<size_t> sortIdxs(nv);
  std::iota(sortIdxs.begin(), sortIdxs.end(), 0);
  std::sort(sortIdxs.begin(), sortIdxs.end(), [&](size_t const i1, size_t const i2) {
    return pTSquaredSum[i1] > pTSquaredSum[i2];
  });
  auto const minFOM_fromFrac = pTSquaredSum[sortIdxs.front()] * fractionSumPt2_;

  for (int j = nv - 1; j >= 0; --j) {
    auto idx = patavtx_soa.sortInd[j];

    if (VtxGood.size() >= maxVtx_) {
      break;
    }
    if (pTSquaredSum[idx] >= minFOM_fromFrac && pTSquaredSum[idx] > minSumPt2_) {
      VtxGood.push_back(idx);
    }
  }
  return VtxGood;
}

std::pair<float, float> L2TauNNProducer::impactParameter(int it,
                                                         const pixelTrack::TrackSoA& patatracks_tsoa,
                                                         float patatrackPhi,
                                                         const reco::BeamSpot& beamspot,
                                                         const MagneticField* magfi) {
  auto const& fit = patatracks_tsoa.stateAtBS;
  /* dxy and dz */
  riemannFit::Vector5d ipar, opar;
  riemannFit::Matrix5d icov, ocov;
  fit.copyToDense(ipar, icov, it);
  riemannFit::transformToPerigeePlane(ipar, icov, opar, ocov);
  LocalTrajectoryParameters lpar(opar(0), opar(1), opar(2), opar(3), opar(4), 1.);
  float sp = std::sin(patatrackPhi);
  float cp = std::cos(patatrackPhi);
  Surface::RotationType Rotation(sp, -cp, 0, 0, 0, -1.f, cp, sp, 0);
  GlobalPoint BeamSpotPoint(beamspot.x0(), beamspot.y0(), beamspot.z0());
  Plane impPointPlane(BeamSpotPoint, Rotation);
  GlobalTrajectoryParameters gp(
      impPointPlane.toGlobal(lpar.position()), impPointPlane.toGlobal(lpar.momentum()), lpar.charge(), magfi);
  GlobalPoint vv = gp.position();
  math::XYZPoint pos(vv.x(), vv.y(), vv.z());
  GlobalVector pp = gp.momentum();
  math::XYZVector mom(pp.x(), pp.y(), pp.z());
  auto lambda = M_PI_2 - pp.theta();
  auto phi = pp.phi();
  float patatrackDxy = -vv.x() * std::sin(phi) + vv.y() * std::cos(phi);
  float patatrackDz =
      (vv.z() * std::cos(lambda) - (vv.x() * std::cos(phi) + vv.y() * std::sin(phi)) * std::sin(lambda)) /
      std::cos(lambda);
  return std::make_pair(patatrackDxy, patatrackDz);
}

void L2TauNNProducer::fillPatatracks(tensorflow::Tensor& cellGridMatrix,
                                     const std::vector<l1t::TauRef>& allTaus,
                                     const pixelTrack::TrackSoA& patatracks_tsoa,
                                     const ZVertexSoA& patavtx_soa,
                                     const reco::BeamSpot& beamspot,
                                     const MagneticField* magfi) {
  using NNInputs = L2TauTagNNv1::NNInputs;
  float deta, dphi;
  int eta_idx = 0;
  int phi_idx = 0;
  int tau_idx = 0;

  auto getCell = [&](NNInputs input) -> float& {
    return getCellImpl(cellGridMatrix, tau_idx, phi_idx, eta_idx, input);
  };
  const int nTaus = static_cast<int>(allTaus.size());
  for (tau_idx = 0; tau_idx < nTaus; tau_idx++) {
    const float tauEta = allTaus[tau_idx]->eta();
    const float tauPhi = allTaus[tau_idx]->phi();

    auto maxTracks = patatracks_tsoa.stride();
    auto const* quality = patatracks_tsoa.qualityData();

    std::vector<int> TrackGood;
    for (int32_t it = 0; it < maxTracks; ++it) {
      auto nHits = patatracks_tsoa.nHits(it);
      if (nHits == 0)
        break;
      auto q = quality[it];
      if (q < pixelTrack::Quality::loose)
        continue;
      if (nHits < 0)
        continue;
      TrackGood.push_back(it);
    }

    std::vector<int> VtxGood = selectGoodVertices(patavtx_soa, patatracks_tsoa, TrackGood);

    for (const auto it : TrackGood) {
      const float patatrackPt = patatracks_tsoa.pt[it];
      if (patatrackPt <= 0)
        continue;
      const float patatrackPhi = patatracks_tsoa.phi(it);
      const float patatrackEta = patatracks_tsoa.eta(it);
      const float patatrackCharge = patatracks_tsoa.charge(it);
      const float patatrackChi2OverNdof = patatracks_tsoa.chi2(it);
      const auto nHits = patatracks_tsoa.nHits(it);
      if (nHits <= 0)
        continue;
      const int patatrackNdof = 2 * nHits - 5;

      const int vtx_idx_assTrk = patavtx_soa.idv[it];
      if (reco::deltaR2(patatrackEta, patatrackPhi, tauEta, tauPhi) < dR2_max) {
        std::tie(deta, dphi, eta_idx, phi_idx) =
            getEtaPhiIndices(patatrackEta, patatrackPhi, allTaus[tau_idx]->polarP4());
        getCell(NNInputs::PatatrackPtSum) += patatrackPt;
        getCell(NNInputs::PatatrackSize) += 1.;
        getCell(NNInputs::PatatrackChargeSum) += patatrackCharge;
        getCell(NNInputs::PatatrackDeltaEta) += deta * patatrackPt;
        getCell(NNInputs::PatatrackDeltaPhi) += dphi * patatrackPt;
        getCell(NNInputs::PatatrackChi2OverNdof) += patatrackChi2OverNdof * patatrackPt;
        getCell(NNInputs::PatatrackNdof) += patatrackNdof * patatrackPt;
        std::pair<float, float> impactParameters = impactParameter(it, patatracks_tsoa, patatrackPhi, beamspot, magfi);
        getCell(NNInputs::PatatrackDxy) += impactParameters.first * patatrackPt;
        getCell(NNInputs::PatatrackDz) += impactParameters.second * patatrackPt;
        if ((std::find(VtxGood.begin(), VtxGood.end(), vtx_idx_assTrk) != VtxGood.end())) {
          getCell(NNInputs::PatatrackPtSumWithVertex) += patatrackPt;
          getCell(NNInputs::PatatrackSizeWithVertex) += 1.;
        }
      }
    }

    // normalize to sum and define stdDev
    for (eta_idx = 0; eta_idx < L2TauTagNNv1::nCellEta; eta_idx++) {
      for (phi_idx = 0; phi_idx < L2TauTagNNv1::nCellPhi; phi_idx++) {
        getCell(NNInputs::nVertices) = VtxGood.size();
        if (getCell(NNInputs::PatatrackPtSum) > 0.) {
          getCell(NNInputs::PatatrackDeltaEta) /= getCell(NNInputs::PatatrackPtSum);
          getCell(NNInputs::PatatrackDeltaPhi) /= getCell(NNInputs::PatatrackPtSum);
          getCell(NNInputs::PatatrackChi2OverNdof) /= getCell(NNInputs::PatatrackPtSum);
          getCell(NNInputs::PatatrackNdof) /= getCell(NNInputs::PatatrackPtSum);
          getCell(NNInputs::PatatrackDxy) /= getCell(NNInputs::PatatrackPtSum);
          getCell(NNInputs::PatatrackDz) /= getCell(NNInputs::PatatrackPtSum);
        }
      }
    }
  }
}

std::vector<float> L2TauNNProducer::getTauScore(const tensorflow::Tensor& cellGridMatrix) {
  std::vector<tensorflow::Tensor> pred_tensor;
  tensorflow::run(L2cacheData_->session, {{inputTensorName_, cellGridMatrix}}, {outputTensorName_}, &pred_tensor);
  const int nTau = cellGridMatrix.shape().dim_size(0);
  std::vector<float> pred_vector(nTau);
  for (int tau_idx = 0; tau_idx < nTau; ++tau_idx) {
    pred_vector[tau_idx] = pred_tensor[0].matrix<float>()(tau_idx, 0);
  }

  return pred_vector;
}

void L2TauNNProducer::produce(edm::Event& event, const edm::EventSetup& eventsetup) {
  std::vector<std::vector<size_t>> TauCollectionMap(L1TauDesc_.size());
  l1t::TauVectorRef allTaus;

  for (size_t inp_idx = 0; inp_idx < L1TauDesc_.size(); inp_idx++) {
    l1t::TauVectorRef l1Taus;
    auto const& l1TriggeredTaus = event.get(L1TauDesc_[inp_idx].inputToken_);
    l1TriggeredTaus.getObjects(trigger::TriggerL1Tau, l1Taus);
    TauCollectionMap.at(inp_idx).resize(l1Taus.size());

    for (size_t l1_idx = 0; l1_idx < l1Taus.size(); l1_idx++) {
      size_t tau_idx;
      const auto iter = std::find(allTaus.begin(), allTaus.end(), l1Taus[l1_idx]);
      if (iter != allTaus.end()) {
        tau_idx = std::distance(allTaus.begin(), iter);
      } else {
        allTaus.push_back(l1Taus[l1_idx]);
        tau_idx = allTaus.size() - 1;
      }
      TauCollectionMap.at(inp_idx).at(l1_idx) = tau_idx;
    }
  }
  const auto ebCal = event.getHandle(ebToken_);
  const auto eeCal = event.getHandle(eeToken_);
  const auto hbhe = event.getHandle(hbheToken_);
  const auto ho = event.getHandle(hoToken_);
  const auto& patatracks_SoA = *event.get(pataTracksToken_);
  const auto& vertices_SoA = *event.get(pataVerticesToken_);
  const auto bsHandle = event.getHandle(beamSpotToken_);

  auto const fieldESH = eventsetup.getHandle(bFieldToken_);
  auto const geometry = eventsetup.getHandle(geometryToken_);

  caloRecHitCollections caloRecHits;
  caloRecHits.hbhe = &*hbhe;
  caloRecHits.ho = &*ho;
  caloRecHits.eb = &*ebCal;
  caloRecHits.ee = &*eeCal;
  caloRecHits.geometry = &*geometry;

  const int nTaus = static_cast<int>(allTaus.size());
  tensorflow::Tensor cellGridMatrix(tensorflow::DT_FLOAT,
                                    {nTaus, L2TauTagNNv1::nCellEta, L2TauTagNNv1::nCellPhi, L2TauTagNNv1::nVars});
  const int n_inputs = nTaus * L2TauTagNNv1::nCellEta * L2TauTagNNv1::nCellPhi * L2TauTagNNv1::nVars;
  for (int input_idx = 0; input_idx < n_inputs; ++input_idx) {
    cellGridMatrix.flat<float>()(input_idx) = 0;
  }
  fillL1TauVars(cellGridMatrix, allTaus);

  fillCaloRecHits(cellGridMatrix, allTaus, caloRecHits);

  fillPatatracks(cellGridMatrix, allTaus, patatracks_SoA, vertices_SoA, *bsHandle, fieldESH.product());

  standardizeTensor(cellGridMatrix);

  if (debugLevel_ > 0) {
    checknan(cellGridMatrix, debugLevel_);
  }

  std::vector<float> tau_score = getTauScore(cellGridMatrix);

  for (size_t inp_idx = 0; inp_idx < L1TauDesc_.size(); inp_idx++) {
    const size_t nTau = TauCollectionMap[inp_idx].size();
    auto tau_tags = std::make_unique<std::vector<float>>(nTau);
    for (size_t tau_pos = 0; tau_pos < nTau; ++tau_pos) {
      const auto tau_idx = TauCollectionMap[inp_idx][tau_pos];
      if (debugLevel_ > 0) {
        edm::LogInfo("DebugInfo") << event.id().event() << " \t " << (allTaus[tau_idx])->pt() << " \t "
                                  << tau_score.at(tau_idx) << std::endl;
      }
      (*tau_tags)[tau_pos] = tau_score.at(tau_idx);
    }
    event.put(std::move(tau_tags), L1TauDesc_[inp_idx].CollectionName);
  }
}
//define this as a plug-in
#include "FWCore/Framework/interface/MakerMacros.h"
DEFINE_FWK_MODULE(L2TauNNProducer);


/*
system.cpp

Copyright (c) 2014 Terumasa Tadano

This file is distributed under the terms of the MIT license.
Please see the file 'LICENCE.txt' in the root directory
or http://opensource.org/licenses/mit-license.php for information.
*/

#include "mpi_common.h"
#include "system.h"
#include "fcs_phonon.h"
#include "error.h"
#include "memory.h"
#include "constants.h"
#include "symmetry_core.h"
#include <string>
#include <iostream>
#include <iomanip>
#include <fstream>
#include "mathfunctions.h"
#include "xml_parser.h"
#include <sstream>
#include <map>
#include <boost/property_tree/ptree.hpp>
#include <boost/property_tree/xml_parser.hpp>
#include <boost/foreach.hpp>
#include <boost/optional.hpp>
#include <boost/lexical_cast.hpp>

using namespace PHON_NS;

System::System(PHON *phon): Pointers(phon) {}

System::~System() {
    memory->deallocate(xr_p);
    memory->deallocate(xr_s);
    memory->deallocate(xr_s_anharm);
    memory->deallocate(kd);
    memory->deallocate(kd_anharm);
    memory->deallocate(xc);
    memory->deallocate(mass);
    memory->deallocate(map_p2s);
    memory->deallocate(map_p2s_anharm);
    memory->deallocate(map_s2p);
    memory->deallocate(map_s2p_anharm);
    memory->deallocate(mass_kd);
    memory->deallocate(magmom);
}

void System::setup()
{
    using namespace std;

    unsigned int i, j;
    double vec_tmp[3][3];
    unsigned int *kd_prim;
    double **xtmp;

    load_system_info_from_XML();

    recips(lavec_s, rlavec_s);
    recips(lavec_s_anharm, rlavec_s_anharm);
    recips(lavec_p, rlavec_p);

    memory->allocate(xr_p, nat, 3);
    memory->allocate(xc, nat, 3);

    for (i = 0; i < nat; ++i){
        rotvec(xc[i], xr_s[i], lavec_s);
        rotvec(xr_p[i], xc[i], rlavec_p);
        for(j = 0; j < 3; ++j){
            xr_p[i][j] /=  2.0 * pi;
        }
    }

    if (mympi->my_rank == 0) {
        cout << " ------------------------------------------------------------" << endl;
        cout << endl;
        cout << " Crystal structure" << endl;
        cout << " =================" << endl << endl;
        cout << " Lattice Vectors:" << endl << endl;
        cout.setf(ios::scientific);

        cout << " * Supercell (from " << fcs_phonon->file_fcs << " )" << endl << endl;
        cout << setw(16) << lavec_s_anharm[0][0] << setw(15) << lavec_s_anharm[1][0] << setw(15) << lavec_s_anharm[2][0] << " : a1" << endl;
        cout << setw(16) << lavec_s_anharm[0][1] << setw(15) << lavec_s_anharm[1][1] << setw(15) << lavec_s_anharm[2][1] << " : a2" << endl;
        cout << setw(16) << lavec_s_anharm[0][2] << setw(15) << lavec_s_anharm[1][2] << setw(15) << lavec_s_anharm[2][2] << " : a3" << endl;
        cout << endl;

        cout << setw(16) << rlavec_s_anharm[0][0] << setw(15) << rlavec_s_anharm[0][1] << setw(15) << rlavec_s_anharm[0][2] << " : b1" << endl;
        cout << setw(16) << rlavec_s_anharm[1][0] << setw(15) << rlavec_s_anharm[1][1] << setw(15) << rlavec_s_anharm[1][2] << " : b2" << endl;
        cout << setw(16) << rlavec_s_anharm[2][0] << setw(15) << rlavec_s_anharm[2][1] << setw(15) << rlavec_s_anharm[2][2] << " : b3" << endl;
        cout << endl;

        cout << " * Primitive cell " << endl << endl;
        cout << setw(16) << lavec_p[0][0] << setw(15) << lavec_p[1][0] << setw(15) << lavec_p[2][0] << " : a1" << endl;
        cout << setw(16) << lavec_p[0][1] << setw(15) << lavec_p[1][1] << setw(15) << lavec_p[2][1] << " : a2" << endl;
        cout << setw(16) << lavec_p[0][2] << setw(15) << lavec_p[1][2] << setw(15) << lavec_p[2][2] << " : a3" << endl;
        cout << endl;

        cout << setw(16) << rlavec_p[0][0] << setw(15) << rlavec_p[0][1] << setw(15) << rlavec_p[0][2] << " : b1" << endl;
        cout << setw(16) << rlavec_p[1][0] << setw(15) << rlavec_p[1][1] << setw(15) << rlavec_p[1][2] << " : b2" << endl;
        cout << setw(16) << rlavec_p[2][0] << setw(15) << rlavec_p[2][1] << setw(15) << rlavec_p[2][2] << " : b3" << endl;
        cout << endl << endl;

        for (i = 0; i < 3; ++i){
            for (j = 0; j < 3; ++j){
                vec_tmp[i][j] = lavec_p[j][i];
            }
        }
        volume_p = volume(vec_tmp[0], vec_tmp[1], vec_tmp[2]);

        cout << "  Volume of the primitive cell : " << volume_p << " (a.u.)^3" << endl << endl;
        cout << "  Number of atoms in the supercell     : " << nat_anharm << endl;
        cout << "  Number of atoms in the primitive cell: " << natmin << endl << endl;

        if (fcs_phonon->update_fc2) {
            cout << endl;
            cout << "  FC2XML is given: Harmonic IFCs will be replaced by the values in " << fcs_phonon->file_fc2 << endl;
            cout << endl;

            cout << " * Supercell for HARMONIC (from " << fcs_phonon->file_fc2 << " )" << endl << endl;
            cout << setw(16) << lavec_s[0][0] << setw(15) << lavec_s[1][0] << setw(15) << lavec_s[2][0] << " : a1" << endl;
            cout << setw(16) << lavec_s[0][1] << setw(15) << lavec_s[1][1] << setw(15) << lavec_s[2][1] << " : a2" << endl;
            cout << setw(16) << lavec_s[0][2] << setw(15) << lavec_s[1][2] << setw(15) << lavec_s[2][2] << " : a3" << endl;
            cout << endl;

            cout << setw(16) << rlavec_s[0][0] << setw(15) << rlavec_s[0][1] << setw(15) << rlavec_s[0][2] << " : b1" << endl;
            cout << setw(16) << rlavec_s[1][0] << setw(15) << rlavec_s[1][1] << setw(15) << rlavec_s[1][2] << " : b2" << endl;
            cout << setw(16) << rlavec_s[2][0] << setw(15) << rlavec_s[2][1] << setw(15) << rlavec_s[2][2] << " : b3" << endl;
            cout << endl;

            cout << "  Number of atoms in the supercell (HARMONIC)   : " << nat << endl;
            cout << endl;
        }

        memory->allocate(xtmp, natmin, 3);

        for (i = 0; i < natmin; ++i) {
            rotvec(xtmp[i], xr_s[map_p2s[i][0]], lavec_s);
            rotvec(xtmp[i], xtmp[i], rlavec_p);
            for (j = 0; j < 3; ++j) xtmp[i][j] /= 2.0 * pi;
        }

        cout << "  Atomic positions in the primitive cell (fractional):" << endl;
        for (i = 0; i < natmin; ++i){
            cout << setw(4) << i + 1 << ":";
            for (j = 0; j < 3; ++j) {
                cout << setw(15) << xtmp[i][j];
            }
            cout << setw(4) << symbol_kd[kd[map_p2s[i][0]]] << endl;
        }
        cout << endl;

        memory->deallocate(xtmp);

        if (lspin) {
            cout << "  MagneticMoments entry found in the XML file. " << endl;
            cout << "  Magnetic moment in Cartesian coordinates: " << endl;
            for (i = 0; i < natmin; ++i) {
                cout << setw(4) << i + 1 << ":";
                for (j = 0; j < 3; ++j) {
                    cout << setw(15) << magmom[i][j];
                }
                cout << endl;
            }
            cout << endl;
            if (noncollinear == 0) {
                cout << "  Collinear calculation: magnetic moments are considered as scalar variables." << endl;
            } else if (noncollinear == 1) {
                cout << "  Noncollinear calculation: magnetic moments are considered as vector variables." << endl;
                if (symmetry->trev_sym_mag) {
                    cout << "  Time-reversal symmetry will be considered for generating magnetic space group" << endl;
                } else {
                    cout << "  Time-reversal symmetry will NOT be considered for generating magnetic space group" << endl;
                }
            }
            cout << endl;
        }

        cout << "  Mass of atomic species (u):" << endl;
        for (i = 0; i < nkd; ++i) {
            cout << setw(4) << symbol_kd[i] << ":";
            cout << fixed << setw(12) << mass_kd[i] << endl;
        }
        cout << endl << endl;
    }

    // Atomic masses in Rydberg unit

    memory->allocate(mass, nat);
    memory->allocate(mass_anharm, nat_anharm);
    for (i = 0; i < nat; ++i){
        mass[i] = mass_kd[kd[i]]*amu_ry;
    }
    for (i = 0; i < nat_anharm; ++i) {
        mass_anharm[i] = mass_kd[kd_anharm[i]]*amu_ry;
    }
    MPI_Bcast(&Tmin, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Bcast(&Tmax, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Bcast(&dT, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Bcast(&volume_p, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);

    memory->allocate(kd_prim, natmin);

    for (i = 0; i < natmin; ++i) {
        kd_prim[i] = kd[map_p2s[i][0]];
    }
    MPI_Bcast(&lspin, 1, MPI_LOGICAL, 0, MPI_COMM_WORLD);
    if (mympi->my_rank > 0) {
        memory->allocate(magmom, natmin, 3);
    }
    MPI_Bcast(&magmom[0][0], 3*natmin, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Bcast(&noncollinear, 1, MPI_INT, 0, MPI_COMM_WORLD);

    setup_atomic_class(natmin, kd_prim, magmom);

    memory->deallocate(kd_prim);
}

void System::load_system_info_from_XML()
{
    if (mympi->my_rank == 0) {

        int i;
        using namespace boost::property_tree;
        ptree pt;
        int nkd_tmp;

        std::map<std::string, int> dict_atomic_kind;

        try {
            read_xml(fcs_phonon->file_fcs, pt);
        }
        catch (std::exception &e) {
            std::string str_error = "Cannot open file FCSXML ( " + fcs_phonon->file_fcs + " )";
            error->exit("load_system_info_from_XML", str_error.c_str());
        }

        // Parse nat and ntran

        nat = boost::lexical_cast<unsigned int>(get_value_from_xml(pt, "Data.Structure.NumberOfAtoms"));
        nkd_tmp = boost::lexical_cast<unsigned int>(get_value_from_xml(pt, "Data.Structure.NumberOfElements"));

        if (nkd != nkd_tmp) error->exit("load_system_info_from_XML",
            "NKD in the FCSXML file is not consistent with that given in the input file.");

        ntran = boost::lexical_cast<unsigned int>(get_value_from_xml(pt, "Data.Symmetry.NumberOfTranslations"));

        natmin = nat / ntran;

        // Parse lattice vectors

        std::stringstream ss;

        for (i = 0; i < 3; ++i) {
            ss.str("");
            ss.clear();
            ss << get_value_from_xml(pt,
                "Data.Structure.LatticeVector.a" + boost::lexical_cast<std::string>(i + 1));
            ss >> lavec_s[0][i] >> lavec_s[1][i] >> lavec_s[2][i];
        }

        // Parse atomic elements and coordinates

        memory->allocate(xr_s, nat, 3);
        memory->allocate(kd, nat);

        BOOST_FOREACH (const ptree::value_type& child_, pt.get_child("Data.Structure.AtomicElements")) {
            const ptree& child = child_.second;
            const unsigned int icount_kd = child.get<unsigned int>("<xmlattr>.number");
            dict_atomic_kind[boost::lexical_cast<std::string>(child_.second.data())] = icount_kd - 1;
        }

        unsigned int index;

        BOOST_FOREACH (const ptree::value_type& child_, pt.get_child("Data.Structure.Position")) {
            const ptree& child = child_.second;
            const std::string str_index = child.get<std::string>("<xmlattr>.index");
            const std::string str_element = child.get<std::string>("<xmlattr>.element");

            ss.str("");
            ss.clear();
            ss << child.data();

            index = boost::lexical_cast<unsigned int>(str_index) - 1;

            if (index >= nat) error->exit("load_system_info_xml", "index is out of range");

            kd[index] = dict_atomic_kind[str_element];
            ss >> xr_s[index][0] >> xr_s[index][1] >> xr_s[index][2];
        }

        dict_atomic_kind.clear();

        // Parse mapping information

        memory->allocate(map_p2s, natmin, ntran);
        memory->allocate(map_s2p, nat);

        unsigned int tran, atom_p, atom_s;

        BOOST_FOREACH (const ptree::value_type& child_, pt.get_child("Data.Symmetry.Translations")) {
            const ptree& child = child_.second;
            const std::string str_tran = child.get<std::string>("<xmlattr>.tran");
            const std::string str_atom = child.get<std::string>("<xmlattr>.atom");

            tran = boost::lexical_cast<unsigned int>(str_tran) - 1;
            atom_p = boost::lexical_cast<unsigned int>(str_atom) - 1;
            atom_s = boost::lexical_cast<unsigned int>(child.data()) - 1;

            if (tran >= ntran || atom_p >= natmin || atom_s >= nat) {
                error->exit("load_system_info_xml", "index is out of range");
            }

            map_p2s[atom_p][tran] = atom_s;
            map_s2p[atom_s].atom_num = atom_p;
            map_s2p[atom_s].tran_num = tran;
        }

        // Parse magnetic moments

        double **magmom_tmp;
        memory->allocate(magmom_tmp, nat, 3);
        memory->allocate(magmom, natmin, 3);

        lspin = true;
        try {
            BOOST_FOREACH (const ptree::value_type& child_, pt.get_child("Data.MagneticMoments")) {
                if (child_.first == "mag") {
                    const ptree& child = child_.second;
                    const std::string str_index = child.get<std::string>("<xmlattr>.index");

                    ss.str("");
                    ss.clear();
                    ss << child.data();

                    index = boost::lexical_cast<unsigned int>(str_index) - 1;

                    if (index >= nat) error->exit("load_system_info_xml", "index is out of range");

                    ss >> magmom_tmp[index][0] >> magmom_tmp[index][1] >> magmom_tmp[index][2];
                }
            }

        } catch(...) {
            lspin = false;
        }

        if (lspin) {
            for (i = 0; i < natmin; ++i) {
                for (int j = 0; j < 3; ++j) {
                    magmom[i][j] = magmom_tmp[map_p2s[i][0]][j];
                }
            }

            try {
                noncollinear = boost::lexical_cast<int>(get_value_from_xml(pt, "Data.MagneticMoments.Noncollinear"));
            } catch(...) {
                noncollinear = 0;
            }

            try {
                symmetry->trev_sym_mag = boost::lexical_cast<int>(get_value_from_xml(pt, "Data.MagneticMoments.TimeReversalSymmetry"));
            } catch(...) {
                symmetry->trev_sym_mag = true;
            }
        } else {
            for (i = 0; i < natmin; ++i) {
                for (int j = 0; j < 3; ++j) {
                    magmom[i][j] = 0.0;
                }
            }
            noncollinear = 0;
            symmetry->trev_sym_mag = true;
        }
        memory->deallocate(magmom_tmp);

        // Now, replicate the information for anharmonic terms.

        int j;
        nat_anharm = nat;
        ntran_anharm = ntran;
        memory->allocate(xr_s_anharm, nat_anharm, 3);
        memory->allocate(kd_anharm, nat_anharm);
        memory->allocate(map_p2s_anharm, natmin, ntran_anharm);
        memory->allocate(map_s2p_anharm, nat_anharm);

        for (i = 0; i < 3; ++i) {
            for (j = 0; j < 3; ++j) lavec_s_anharm[i][j] = lavec_s[i][j];
        }
        for (i = 0; i < nat_anharm; ++i) {
            for (j = 0; j < 3; ++j) xr_s_anharm[i][j] = xr_s[i][j];
            kd_anharm[i] = kd[i];
            map_s2p_anharm[i] = map_s2p[i];
        }
        for (i = 0; i < natmin; ++i) {
            for (j = 0; j < ntran_anharm; ++j) {
                map_p2s_anharm[i][j] = map_p2s[i][j];
            }
        }

        if (fcs_phonon->update_fc2) {

            // When FC2XML is given, structural information is updated only for harmonic terms.

            int natmin_tmp;

            try {
                read_xml(fcs_phonon->file_fc2, pt);
            }
            catch (std::exception &e) {
                std::string str_error = "Cannot open file FC2XML ( " + fcs_phonon->file_fc2 + " )";
                error->exit("load_system_info_from_XML", str_error.c_str());
            }

            // Parse nat and ntran

            nat = boost::lexical_cast<unsigned int>(get_value_from_xml(pt, "Data.Structure.NumberOfAtoms"));
            nkd_tmp = boost::lexical_cast<unsigned int>(get_value_from_xml(pt, "Data.Structure.NumberOfElements"));

            if (nkd != nkd_tmp) error->exit("load_system_info_from_XML",
                "NKD in the FC2XML file is not consistent with that given in the input file.");

            ntran = boost::lexical_cast<unsigned int>(get_value_from_xml(pt, "Data.Symmetry.NumberOfTranslations"));

            natmin_tmp = nat / ntran;

            if (natmin_tmp != natmin) error->exit("load_system_info_from_XML",
                "Number of atoms in a primitive cell is different in FCSXML and FC2XML.");

            memory->deallocate(xr_s);
            memory->deallocate(kd);
            memory->deallocate(map_p2s);
            memory->deallocate(map_s2p);

            // Parse lattice vectors

            std::stringstream ss;

            for (i = 0; i < 3; ++i) {
                ss.str("");
                ss.clear();
                ss << get_value_from_xml(pt,
                    "Data.Structure.LatticeVector.a" + boost::lexical_cast<std::string>(i + 1));
                ss >> lavec_s[0][i] >> lavec_s[1][i] >> lavec_s[2][i];
            }

            // Parse atomic elements and coordinates

            memory->allocate(xr_s, nat, 3);
            memory->allocate(kd, nat);

            BOOST_FOREACH (const ptree::value_type& child_, pt.get_child("Data.Structure.AtomicElements")) {
                const ptree& child = child_.second;
                const unsigned int icount_kd = child.get<unsigned int>("<xmlattr>.number");
                dict_atomic_kind[boost::lexical_cast<std::string>(child_.second.data())] = icount_kd - 1;
            }

            unsigned int index;

            BOOST_FOREACH (const ptree::value_type& child_, pt.get_child("Data.Structure.Position")) {
                const ptree& child = child_.second;
                const std::string str_index = child.get<std::string>("<xmlattr>.index");
                const std::string str_element = child.get<std::string>("<xmlattr>.element");

                ss.str("");
                ss.clear();
                ss << child.data();

                index = boost::lexical_cast<unsigned int>(str_index) - 1;

                if (index >= nat) error->exit("load_system_info_xml", "index is out of range");

                kd[index] = dict_atomic_kind[str_element];
                ss >> xr_s[index][0] >> xr_s[index][1] >> xr_s[index][2];
            }

            dict_atomic_kind.clear();

            // Parse mapping information

            memory->allocate(map_p2s, natmin, ntran);
            memory->allocate(map_s2p, nat);

            unsigned int tran, atom_p, atom_s;

            BOOST_FOREACH (const ptree::value_type& child_, pt.get_child("Data.Symmetry.Translations")) {
                const ptree& child = child_.second;
                const std::string str_tran = child.get<std::string>("<xmlattr>.tran");
                const std::string str_atom = child.get<std::string>("<xmlattr>.atom");

                tran = boost::lexical_cast<unsigned int>(str_tran) - 1;
                atom_p = boost::lexical_cast<unsigned int>(str_atom) - 1;
                atom_s = boost::lexical_cast<unsigned int>(child.data()) - 1;

                if (tran >= ntran || atom_p >= natmin || atom_s >= nat) {
                    error->exit("load_system_info_xml", "index is out of range");
                }

                map_p2s[atom_p][tran] = atom_s;
                map_s2p[atom_s].atom_num = atom_p;
                map_s2p[atom_s].tran_num = tran;
            }
        }
    }

    MPI_Bcast(&lavec_s[0][0], 9, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Bcast(&lavec_p[0][0], 9, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Bcast(&lavec_s_anharm[0][0], 9, MPI_DOUBLE, 0, MPI_COMM_WORLD);

    MPI_Bcast(&nkd, 1, MPI_UNSIGNED, 0, MPI_COMM_WORLD);
    MPI_Bcast(&nat, 1, MPI_UNSIGNED, 0, MPI_COMM_WORLD);
    MPI_Bcast(&nat_anharm, 1, MPI_UNSIGNED, 0, MPI_COMM_WORLD);
    MPI_Bcast(&natmin, 1, MPI_UNSIGNED, 0, MPI_COMM_WORLD);
    MPI_Bcast(&ntran, 1, MPI_UNSIGNED, 0, MPI_COMM_WORLD);
    MPI_Bcast(&ntran_anharm, 1, MPI_UNSIGNED, 0, MPI_COMM_WORLD);
    MPI_Bcast(&lspin, 1, MPI_LOGICAL, 0, MPI_COMM_WORLD);

    if (mympi->my_rank > 0){
        memory->allocate(mass_kd, nkd);
        memory->allocate(xr_s, nat, 3);
        memory->allocate(xr_s_anharm, nat_anharm, 3);
        memory->allocate(kd, nat);
        memory->allocate(kd_anharm, nat_anharm);
        memory->allocate(map_p2s, natmin, ntran);
        memory->allocate(map_p2s_anharm, natmin, ntran_anharm);
        memory->allocate(map_s2p, nat);
        memory->allocate(map_s2p_anharm, nat_anharm);
        if (lspin) {
            memory->allocate(magmom, natmin, 3);
        }
    }

    MPI_Bcast(&mass_kd[0], nkd, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Bcast(&xr_s[0][0], 3*nat, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Bcast(&xr_s_anharm[0][0], 3*nat_anharm, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Bcast(&kd[0], nat, MPI_UNSIGNED, 0, MPI_COMM_WORLD);
    MPI_Bcast(&kd_anharm[0], nat_anharm, MPI_UNSIGNED, 0, MPI_COMM_WORLD);
    MPI_Bcast(&map_p2s[0][0], natmin*ntran, MPI_UNSIGNED, 0, MPI_COMM_WORLD);
    MPI_Bcast(&map_p2s_anharm[0][0], natmin*ntran_anharm, MPI_UNSIGNED, 0, MPI_COMM_WORLD);
    MPI_Bcast(&map_s2p[0], nat*sizeof(map_s2p[0]), MPI_BYTE, 0, MPI_COMM_WORLD);
    MPI_Bcast(&map_s2p_anharm[0], nat_anharm*sizeof(map_s2p_anharm[0]), MPI_BYTE, 0, MPI_COMM_WORLD);
    if (lspin) MPI_Bcast(&magmom[0][0], 3*natmin, MPI_DOUBLE, 0, MPI_COMM_WORLD);

}

void System::recips(double vec[3][3], double inverse[3][3])
{
    double det;
    det = vec[0][0] * vec[1][1] * vec[2][2]
    + vec[1][0] * vec[2][1] * vec[0][2]
    + vec[2][0] * vec[0][1] * vec[1][2]
    - vec[0][0] * vec[2][1] * vec[1][2]
    - vec[2][0] * vec[1][1] * vec[0][2]
    - vec[1][0] * vec[0][1] * vec[2][2];

    if(std::abs(det) < eps12) {
        error->exit("recips", "Lattice Vector is singular");
    }

    double factor = 2.0 * pi / det;

    inverse[0][0] = (vec[1][1] * vec[2][2] - vec[1][2] * vec[2][1]) * factor;
    inverse[0][1] = (vec[0][2] * vec[2][1] - vec[0][1] * vec[2][2]) * factor;
    inverse[0][2] = (vec[0][1] * vec[1][2] - vec[0][2] * vec[1][1]) * factor;

    inverse[1][0] = (vec[1][2] * vec[2][0] - vec[1][0] * vec[2][2]) * factor;
    inverse[1][1] = (vec[0][0] * vec[2][2] - vec[0][2] * vec[2][0]) * factor;
    inverse[1][2] = (vec[0][2] * vec[1][0] - vec[0][0] * vec[1][2]) * factor;

    inverse[2][0] = (vec[1][0] * vec[2][1] - vec[1][1] * vec[2][0]) * factor;
    inverse[2][1] = (vec[0][1] * vec[2][0] - vec[0][0] * vec[2][1]) * factor;
    inverse[2][2] = (vec[0][0] * vec[1][1] - vec[0][1] * vec[1][0]) * factor;
}

double System::volume(double vec1[3], double vec2[3], double vec3[3])
{
    double vol;

    vol = std::abs(vec1[0]*(vec2[1]*vec3[2] - vec2[2]*vec3[1])
        + vec1[1]*(vec2[2]*vec3[0] - vec2[0]*vec3[2])
        + vec1[2]*(vec2[0]*vec3[1] - vec2[1]*vec3[0]));

    return vol;
}

void System::setup_atomic_class(unsigned int N, unsigned int *kd, double **magmom_in) {

    // In the case of collinear calculation, spin moments are considered as scalar
    // variables. Therefore, the same elements with different magnetic moments are
    // considered as different types. In noncollinear calculations,
    // magnetic moments are not considered in this stage. They will be treated
    // separately in symmetry.cpp where spin moments will be rotated and flipped
    // using time-reversal symmetry.

    unsigned int i;
    AtomType type_tmp;
    std::set<AtomType> set_type;
    set_type.clear();

    for (i = 0; i < N; ++i) {
        type_tmp.element = kd[i];

        if (noncollinear == 0) {
            type_tmp.magmom = magmom_in[i][2];
        } else {
            type_tmp.magmom = 0.0;
        }
        set_type.insert(type_tmp);
    }

    nclassatom = set_type.size();

    memory->allocate(atomlist_class, nclassatom);

    for (i = 0; i < N; ++i) {
        int count = 0;
        for (std::set<AtomType>::iterator it = set_type.begin(); it != set_type.end(); ++it) {
            if (noncollinear) {
                if (kd[i] == (*it).element) {
                    atomlist_class[count].push_back(i);
                }
            } else {
                if (kd[i] == (*it).element && std::abs(magmom[i][2] - (*it).magmom) < eps6) {
                    atomlist_class[count].push_back(i);
                }
            }
            ++count;
        }
    }
    set_type.clear();
}