/*
 * Copyright (c) 2023 - 2025 Chair for Design Automation, TUM
 * Copyright (c) 2025 Munich Quantum Software Company GmbH
 * All rights reserved.
 *
 * SPDX-License-Identifier: MIT
 *
 * Licensed under the MIT License
 */

#include "DeterministicNoiseSimulator.hpp"

#include "Simulator.hpp"
#include "dd/ComplexNumbers.hpp"
#include "dd/DDDefinitions.hpp"
#include "dd/Operations.hpp"
#include "ir/operations/NonUnitaryOperation.hpp"
#include "ir/operations/OpType.hpp"
#include "ir/operations/Operation.hpp"

#include <cstddef>
#include <cstdint>
#include <iterator>
#include <map>
#include <memory>
#include <random>
#include <string>
#include <vector>

using CN = dd::ComplexNumbers;

void DeterministicNoiseSimulator::initializeSimulation(
    const std::size_t nQubits) {
  rootEdge = dd->makeZeroDensityOperator(static_cast<dd::Qubit>(nQubits));
}

void DeterministicNoiseSimulator::applyOperationToState(
    std::unique_ptr<qc::Operation>& op) {
  auto operation = dd::getDD(*op, *Simulator::dd);
  dd->applyOperationToDensity(DeterministicNoiseSimulator::rootEdge, operation);
  deterministicNoiseFunctionality.applyNoiseEffects(
      DeterministicNoiseSimulator::rootEdge, op);
}

char DeterministicNoiseSimulator::measure(const dd::Qubit i) {
  return Simulator::dd->measureOneCollapsing(
      rootEdge, static_cast<dd::Qubit>(i), Simulator::mt);
}

void DeterministicNoiseSimulator::reset(qc::NonUnitaryOperation* nonUnitaryOp) {
  for (const auto& qubit : nonUnitaryOp->getTargets()) {
    auto const result =
        dd->measureOneCollapsing(rootEdge, static_cast<dd::Qubit>(qubit), mt);
    if (result == '1') {
      const auto x = qc::StandardOperation(qubit, qc::X);
      const auto operation = dd::getDD(x, *dd);
      rootEdge = dd->applyOperationToDensity(rootEdge, operation);
    }
  }
}

std::map<std::string, std::size_t>
DeterministicNoiseSimulator::sampleFromProbabilityMap(
    const dd::SparsePVecStrKeys& resultProbabilityMap,
    const std::size_t shots) {
  std::vector<dd::fp> weights;
  weights.reserve(resultProbabilityMap.size());
  for (const auto& [state, prob] : resultProbabilityMap) {
    weights.emplace_back(prob);
  }
  std::discrete_distribution<std::size_t> d(
      weights.begin(),
      weights.end()); // NOLINT(misc-const-correctness) false-positive

  // Create the final map containing the measurement results and the
  // corresponding shots
  std::map<std::string, std::size_t> results;
  for (size_t n = 0; n < shots; ++n) {
    const auto sampleIdx = d(mt);
    const auto state = (std::next(resultProbabilityMap.begin(),
                                  static_cast<std::int64_t>(sampleIdx)))
                           ->first;
    results[state] += 1;
  }

  return results;
}
