"""
********************************************************************************
* Copyright (c) 2024 the Qrisp authors
*
* This program and the accompanying materials are made available under the
* terms of the Eclipse Public License 2.0 which is available at
* http://www.eclipse.org/legal/epl-2.0.
*
* This Source Code may also be made available under the following Secondary
* Licenses when the conditions for such availability set forth in the Eclipse
* Public License, v. 2.0 are satisfied: GNU General Public License, version 2
* with the GNU Classpath Exception which is
* available at https://www.gnu.org/software/classpath/license.html.
*
* SPDX-License-Identifier: EPL-2.0 OR GPL-2.0 WITH Classpath-exception-2.0
********************************************************************************
"""

import math
import threading

import numpy as np
from numba import njit

from qrisp.core import QuantumVariable, QuantumArray
from qrisp.core.compilation import qompiler
from qrisp.interface import BatchedBackend


def get_measurement(
    hamiltonian,
    qarg,
    precision=0.01,
    backend=None,
    compile=True,
    compilation_kwargs={},
    subs_dic={},
    precompiled_qc=None,
    diagonalisation_method="commuting_qw",
    measurement_data=None,  # measurement settings
):
    r"""
    This method returns the expected value of a Hamiltonian for the state of a quantum argument.

    Parameters
    ----------
    qarg : :ref:`QuantumVariable` or list[Qubit]
        The quantum argument to evaluate the Hamiltonian on.
    precision: float, optional
        The precision with which the expectation of the Hamiltonian is to be evaluated.
        The default is 0.01. The number of shots scales quadratically with the inverse precision.
    backend : :ref:`BackendClient`, optional
        The backend on which to evaluate the quantum circuit. The default can be
        specified in the file default_backend.py.
    compile : bool, optional
        Boolean indicating if the .compile method of the underlying QuantumSession
        should be called before. The default is ``True``.
    compilation_kwargs  : dict, optional
        Keyword arguments for the compile method. For more details check
        :meth:`QuantumSession.compile <qrisp.QuantumSession.compile>`. The default
        is ``{}``.
    subs_dic : dict, optional
        A dictionary of Sympy symbols and floats to specify parameters in the case
        of a circuit with unspecified, :ref:`abstract parameters<QuantumCircuit>`.
        The default is ``{}``.
    precompiled_qc : QuantumCircuit, optional
            A precompiled quantum circuit.
    diagonalisation_method : str, optional
        Specifies the method for grouping and diagonalizing the QubitOperator.
        Available are ``commuting_qw``, i.e., the operator is grouped based on qubit-wise commutativity of terms,
        and ``commuting``, i.e., the operator is grouped based on commutativity of terms.
        The default is ``commuting_qw``.
    measurement_data : QubitOperatorMeasurement
        Cached data to accelerate the measurement procedure. Automatically generated by default.

    Raises
    ------
    Exception
        If the containing QuantumSession is in a quantum environment, it is not
        possible to execute measurements.

    Returns
    -------
    float
        The expected value of the Hamiltonian.

    """

    from qrisp import QuantumSession, merge

    if isinstance(qarg, QuantumVariable):
        if qarg.is_deleted():
            raise Exception("Tried to get measurement from deleted QuantumVariable")
        qs = qarg.qs

    elif isinstance(qarg, QuantumArray):
        for qv in qarg.flatten():
            if qv.is_deleted():
                raise Exception(
                    "Tried to measure QuantumArray containing deleted QuantumVariables"
                )
        qs = qarg.qs
    elif isinstance(qarg, list):
        qs = QuantumSession()
        for arg in qarg:
            if isinstance(arg, QuantumVariable) and qv.is_deleted():
                raise Exception(
                    "Tried to measure QuantumArray containing deleted QuantumVariables"
                )
            merge(qs, arg)

    if backend is None:
        if qs.backend is None:
            from qrisp.default_backend import def_backend

            backend = def_backend
        else:
            backend = qarg.qs.backend

    if len(qs.env_stack) != 0:
        raise Exception("Tried to get measurement within open environment")

    hamiltonian = hamiltonian.hermitize()
    hamiltonian = hamiltonian.eliminate_ladder_conjugates()
    hamiltonian = hamiltonian.apply_threshold(0)
    if len(hamiltonian.terms_dict) == 0:
        return 0

    # Copy circuit in over to prevent modification
    if precompiled_qc is None:
        if compile:
            qc = qompiler(qs, **compilation_kwargs)
        else:
            qc = qs.copy()
        qubit_list = qarg
    else:
        qc = precompiled_qc.copy()
        qubit_list = qc.qubits[: len(qarg.reg)]

    # Bind parameters
    if subs_dic:
        qc = qc.bind_parameters(subs_dic)
        from qrisp.core.compilation import combine_single_qubit_gates

        qc = combine_single_qubit_gates(qc)

    qc = qc.transpile()

    if measurement_data is None:
        measurement_data = QubitOperatorMeasurement(
            hamiltonian, diagonalisation_method=diagonalisation_method
        )

    return measurement_data.get_measurement(qc, qubit_list, precision, backend)


class QubitOperatorMeasurement:

    def __init__(self, hamiltonian, diagonalisation_method="commuting_qw"):

        n = hamiltonian.find_minimal_qubit_amount()

        if diagonalisation_method == "commuting_qw":
            temp_groups = hamiltonian.commuting_qw_groups()
            self.groups = []
            # In order for the change of basis function (below) to work properly,
            # the ladder terms either need to completely agree or completely disagree
            for group in temp_groups:
                self.groups.extend(
                    group.group_up(
                        lambda a, b: a.ladders_agree(b) or not a.ladders_intersect(b)
                    )
                )

        elif diagonalisation_method == "commuting":
            temp_groups = hamiltonian.group_up(lambda a, b: a.commute_pauli(b))
            self.groups = []
            # In order for the change of basis function (below) to work properly,
            # the ladder terms either need to completely agree or completely disagree
            for group in temp_groups:
                self.groups.extend(
                    group.group_up(
                        lambda a, b: a.ladders_agree(b) or not a.ladders_intersect(b)
                    )
                )

        else:
            raise Exception(
                f"Unknown diagonalisation method: {diagonalisation_method}."
            )

        self.stds = []
        self.change_of_basis_gates = []
        self.measurement_operators = []

        for group in self.groups:

            qv = QuantumVariable(n)

            meas_op = group.change_of_basis(qv, diagonalisation_method)
            self.change_of_basis_gates.append(qv.qs.to_gate())
            self.measurement_operators.append(meas_op)

            # Collect standard deviation
            self.stds.append(np.sqrt(meas_op.get_operator_variance(n=n)))

        N = sum(self.stds)
        self.shots_list = [N * s for s in self.stds]

    def get_measurement(self, qc, qubit_list, precision, backend):

        from qrisp.misc import get_measurement_from_qc

        results = [0]*len(self.measurement_operators)
        meas_coeffs = list(results)
        meas_ops = list(results)
        
        def measurement_thread(i):
            
            group = self.measurement_operators[i]

            shots = int(self.shots_list[i] / precision**2)

            qubits = [
                qubit_list[j] for j in range(self.change_of_basis_gates[i].num_qubits)
            ]

            curr = qc.copy()
            curr.append(self.change_of_basis_gates[i], qubits)

            res = get_measurement_from_qc(curr, list(qubit_list), backend, shots)
            results[i] = res

            temp_meas_ops = []
            temp_coeff = []
            for term, coeff in group.terms_dict.items():
                temp_meas_ops.append(term.serialize())
                temp_coeff.append(coeff)

            meas_coeffs[i] = temp_coeff
            meas_ops[i] = temp_meas_ops
        
        
        threads = []
        
        for i in range(len(self.measurement_operators)):
            
            if isinstance(backend, BatchedBackend):
                thread = threading.Thread(target = measurement_thread, args = (i,))
                thread.start()
                threads.append(thread)
            else:
                measurement_thread(i)
        
        if isinstance(backend, BatchedBackend):
            backend.dispatch(min_calls = len(self.measurement_operators))
            for thread in threads:
                thread.join()

        samples = create_padded_array([list(res.keys()) for res in results]).astype(
            np.int64
        )
        probs = create_padded_array([list(res.values()) for res in results])
        meas_ops = create_padded_array(meas_ops, use_tuples=True).astype(np.int64)
        meas_coeffs = create_padded_array(meas_coeffs)

        return evaluate_expectation_jitted(samples, probs, meas_ops, meas_coeffs)


def create_padded_array(list_of_lists, use_tuples=False):
    """
    Create a padded numpy array from a list of lists with varying lengths.

    Parameters:
    list_of_lists (list): A list of lists with potentially different lengths.

    Returns:
    numpy.ndarray: A 2D numpy array with padded rows.
    """
    # Find the maximum length of any list in the input
    max_length = max(len(lst) for lst in list_of_lists)

    # Create a padded list of lists
    if not use_tuples:
        padded_lists = [lst + [0] * (max_length - len(lst)) for lst in list_of_lists]
    else:
        padded_lists = [
            lst + [(0, 0, 0, 0)] * (max_length - len(lst)) for lst in list_of_lists
        ]

    # Convert to numpy array
    return np.array(padded_lists)


#
# Evaluate expectation
#


def evaluate_expectation(samples, probs, operators, coefficients):
    """
    Evaluate the expectation.

    """
    # print(results)
    # print(operators)
    # print(coefficients)
    # raise

    expectation = 0

    for index1, ops in enumerate(operators):
        for index2, op in enumerate(ops):
            for i in range(len(samples[index1])):
                outcome, probability = samples[index1, i], probs[index1, i]
                expectation += (
                    probability
                    * evaluate_observable(op, outcome)
                    * np.real(coefficients[index1][index2])
                )

    return expectation


def evaluate_observable(observable: tuple, x: int):
    # This function evaluates how to compute the energy of a measurement sample x.
    # Since we are also considering ladder operators, this energy can either be
    # 0, -1 or 1. For more details check out the comments of QubitOperator.get_conjugation_circuit

    # The observable is given as tuple, containing for integers and a boolean.
    # To understand the meaning of these integers check QubitTerm.serialize.

    # Unwrap the tuple
    z_int, AND_bits, AND_ctrl_state, contains_ladder = observable

    # Compute whether the sign should be sign flipped based on the Z operators
    sign_flip_int = z_int & x
    sign_flip = 0
    while sign_flip_int:
        sign_flip += sign_flip_int & 1
        sign_flip_int >>= 1

    # If there is a ladder operator in the term, we need to half the energy
    # because we want to measure (|110><110| - |111><111|)/2
    if contains_ladder:
        prefactor = 0.5
    else:
        prefactor = 1

    # If there are no and bits, we return the result
    if AND_bits == 0:
        return prefactor * (-1) ** sign_flip

    # Otherwise we apply the AND_ctrl_state to flip the appropriate bits.
    corrected_x = x ^ AND_ctrl_state

    # If all bits are in the 0 state the AND is true.
    if corrected_x & AND_bits == 0:
        return prefactor * (-1) ** sign_flip
    else:
        return 0


evaluate_observable_jitted = njit(cache=True)(evaluate_observable)


@njit(cache=True)
def evaluate_expectation_jitted(samples, probs, operators, coefficients):
    """
    Evaluate the expectation.

    """
    expectation = 0

    for index1, ops in enumerate(operators):
        for index2, op in enumerate(ops):
            for i in range(len(samples[index1])):
                outcome, probability = samples[index1, i], probs[index1, i]
                expectation += (
                    probability
                    * evaluate_observable_jitted(op, outcome)
                    * np.real(coefficients[index1][index2])
                )

    return expectation


def partition(values, num_qubits):
    """
    Partitions a list of integers into a list of lists of integers with size 64 bit.

    Parameters
    ----------
    values : list[int]
        A list of integers.
    num_qubits : int
        The maximal number of bits.

    Returns
    -------
    partition : list[np.array]
        A list of NumPy numpy.uint64 arrays.

    """

    M = math.ceil(num_qubits / 64)
    N = len(values)

    zeros = [[0] * N for k in range(M)]
    partition = [values]
    partition.extend(zeros)

    lower_mask = (1 << 64) - 1

    for j in range(1, M):
        for i in range(N):
            partition[j][i] = partition[j - 1][i] & lower_mask
            partition[j + 1][i] = partition[j - 1][i] >> 64

    if M == 1:
        return [np.array(partition[0], dtype=np.uint64)]
    else:
        return [np.array(part, dtype=np.uint64) for part in partition[1:]]
