/**
 * Copyright 2017, IBM.
 *
 * This source code is licensed under the Apache License, Version 2.0 found in
 * the LICENSE.txt file in the root directory of this source tree.
 */

/**
 * @file    noise_models.hpp
 * @brief   Noise Models for Simulator Backends
 * @author  Christopher J. Wood <cjwood@us.ibm.com>
 */

#ifndef _NoiseModels_h_
#define _NoiseModels_h_

#include <complex>
#include <map>
#include <random>
#include <stdexcept>
#include <vector>

#include "misc.hpp"
#include "types.hpp"

namespace QISKIT {

/*******************************************************************************
 *
 * Relaxation Class
 *
 ******************************************************************************/

class Relaxation {
public:
  /**
   * Relaxation data members
   */
  double rate = 0.;                            // relxation rate
  std::discrete_distribution<> populations{1}; // equilibrium populations

  /**
   * Computes the relaxation probability for a length of time assuming
   * exponential relaxation: p(t) = 1-exp(-t*rate).
   * @param t: the length of time
   * @returns: relaxation probability p(t)
   */
  double p(double t) const;

  /**
   * Default constructor: no relaxation error
   */
  Relaxation(){};

  /**
   * Constructor: sets the relaxation rate and equilibrium population that
   * system will relax towards.
   * @param rate1: the relaxation rate
   * @param pops: equilibrium populations as a discrete_distribution
   */
  Relaxation(double rate1, const std::discrete_distribution<> &pops);

  /**
   * Constructor: sets the relaxation rate and equilibrium population that
   * system will relax towards. If the input populations are not normalized
   * they will be rescaled when converted to a discrete_distribution.
   * @param rate1: the relaxation rate
   * @param pops: vector of equilibrium populations
   */
  Relaxation(double rate1, const rvector_t &pops);

  /**
   * Verify that the set noise parameters are valid for the simulation by
   * checking that the relaxation rate is non-negative, and that the vector of
   * probabilities for the equilibrium distribution is not greater than the
   * subsystem dimension dim.
   * @param dim: dimension of each subsystem (typically 2 for a qubit)
   * @returns: true if noise parameters are valid
   */
  bool verify(uint_t dim) const;
};

/*******************************************************************************
 *
 * Reset Error Class
 *
 * This class specifies errors in the QASM reset operation. The reset error sets
 * a system to the wrong state with some probability following a reset.
 *
 ******************************************************************************/

class ResetError {
public:
  /***
   * ResetError data members
   ***/
  bool ideal = true;                 // flag if reset error is set
  std::discrete_distribution<> p{1}; // the reset error probabilities

  /***
   * Default constructor: no reset error
   ***/
  ResetError(){};

  /***
   * Constructor: sets the probabilities of resetting to each energy level using
   * a discrete_distribution input.
   * @param probs: reset probabilities as a discrete_distribution
   ***/
  ResetError(const std::discrete_distribution<> &probs);

  /***
   * Constructor: sets the probabilities of resetting to each energy level using
   * a vector of probabilities as input. If the input probabilities are not
   * normalized they will be rescaled when converted to a discrete_distribution.
   * @param probs: reset probabilities as a vector
   ***/
  ResetError(const rvector_t &probs);

  /***
   * Verify that the set noise parameters are valid for the simulation by
   * checking that the length of the vector of reset probabilities is not
   *greater
   * than the subsystem dimension dim.
   * @param dim: dimension of each subsystem (typically 2 for a qubit)
   * @returns: true if noise parameters are valid
   ***/
  bool verify(uint_t dim) const;
};

/*******************************************************************************
 *
 * ReadoutError Class
 *
 * This class specifies errors in recording the value of the classical bit after
 * performing a QASM measure operation. It is an assignment error where the
 * probability of recording a measurement outcome as k, given that m was
 * actually obtained is given by P(k|m).
 *
 ******************************************************************************/

class ReadoutError {
public:
  /***
   * ResetError data members
   ***/
  bool ideal = true;                           // flag if reset error is set
  std::vector<std::discrete_distribution<>> p; // the assignment error probs

  /***
   * Default constructor: no readout error
   ***/
  ReadoutError(){};

  /***
   * Constructor: sets the readout assignment probablities from a vector of
   * discrete_distributions for each measurement outcome.
   * @param probs: vector of discrete_distributions
   ***/
  ReadoutError(const std::vector<std::discrete_distribution<>> &probs);

  /***
   * Constructor: sets the readout assignment probablities to be of the form
   *  P(0|0) = 1-p1, P(1|0) = p1 for a 0 measurement outcome
   *  P(0|1) = p1, P(1|1) = 1-p1 for a 1 measurement outcome
   * @param p1: assignement error probability
   ***/
  explicit ReadoutError(double p1);

  /***
   * Constructor: sets the readout assignment probablities from a vector of
   * probabilty vectors for each measurement outcome. If any probability
   * vector is not correctly normalized it will be rescaled when converted to
   * a discrete_distribution.
   * @param probs: vector of probability vectors
   ***/
  ReadoutError(const std::vector<rvector_t> &probs);

  /***
   * Verify that the set noise parameters are valid for the simulation by
   * checking that the vector of assignment probabilities is the correct length
   *for
   * the subsystem dimension dim.
   * @param dim: dimension of each subsystem (typically 2 for a qubit)
   * @returns: true if noise parameters are valid
   ***/
  bool verify(uint_t dim) const;
};

/*******************************************************************************
 *
 * PauliChannel Class
 *
 ******************************************************************************/

class PauliChannel {
public:
  uint_t n = 0;
  bool ideal = true;
  std::discrete_distribution<> p;

  // Constructors
  PauliChannel(){};
  explicit PauliChannel(uint_t nq) : n(nq){};
  PauliChannel(uint_t nq, rvector_t p_pauli);

  // Get vector of error probabilities eg {pX, pY, pZ}
  rvector_t p_errors() const;
  void set_p(rvector_t p_pauli);
};

/*******************************************************************************
 *
 * GateError Class
 *
 ******************************************************************************/

class GateError {
public:
  /***
   * GateError data members
   ***/
  std::string label; // label for gate error
  bool ideal = true; // flag if any gate error is set

  /***
   * Incoherent Error Channel parameters
   ***/
  PauliChannel pauli;
  double gate_time = 0.;

  /***
   * Unitary Coherent Error Channel parameters
   ***/
  bool coherent_error = false; // flag if coherent gate error is set
  cmatrix_t Uerr;              // Coherent error to be applied after gate

  /***
   * Default constructor: no gate error
   ***/
  GateError(){};

  /***
   * Verify that the set noise parameters are valid for the simulation by
   * checking that all probabilities are non-negative, the length of pauli
   *error
   * probabiltity vectors are the correct for the size of the gate, and the
   * coherent error matrices are the correct dimension for the subsystem size
   * @param dim: dimension of each subsystem (typically 2 for a qubit)
   * @returns: true if noise parameters are valid
   ***/
  bool verify(uint_t dim) const;
};

/*******************************************************************************
 *
 * QubitNoise Class
 *
 ******************************************************************************/

class QubitNoise {
public:
  /***
   * Noise members
   ***/
  bool ideal = true;    // flag if any error is set
  ResetError reset;     // noise params used by QASM reset
  ReadoutError readout; // noise params used by QASM measure
  Relaxation relax;     // relaxation parameters used by other errors
  std::map<std::string, GateError> gate; // noise params for each gate type
  const static std::vector<std::string> gate_names;

  /***
   * Default constructor: no gate error
   ***/
  QubitNoise(){};

  /***
   * Verify that each of the noise class members has valid parameters by
   *calling
   * the corresponding verify method of each member.
   * @param dim: dimension of each subsystem (typically 2 for a qubit)
   * @returns: true if noise parameters are valid
   ***/
  bool verify(uint_t dim = 2);
};

/*******************************************************************************
 *
 * Relaxation Class Methods
 *
 ******************************************************************************/

double Relaxation::p(double t) const {
  return (rate > 0. && t > 0.) ? 1. - std::exp(-t * rate) : 0.;
}

Relaxation::Relaxation(double rate1, const std::discrete_distribution<> &pops)
    : rate(rate1), populations(pops) {
#ifdef DEBUG
  std::cout << "DEBUG Constructor: Relaxation" << std::endl;
  std::cout << "DEBUG relaxation_rate = " << rate << std::endl;
  std::cout << "DEBUG thermal_populations = " << populations.probabilities()
            << std::endl;
#endif
}

Relaxation::Relaxation(double rate1, const rvector_t &pops)
    : Relaxation(rate1,
                 std::discrete_distribution<>(pops.begin(), pops.end())) {}

bool Relaxation::verify(uint_t dim) const {
  // Check Relaxation
  if (populations.probabilities().size() > dim) {
    std::cerr << "error: thermal_populations vector is too long" << std::endl;
    return false;
  } else if (rate < 0) {
    std::cerr << "error: relaxation_rate is negative" << std::endl;
    return false;
  } else
    return true;
}

inline void to_json(json_t &js, const Relaxation &error) {
  if (error.rate > 0.) {
    js["relaxation_rate"] = error.rate;
    js["thermal_populations"] = error.populations.probabilities();
  }
}

inline void from_json(const json_t &noise, Relaxation &error) {
  if (JSON::check_key("relaxation_rate", noise)) {
    double rate = noise["relaxation_rate"];
    rvector_t pops{1.};
    if (JSON::check_key("thermal_populations", noise)) {
      const json_t &node = noise["thermal_populations"];
      if (node.is_number())
        pops[0] = node.get<double>();
      else
        pops = node.get<rvector_t>();
    }
    error = Relaxation(rate, pops);
  } else
    error = Relaxation();
}

/*******************************************************************************
 *
 * ResetError Methods
 *
 ******************************************************************************/

ResetError::ResetError(const std::discrete_distribution<> &probs) : p(probs) {
  ideal = !(p.probabilities()[0.] < 1.);
#ifdef DEBUG
  std::cout << "DEBUG Constructor: ResetError" << std::endl;
  std::cout << "DEBUG p = " << p.probabilities() << std::endl;
#endif
}

ResetError::ResetError(const rvector_t &probs)
    : ResetError(std::discrete_distribution<>(probs.begin(), probs.end())) {}

bool ResetError::verify(uint_t dim) const {
  if (p.probabilities().size() > dim) {
    std::cerr << "error: reset.p error vector is too long" << std::endl;
    return false;
  } else
    return true;
}

inline void to_json(json_t &js, const ResetError &error) {
  if (!error.ideal && error.p.probabilities()[0] < 1.) {
    js["reset_error"] = error.p.probabilities();
  }
}

inline void from_json(const json_t &noise, ResetError &error) {
  if (JSON::check_key("reset_error", noise)) {
    // Scalar value is p0, 1-p0 reset error
    if (noise["reset_error"].is_number()) {
      double p1 = noise["reset_error"].get<double>();
      error = (p1 > 0.) ? ResetError(std::discrete_distribution<>({1 - p1, p1}))
                        : ResetError();
    }
    // node is a vector of reset probs
    else if (noise["reset_error"].is_array()) {
      error = ResetError(noise["reset_error"].get<rvector_t>());
    } else {
      throw std::runtime_error(std::string("p_reset error invalid input"));
    }
  } else
    error = ResetError();
}

/*******************************************************************************
 *
 * ReadoutError Methods
 *
 ******************************************************************************/

ReadoutError::ReadoutError(
    const std::vector<std::discrete_distribution<>> &probs)
    : ideal(false), p(probs) {
#ifdef DEBUG
  std::cout << "DEBUG Constructor: ReadoutError" << std::endl;
  for (const auto &v : p)
    std::cout << "DEBUG p = " << v.probabilities() << std::endl;
#endif
}

ReadoutError::ReadoutError(double p1) {
  if (p1 > 0.) {
    ideal = false;
    p.push_back(std::discrete_distribution<>{1. - p1, p1});
    p.push_back(std::discrete_distribution<>{p1, 1. - p1});
  }
#ifdef DEBUG
  std::cout << "DEBUG Constructor: ReadoutError" << std::endl;
  for (const auto &v : p)
    std::cout << "DEBUG p = " << v.probabilities() << std::endl;
#endif
}

ReadoutError::ReadoutError(const std::vector<rvector_t> &probs) {
  ideal = false;
  for (auto const &v : probs)
    p.push_back(std::discrete_distribution<>(v.begin(), v.end()));

#ifdef DEBUG
  std::cout << "DEBUG Constructor: ReadoutError" << std::endl;
  for (const auto &v : p)
    std::cout << "DEBUG p = " << v.probabilities() << std::endl;
#endif
}

bool ReadoutError::verify(uint_t dim) const {
  if (p.size() > dim) {
    std::cerr << "error: readout_error.p error vector is too long" << std::endl;
    return false;
  } else {
    for (const auto &d : p) {
      if (d.probabilities().size() > dim) {
        std::cerr << "error: readout_error.p error vector is too long"
                  << std::endl;
        return false;
      }
    }
    return true;
  }
}

inline void to_json(json_t &js, const ReadoutError &error) {
  if (!error.ideal) {
    std::vector<rvector_t> probs;
    for (auto const &dist : error.p)
      probs.push_back(dist.probabilities());
    js["readout_error"] = probs;
  }
}

inline void from_json(const json_t &noise, ReadoutError &error) {
  if (JSON::check_key("readout_error", noise)) {
    const json_t &node = noise["readout_error"];
    // Scalar value is p0, 1-p0 reset error
    if (node.is_number()) {
      error = ReadoutError(node.get<double>());
    } else if (node.is_array()) {
      // check length of array
      if (node[0].is_number() && node.size() == 2) {
        // 2-outcome only assignment errors
        double p0 = node[0];
        double p1 = node[1];
        if (p1 > 0. || p0 > 0.) {
          std::vector<std::discrete_distribution<>> ps;
          ps.push_back(std::discrete_distribution<>({1 - p0, p0}));
          ps.push_back(std::discrete_distribution<>({p1, 1 - p1}));
          error = ReadoutError(ps);
        } else {
          error = ReadoutError();
        }
      } else {
        std::vector<rvector_t> ps = node;
        error = ReadoutError(ps);
      }
    } else {
      throw std::runtime_error(std::string("p_meas_error vector invalid"));
    }
  } else
    error = ReadoutError();
}

/*******************************************************************************
 *
 * PauliChannel Methods
 *
 ******************************************************************************/

PauliChannel::PauliChannel(uint_t nq, rvector_t p_pauli) : n(nq) {
  set_p(p_pauli);
}

void PauliChannel::set_p(rvector_t p_pauli) {
  // Get pI error prob
  uint_t N = 1ULL << (2 * n);
  double tot = 0;
  for (auto &pj : p_pauli)
    tot += pj;
  p_pauli.insert(p_pauli.begin(), std::max(0., 1. - tot));
  // Check vector
  if (p_pauli.size() > N || tot > 1. || tot < 0.)
    throw std::runtime_error("invalid Pauli vector");
  if (p_pauli[0] < 1.) {
    p = std::discrete_distribution<>(p_pauli.begin(), p_pauli.end());
    ideal = false;
  }
}

rvector_t PauliChannel::p_errors() const {
  const auto d = (1ULL << n);
  rvector_t ret = p.probabilities();
  ret.erase(ret.begin()); // remove pI element
  ret.resize(d * d - 1);  // resize
  return ret;
}

/*******************************************************************************
 *
 * GateError Methods
 *
 ******************************************************************************/

bool GateError::verify(uint_t dim) const {
  if (pauli.p.probabilities().size() > dim * dim) {
    std::cerr << "error: pauli error vector is wrong length." << std::endl;
    return false;
  } else if (gate_time < 0.) {
    std::cerr << "error: gate_time must be non-negative" << std::endl;
    return false;
  } else
    return true;
}

inline void to_json(json_t &js, const GateError &error) {
  if (!error.ideal) {
    json_t node;
    if (error.gate_time > 0.)
      node["gate_time"] = error.gate_time;
    if (error.pauli.ideal == false)
      node["p_pauli"] = error.pauli.p_errors();
    if (error.coherent_error && error.Uerr.size() > 0)
      node["U_error"] = error.Uerr;
    js[error.label] = node;
  }
}

inline GateError load_gate_error(std::string key, uint_t nq, const json_t &js) {

  GateError error = GateError();
  error.label = key;
  json_t noise;
  if (JSON::check_key(key, js))
    noise = js[key];
  else
    noise = js;

  // Load Gate Time
  JSON::get_value(error.gate_time, "gate_time", noise);
  if (error.gate_time > 0.)
    error.ideal = false;

  // Load Coherent Error
  if (JSON::check_key("U_error", noise)) {
    error.ideal = false;
    error.coherent_error = true;
    error.Uerr = noise["U_error"];
  }

  // Load Pauli Error
  double p_depol = 0.;
  rvector_t p_pauli;
  uint_t N = 1ULL << (2 * nq);
  JSON::get_value(p_pauli, "p_pauli", noise);
  if (JSON::get_value(p_depol, "p_depol", noise) && p_depol > 0.) {
    p_pauli.resize(N - 1);
    for (auto &p : p_pauli)
      p = p + (p_depol / double(N)) - (p_depol * p);
  }
  error.pauli = PauliChannel(nq, p_pauli);
  error.ideal &= error.pauli.ideal;
  return error;
}

/*******************************************************************************
 *
 * QubitNoise Methods
 *
 ******************************************************************************/

bool QubitNoise::verify(uint_t dim) {
  bool pass = reset.verify(dim) && readout.verify(dim) && relax.verify(dim);

  for (const auto &g : gate) {
    uint_t dim2 = (g.first == "CX" || g.first == "CZ") ? dim * dim : dim;
    pass = pass && g.second.verify(dim2);
  }
  return pass;
}

const std::vector<std::string>
    QubitNoise::gate_names({"X90", "CX", "CZ", "id", "U", "measure", "reset"});

inline void to_json(json_t &js, const QubitNoise &noise) {
  json_t node;
  to_json(node, noise.relax);
  to_json(node, noise.reset);
  to_json(node, noise.readout);
  for (const auto &g : noise.gate) {
    to_json(node, g.second);
  }
  js = node;
}

inline void from_json(const json_t &js, QubitNoise &noise) {
  if (JSON::check_key("noise_params", js)) {
    from_json(js["noise_params"], noise);
  } else {
    noise = QubitNoise();
    from_json(js, noise.reset);
    noise.ideal &= noise.reset.ideal;
    noise.readout = js;
    noise.ideal &= noise.readout.ideal;
    noise.relax = js;
    noise.ideal &= !(noise.relax.rate > 0.);
    for (const auto &n : QubitNoise::gate_names) {
      if (JSON::check_key(n, js)) {
        GateError g;
        if (n == "CX" || n == "CZ")
          g = load_gate_error(n, 2, js);
        else
          g = load_gate_error(n, 1, js);

        // Check coherent error
        if (g.coherent_error) {
          cmatrix_t check = MOs::Dagger(g.Uerr) * g.Uerr;
          double threshold = 1e-10;
          double delta = 0.;
          for (size_t i=0; i < check.GetRows(); i++)
            for (size_t j=0; j < check.GetColumns(); j++) {
              complex_t val = (i==j) ? 1. : 0.;
              delta += std::real(std::abs(check(i, j) - val));
            }
          if (delta > threshold) {
            throw std::runtime_error(std::string(g.label + " U_error is not unitary"));
          }
        }

        // Add gate to noise model
        noise.gate.insert(std::make_pair(n, g));
        noise.ideal &= g.ideal;
      }
    }
  }
}

} // end namespace QISKIT
#endif