import json

import cirq
import numpy as np
from pytket import Circuit
from pytket.circuit import Unitary1qBox, Unitary2qBox, Unitary3qBox
from pytket.extensions.qiskit import AerBackend
from qiskit import QuantumCircuit, transpile
from qiskit.circuit.library import UnitaryGate
from qiskit.quantum_info import Statevector
from qiskit_aer import AerSimulator
from qiskit_ibm_runtime import QiskitRuntimeService
from qiskit_ibm_runtime import SamplerV2 as Sampler


class State:
    def __init__(self, n_qubits: int = 1):
        self.n_qubits: int = n_qubits
        self.state: np.ndarray = np.zeros(2**n_qubits, dtype=complex)
        self.H: np.ndarray = 1 / np.sqrt(2) * np.array([[1, 1], [1, -1]])
        self.X: np.ndarray = np.array([[0, 1], [1, 0]])
        self.S: np.ndarray = np.array([[1, 0], [0, 1j]])
        self.CNOT: np.ndarray = np.array(
            [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], dtype=complex
        )
        self.CH: np.ndarray = np.array(
            [
                [1, 0, 0, 0],
                [0, 1, 0, 0],
                [0, 0, 1 / np.sqrt(2), 1 / np.sqrt(2)],
                [0, 0, 1 / np.sqrt(2), -1 / np.sqrt(2)],
            ],
            dtype=complex,
        )
        self.SWAP: np.ndarray = np.array(
            [[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]], dtype=complex
        )
        self.TOFFOLI: np.ndarray = np.array(
            [
                [1, 0, 0, 0, 0, 0, 0, 0],
                [0, 1, 0, 0, 0, 0, 0, 0],
                [0, 0, 1, 0, 0, 0, 0, 0],
                [0, 0, 0, 1, 0, 0, 0, 0],
                [0, 0, 0, 0, 1, 0, 0, 0],
                [0, 0, 0, 0, 0, 1, 0, 0],
                [0, 0, 0, 0, 0, 0, 0, 1],
                [0, 0, 0, 0, 0, 0, 1, 0],
            ],
            dtype=complex,
        )
        self.__custom_gates: dict[str, np.ndarray] = {}
        self.__circuit: list[dict[str, list[int] | str]] = []
        self.__initial_qubits: list[int] = []
        self.__initial_state: np.ndarray = np.zeros(2**n_qubits, dtype=complex)

    def __repr__(self):
        return f"State(n_qubits={self.n_qubits}, state={self.state})"

    def __str__(self):
        def clear_space(proto: list[str], n_qubits: int):
            for i in range(2 * n_qubits + 1):
                if i % 2 == 1:
                    proto[i] += "-------"
                else:
                    proto[i] += "       "

            return proto

        def draw_gate(proto: list[str], gate: list[dict[str, list[int] | str]]):
            list_proto = [list(proto[i]) for i in range(2 * self.n_qubits + 1)]

            for qubit in gate["qubits"]:
                list_proto[2 * qubit + 1][-1] = "|"
                list_proto[2 * qubit + 1][-5] = "|"
                for i in range(-4, -1, 1):
                    list_proto[2 * qubit][i] = "-"
                    list_proto[2 * qubit + 2][i] = "-"

                if len(gate["name"]) == 1:
                    list_proto[2 * qubit + 1][-3] = gate["name"]
                elif len(gate["name"]) == 2:
                    list_proto[2 * qubit + 1][-4] = gate["name"][0]
                    list_proto[2 * qubit + 1][-3] = gate["name"][1]
                else:
                    list_proto[2 * qubit + 1][-4] = gate["name"][0]
                    list_proto[2 * qubit + 1][-3] = gate["name"][1]
                    list_proto[2 * qubit + 1][-2] = gate["name"][2]

            for i in range(2 * self.n_qubits + 1):
                proto[i] = "".join(list_proto[i])

            return proto

        proto = ["" for _ in range(2 * self.n_qubits + 1)]

        for i in range(2 * self.n_qubits + 1):
            if i % 2 == 1:
                proto[i] = str(self.__initial_qubits[i // 2]) + "--"
            else:
                proto[i] = "   "

        for gate in self.__circuit:
            proto = clear_space(proto, self.n_qubits)

            proto = draw_gate(proto, gate)

        return "\n".join(proto)

    def initialize_state(self, qubit_values: list[int]):
        assert len(qubit_values) == self.n_qubits
        index = int("".join(map(str, qubit_values)), 2)
        self.state[index] = 1
        self.__initial_state = self.state.copy()
        self.__initial_qubits = qubit_values

    def __apply_gate(self, gate: np.ndarray, n_gate: int, starting_qubit: int):
        U = np.array([[1]], dtype=complex)

        for qubit in range(self.n_qubits):
            if starting_qubit <= qubit < starting_qubit + n_gate:
                if qubit == starting_qubit:
                    U = np.kron(U, gate)
            else:
                U = np.kron(U, np.eye(2))
        self.state = U @ self.state

    def apply_H(self, target_qubit: int):
        self.__circuit.append({"name": "H", "qubits": [target_qubit]})
        self.__apply_gate(self.H, 1, target_qubit)

    def apply_X(self, target_qubit: int):
        self.__circuit.append({"name": "X", "qubits": [target_qubit]})
        self.__apply_gate(self.X, 1, target_qubit)

    def apply_S(self, target_qubit: int):
        self.__circuit.append({"name": "S", "qubits": [target_qubit]})
        self.__apply_gate(self.S, 1, target_qubit)

    def make_adjacency(self, qubit1: int, qubit2: int):
        for q in range(qubit1, qubit2 - 1):
            self.__apply_SWAP_adjacent(q)

    def destroy_adjacency(self, qubit1: int, qubit2: int):
        for q in range(qubit2 - 2, qubit1 - 1, -1):
            self.__apply_SWAP_adjacent(q)

    def non_adjacent_helper(self, gate: np.ndarray, control: int, target: int):
        if control == target:
            return

        q1, q2, reverse = (
            (control, target, 0) if control < target else (target, control, 1)
        )

        self.make_adjacency(q1, q2)

        if reverse == 0:
            self.__apply_gate(gate, 2, q2 - 1)
        else:
            self.__apply_SWAP_adjacent(q2 - 1)
            self.__apply_gate(gate, 2, q2 - 1)
            self.__apply_SWAP_adjacent(q2 - 1)

        self.destroy_adjacency(q1, q2)

    def apply_CNOT(self, control: int, target: int):
        self.__circuit.append({"name": "CNOT", "qubits": [control, target]})
        self.non_adjacent_helper(self.CNOT, control, target)

    def apply_CH(self, control: int, target: int):
        self.__circuit.append({"name": "CH", "qubits": [control, target]})
        self.non_adjacent_helper(self.CH, control, target)

    def __apply_SWAP_adjacent(self, target_qubit: int):
        self.__apply_gate(self.SWAP, 2, target_qubit)

    def __apply_SWAP(self, q1: int, q2: int):
        self.non_adjacent_helper(self.SWAP, q1, q2)

    def apply_SWAP(self, q1: int, q2: int):
        self.__circuit.append({"name": "SWAP", "qubits": [q1, q2]})
        self.non_adjacent_helper(self.SWAP, q1, q2)

    def __make_linear(self, qubits: list[int]):
        pos = list(range(self.n_qubits))
        swap_sequence: list[tuple[int, int]] = []

        for target_pos, qubit in enumerate(qubits):
            current_index = pos.index(qubit)
            while current_index != target_pos:
                if current_index > target_pos:
                    self.__apply_SWAP(current_index - 1, current_index)
                    swap_sequence.append((current_index - 1, current_index))
                    pos[current_index], pos[current_index - 1] = (
                        pos[current_index - 1],
                        pos[current_index],
                    )
                    current_index -= 1
                else:
                    self.__apply_SWAP(current_index, current_index + 1)
                    swap_sequence.append((current_index, current_index + 1))
                    pos[current_index], pos[current_index + 1] = (
                        pos[current_index + 1],
                        pos[current_index],
                    )
                    current_index += 1
        return swap_sequence

    def __make_non_linear(self, swap_sequence: list[tuple[int, int]]):
        for a, b in reversed(swap_sequence):
            self.__apply_SWAP(a, b)

    def apply_TOFFOLI(self, control1: int, control2: int, target: int):
        qubits = [control1, control2, target]
        self.__circuit.append({"name": "TOFFOLI", "qubits": qubits})
        assert len({control1, control2, target}) == 3

        swap_sequence = self.__make_linear(qubits)

        self.__apply_gate(self.TOFFOLI, 3, 0)

        self.__make_non_linear(swap_sequence)

    def create_custom_gate(self, name: str, gate_matrix: np.ndarray):
        if gate_matrix.shape[0] != gate_matrix.shape[1]:
            raise ValueError("Gate must be a square matrix!")

        identity = np.eye(gate_matrix.shape[0])
        if not np.allclose(gate_matrix.conj().T @ gate_matrix, identity):
            raise ValueError("Gate must be an unitary matrix!")
        self.__custom_gates[name.upper()] = gate_matrix
        return gate_matrix

    def create_custom_gate_from_other_gates(self, name: str, *gates):
        gates = list(gates)
        if len(gates) == 0:
            raise ValueError("Empty list of gates!")
        U = np.eye(gates[0].shape[0])

        for g in gates:
            if g.shape[0] != U.shape[0]:
                raise ValueError("Gates must take the same number of qubits each!")
            if not isinstance(g, np.ndarray):
                raise ValueError("Not a valid gate!")
            U = g @ U
        self.__custom_gates[name.upper()] = U
        return U

    def apply_custom_gate(self, name: str, *qubits):
        name = name.upper()
        if self.__custom_gates.get(name) is None:
            raise ValueError("Gate not defined!")
        gate = self.__custom_gates[name]
        qubits = list(qubits)
        k = int(np.log2(gate.shape[0]))
        if len(qubits) != k:
            raise ValueError("Invalid number of qubits")
        swap_sequence = self.__make_linear(qubits)

        self.__apply_gate(gate, k, 0)

        self.__make_non_linear(swap_sequence)
        self.__circuit.append({"name": name, "qubits": qubits})

    def make_controlled(self, name: str, gate: np.ndarray):
        if isinstance(gate, str):
            gate = self.__custom_gates[gate]

        dim_target = gate.shape[0]
        dim_total = 2 * dim_target

        CU = np.zeros((dim_total, dim_total), dtype=complex)

        CU[:dim_target, :dim_target] = np.eye(dim_target)

        CU[dim_target:, dim_target:] = gate
        self.__custom_gates[name.upper()] = CU
        return CU

    def produce_measurement(self):
        n = self.n_qubits
        probs = np.abs(self.state) ** 2
        probs /= np.sum(probs)
        index = np.random.choice(2**n, p=probs)
        measured_state = [int(x) for x in format(index, f"0{n}b")]
        return measured_state

    # different from normal phase function; used only in QFT
    def __phase(self, m):
        return np.array([[1, 0], [0, np.exp((2 * np.pi * 1j) / (2**m))]], dtype=complex)

    def apply_QFT(self, qubits):
        swap_sequence = self.__make_linear(qubits)

        n = len(qubits)

        for i in range(n):
            self.apply_H(i)
            self.__circuit[::-1][0]["qubits"] = [qubits[i]]
            for j in range(n - i - 1):
                self.make_controlled(f"R{i}{j + 2}", self.__phase(j + 2))
                self.apply_custom_gate(f"R{i}{j + 2}", i + j + 1, i)
                self.__circuit[::-1][0]["qubits"] = [qubits[i + j + 1], qubits[i]]

        for i in range(n // 2):
            self.apply_SWAP(i, n - i - 1)
            self.__circuit[::-1][0]["qubits"] = [qubits[i], qubits[n - i - 1]]
        self.__make_non_linear(swap_sequence)

    def __complex_to_list(self, arr):
        arr = np.array(arr)
        if arr.ndim == 1:
            return [[c.real, c.imag] for c in arr]
        elif arr.ndim == 2:
            return [[[c.real, c.imag] for c in row] for row in arr]

    def __list_to_complex(self, arr):
        arr = np.array(arr)
        if arr.ndim == 2:
            return np.array([complex(c[0], c[1]) for c in arr], dtype=complex)
        elif arr.ndim == 3:
            return np.array(
                [[complex(c[0], c[1]) for c in row] for row in arr], dtype=complex
            )

    def export_circuit(self, filename):
        data = {
            "n_qubits": self.n_qubits,
            "initial_state": self.__complex_to_list(self.__initial_state.tolist()),
            "circuit": self.__circuit,
            "custom_gates": {
                name: self.__complex_to_list(g)
                for name, g in self.__custom_gates.items()
            },
        }
        with open(filename, "w") as f:
            json.dump(data, f)

    def bra_state(self):
        if self.state is None:
            raise ValueError("No state vector available.")

        return self.state.conj()

    def ket_state(self):
        if self.state is None:
            raise ValueError("No state vector available.")

        return self.state.T

    def compute_density_matrix(self):
        return np.outer(self.ket_state(), self.bra_state())

    def compute_fidelity(self, state_two):
        return np.abs(np.vdot(self.state, state_two.state)) ** 2

    def __run(self, circ):
        for c in circ:
            match c["name"]:
                case "H":
                    self.apply_H(c["qubits"][0])
                case "X":
                    self.apply_X(c["qubits"][0])
                case "S":
                    self.apply_S(c["qubits"][0])
                case "CNOT":
                    self.apply_CNOT(c["qubits"][0], c["qubits"][1])
                case "TOFFOLI":
                    self.apply_TOFFOLI(c["qubits"][0], c["qubits"][1], c["qubits"][2])
                case "CH":
                    self.apply_CH(c["qubits"][0], c["qubits"][1])
                case "SWAP":
                    self.apply_SWAP(c["qubits"][0], c["qubits"][1])
                case _:
                    if self.__custom_gates.get(c["name"].upper()) is not None:
                        self.apply_custom_gate(c["name"], *c["qubits"])
                    else:
                        raise ValueError("Gate not found!")

    def import_circuit(self, filename):
        with open(filename, "r") as f:
            data = json.load(f)
        self.__custom_gates = {
            name: self.__list_to_complex(data["custom_gates"][name])
            for name in data["custom_gates"]
        }

        circ = data["circuit"]
        self.state = self.__list_to_complex(data["initial_state"])
        self.__initial_state = self.state.copy()
        self.n_qubits = data["n_qubits"]

        self.__run(circ)

    def get_probabilities(self):
        return abs(self.state) ** 2

    def run(self, shots=1024):
        prob = self.get_probabilities()
        prob /= np.sum(prob)
        index = np.random.choice(len(prob), size=shots, p=prob)
        unique_index, counts = np.unique(index, return_counts=True)
        cs = {}
        for idx, count in zip(unique_index, counts):
            bin_string = format(idx, f"0{self.n_qubits}b")
            cs[bin_string] = int(count)

        return cs


class QiskitExport:
    def __init__(self, og):
        self.og = og
        self.og_circuit = og._State__circuit
        self.og_custom_gates = og._State__custom_gates
        self.circuit = self.export()
        self.state = Statevector.from_instruction(self.circuit)
        self.probabilities = np.abs(self.state) ** 2

    def export(self):
        n = self.og.n_qubits
        qc = QuantumCircuit(n)

        qc.initialize(self.og._State__initial_state, range(n))

        for op in self.og_circuit:
            qubits = [n - 1 - q for q in op["qubits"]]

            match op["name"]:
                case "H":
                    qc.h(qubits[0])
                case "X":
                    qc.x(qubits[0])
                case "S":
                    qc.s(qubits[0])
                case "CNOT":
                    qc.cx(qubits[0], qubits[1])
                case "SWAP":
                    qc.swap(qubits[0], qubits[1])
                case "TOFFOLI":
                    qc.ccx(qubits[0], qubits[1], qubits[2])
                case "CH":
                    qc.ch(qubits[0], qubits[1])
                case _:
                    if op["name"] in self.og_custom_gates:
                        matrix = self.og_custom_gates[op["name"]]

                        gate = UnitaryGate(matrix, label=op["name"])
                        qc.append(gate, qubits[::-1])
                    else:
                        raise ValueError("Gate not found!")

        return qc

    def run(self, shots=1024):
        qc_run = self.circuit.copy()
        qc_run.measure_all()
        sim = AerSimulator()
        t_qc = transpile(qc_run, sim)
        result = sim.run(t_qc, shots=shots).result()
        return result.get_counts()


class CirqExport:
    def __init__(self, og):
        self.og = og
        self.og_circuit = og._State__circuit
        self.og_custom_gates = og._State__custom_gates
        self.qubits = cirq.LineQubit.range(og.n_qubits)
        self.circuit = self.export()

    def export(self):
        circuit = cirq.Circuit()

        index = np.argmax(self.og._State__initial_state)
        bin = format(index, f"0{self.og.n_qubits}b")

        for i, bit in enumerate(bin):
            if bit == "1":
                circuit.append(cirq.X(self.qubits[i]))

        for op in self.og_circuit:
            op_name = op["name"]

            q_args = [self.qubits[i] for i in op["qubits"]]

            match op_name:
                case "H":
                    circuit.append(cirq.H(*q_args))
                case "X":
                    circuit.append(cirq.X(*q_args))
                case "S":
                    circuit.append(cirq.S(*q_args))
                case "CNOT":
                    circuit.append(cirq.CNOT(q_args[0], q_args[1]))
                case "SWAP":
                    circuit.append(cirq.SWAP(q_args[0], q_args[1]))
                case "TOFFOLI":
                    circuit.append(cirq.TOFFOLI(q_args[0], q_args[1], q_args[2]))
                case "CH":
                    circuit.append(cirq.H(q_args[1]).controlled_by(q_args[0]))
                case _:
                    if op_name in self.og_custom_gates:
                        matrix = self.og_custom_gates[op_name]
                        matrix = np.array(matrix, dtype=complex)
                        custom_gate = cirq.MatrixGate(matrix)

                        circuit.append(custom_gate.on(*q_args))
                    else:
                        raise ValueError(f"Gate {op_name} not found or supported.")

        return circuit

    def run(self, shots=1024):
        qc_run = self.circuit.copy()

        qc_run.append(cirq.measure(*self.qubits, key="m"))

        sim = cirq.Simulator()
        result = sim.run(qc_run, repetitions=shots)

        measurements = result.measurements["m"]

        counts = {}
        for row in measurements:
            bin_string = "".join(str(bit) for bit in row)

            counts[bin_string] = counts.get(bin_string, 0) + 1

        return counts


class TketExport:
    def __init__(self, og):
        self.og = og
        self.og_circuit = og._State__circuit
        self.og_custom_gates = og._State__custom_gates
        self.circuit = self.export()

    def __add_unitary_box(self, matrix, qubits, qc):
        rows = matrix.shape[0]
        n_qubits = int(np.log2(rows))

        if n_qubits == 1:
            box = Unitary1qBox(matrix)

        elif n_qubits == 2:
            box = Unitary2qBox(matrix)

        elif n_qubits == 3:
            box = Unitary3qBox(matrix)

        else:
            raise ValueError("Too many qubits for a single TKET gate (max 3)")

        qc.add_gate(box, qubits)

    def export(self):
        n = self.og.n_qubits
        qc = Circuit(n)

        index = np.argmax(np.abs(self.og._State__initial_state))
        bin_str = format(index, f"0{n}b")

        for i, bit in enumerate(bin_str):
            if bit == "1":
                qc.X(n - 1 - i)

        for op in self.og_circuit:
            op_name = op["name"]
            qubits = [n - 1 - q for q in op["qubits"]]

            match op_name:
                case "H":
                    qc.H(qubits[0])
                case "X":
                    qc.X(qubits[0])
                case "S":
                    qc.S(qubits[0])
                case "CNOT":
                    qc.CX(qubits[0], qubits[1])
                case "SWAP":
                    qc.SWAP(qubits[0], qubits[1])
                case "TOFFOLI":
                    qc.CCX(qubits[0], qubits[1], qubits[2])
                case "CH":
                    qc.CH(qubits[0], qubits[1])
                case _:
                    if op_name in self.og_custom_gates:
                        matrix = self.og_custom_gates[op_name]
                        self.__add_unitary_box(matrix, qubits, qc)
                    else:
                        raise ValueError("Gate not found.")

        return qc

    def run(self, shots=1024):
        backend = AerBackend()
        qc_run = backend.get_compiled_circuit(self.circuit)

        qc_run.measure_all()

        handle = backend.process_circuit(qc_run, n_shots=shots)
        result = backend.get_result(handle)
        counts_aux = result.get_counts()

        counts = {}
        for outcome, c in counts_aux.items():
            bin_string = "".join(map(str, outcome))[::-1]
            counts[bin_string] = c

        return counts


class RealQC(QiskitExport):
    def __init__(self, og, token, instance=None):
        super().__init__(og)

        try:
            if instance:
                self.service = QiskitRuntimeService(
                    channel="ibm_cloud", token=token, instance=instance.strip()
                )
            else:
                self.service = QiskitRuntimeService(
                    channel="ibm_quantum_platform", token=token
                )
        except Exception as e:
            raise ValueError(f"IBM Auth Failed: {e}")

    def run(self, shots=1024, backend_name=None):
        if backend_name:
            backend = self.service.backend(backend_name)
        else:
            print("Searching for least busy qc...")
            backend = self.service.least_busy(operational=True, simulator=False)

        print(f"Running on: {backend.name}")

        qc_run = self.circuit.copy()
        qc_run.measure_all()

        print("Transpiling circuit...")
        t_qc = transpile(qc_run, backend)

        sampler = Sampler(mode=backend)
        job = sampler.run([t_qc], shots=shots)
        print(f"Job submitted! ID: {job.job_id()}")
        print("Waiting for results...")

        result = job.result()

        counts = result[0].data.meas.get_counts()
        return counts
