import os
import time
from typing import Dict, Optional

import numpy as np
from tqdm import tqdm
import qiskit
from qiskit import Aer, IBMQ
from qiskit.providers.ibmq.job.ibmqjob import JobStatus

from qrem.functions_qrem.povmtools import reorder_probabilities
from qrem.functions_qrem import ancillary_functions as anf
from qrem.functions_qrem import povmtools

from qrem.common import io
from qrem.common.printer import qprint

#ORGANIZATION - poorly documented code (MO)

def get_qiskit_provider(backend_name: str,
                        provider_data: Dict[str, str]=None):
    IBMQ.load_account()
    if backend_name.lower() in ['ionq.qpu', 'ionq.simulator']:
        from azure.quantum.qiskit import AzureQuantumProvider

        if provider_data is None:
            AZURE_RESOURCE_ID = os.environ['AZURE_RESOURCE_ID_IONQ']
            AZURE_LOCATION = os.environ['AZURE_LOCATION_IONQ']
        else:
            AZURE_RESOURCE_ID = provider_data['AZURE_RESOURCE_ID_IONQ']
            AZURE_LOCATION = provider_data['AZURE_LOCATION_IONQ']

        provider = AzureQuantumProvider(
            resource_id=AZURE_RESOURCE_ID,
            location=AZURE_LOCATION)

    else:
        if provider_data is None:
            IBMQ_hub = os.environ['IBMQ_HUB']
            IBMQ_group = os.environ['IBMQ_GROUP']
            IBMQ_project = os.environ['IBMQ_PROJECT']
        else:
            IBMQ_hub = provider_data['IBMQ_HUB']
            IBMQ_group = provider_data['IBMQ_GROUP']
            IBMQ_project = provider_data['IBMQ_PROJECT']

        provider = IBMQ.get_provider(hub=IBMQ_hub,
                                     group=IBMQ_group,
                                     project=IBMQ_project)

    return provider

def add_gate_to_circuit(circuit,
                        qreg,
                        q_index,
                        unitary):
    # TODO: Check whether this is needed - at some point I remember there were troubles with parametrizing those two unitaries
    if unitary[0, 0] == 1 and unitary[1, 1] == 1:
        pass
    elif unitary[0, 1] == 1 and unitary[1, 0] == 1:
        circuit.X(qreg[q_index])

    else:
        # get angles for single-qubit state change unitary
        current_angles = povmtools.get_su2_parametrizing_angles(unitary)

        # implement unitary
        circuit.u3(current_angles[0],
                   current_angles[1],
                   current_angles[2],
                   qreg[q_index])
    return circuit, qreg


#ORGANIZE: copied to experiment_results. Delete when refactoring complete
def get_frequencies_from_counts(counts_dict,
                                crs=None,
                                classical_register=None,
                                shots_number=None,
                                reorder_bits=True):
    if crs is None:
        crs = len(list(list(counts_dict.keys())[0]))

    d = 2 ** crs

    if classical_register is None:
        classical_register = ["{0:b}".format(i).zfill(crs) for i in range(d)]

    normal_order = []

    for j in range(d):
        if classical_register[j] in counts_dict.keys():
            counts_now = counts_dict[classical_register[j]]
            normal_order.append(counts_now)

        else:
            normal_order.append(0)
    if reorder_bits:
        frequencies = reorder_probabilities(normal_order, range(crs)[::-1])
    else:
        frequencies = normal_order

    if shots_number is None:
        frequencies = frequencies / np.sum(frequencies)
    else:
        frequencies = frequencies / shots_number

    return frequencies


def download_multiple_jobs(backend_name,
                           job_IDs_list,
                           provider_data=None):
    IBMQ.load_account()
    if backend_name in ['qasm_simulator', 'statevector_simulator', 'unitary_simulator']:
        raise ValueError('Local simulators do not store jobs online.')

    provider = get_qiskit_provider(backend_name=backend_name,
                                   provider_data=provider_data)

    backend = provider.get_backend(backend_name)

    all_jobs = []
    for job_ID in tqdm(job_IDs_list):
        # qprint('Getting job with ID:', job_ID)
        job = backend.retrieve_job(job_ID)
        # qprint('Got it!')
        all_jobs.append(job)
    return all_jobs


#ORGANIZE: copied to experiment_results, but possibly not needed.
def add_counts_dicts(all_counts, modulo, dimension):
    frequencies = [np.zeros(dimension) for i in range(modulo)]

    for counts_index in tqdm(range(len(all_counts))):
        true_index = counts_index % modulo

        freqs_now = povmtools.counts_dict_to_frequencies_vector(all_counts[counts_index], True)
        frequencies[true_index][:] += freqs_now[:]

        # print(freqs_now)
    for i in range(modulo):
        frequencies[i] *= 1 / np.sum(frequencies[i])

    return frequencies

#MOcomm is it just for qiskit or more general

def run_batches(batches,
                backend_name,
                shots,
                saving_IDs_dictionary=dict(saving=False, directory=None, file_name=None, dictionary_to_save={}),
                provider_data:Optional[Dict[str,str]] = None, verbose_log = True):
    """
    Description:
        This function takes a list of circuit batches and sends them for execution on the specified backend.
        User can choose to save the resulting jobs' IDs, using parameters in saving_IDs_dictionary
    Parameters:
        :param batches: (list(list(qiskit.circuit.quantumcircuit.QuantumCircuit))): list of batches, where each batch is
        a list of qiskit circuits and each batch has length <= maximum number of circuits in job
        :param backend_name: (str): if value is of the names of a qiskit simulator: ('qasm_simulator',
        'statevector_simulator', 'unitary_simulator'), provider is chosen to be Aer; in other cases provider is taken
         from 'provider_data'
        #TODO: consider cases where other backend names are provided (e.g. aer_simulator)
        :param shots: number of repetitions of each circuit from batches
        :param saving_IDs_dictionary:
        :provider_data: Dict[str]: contains access data required by providers, like user ID etc. If None, as is default,
        this data is obtained from environmental variables by function get_qiskit_provider (see function for details)
    Returns:
        jobs: list of qiskit jobs (results of qiskit.execute())
    Notes:

    """
    saving = saving_IDs_dictionary['saving']

    qprint('\nSending jobs to execution on: ', backend_name + '.')
    qprint('Number of shots: ', str(shots) + ' .')
    qprint('Target number of jobs: ', str(len(batches)) + ' .')
    print()

    iterations_done = 0
    wait_time_in_minutes = 10
    jobs = []

    qiskit_simulators = ['qasm_simulator', 'statevector_simulator', 'unitary_simulator']
    braket_simulators = ['BraketLocalBackend']
    braket_quantum_devices = ['Aspen-M-1', 'Aspen-M-2']

    # set provider and backend according to passed argument 'backend_name':
    if backend_name in qiskit_simulators:
        backend = Aer.get_backend(backend_name)
    else:
        provider = get_qiskit_provider(backend_name=backend_name,
                                       provider_data=provider_data)
        backend = provider.get_backend(backend_name)

    # Run the circuits on backend:
    while iterations_done < len(batches):
        qprint('job number:', str(iterations_done))
        circuits = batches[iterations_done]
        try:
            time.sleep(2)
            qprint("Sending quantum program to: ", backend_name + '.')
            # below we use pure qiskit:
            if backend_name not in braket_simulators and backend_name not in braket_quantum_devices:
                job = qiskit.execute(circuits, backend, shots=shots, max_credits=200)#, qobj_id=qobj_id)

            # below we use braket (note that we use 'job' here in the qiskit sense, in braket this is called 'task'):
            else:
                job = backend.run(circuits, shots=shots, disable_qubit_rewiring=True)
                # save the circuit names as metadata to use in get_counts_from_jobs and get_counts_from_result_object:
                job.metadata['circuit_names'] = [circuit.name for circuit in circuits]
                # alternatives: job._tasks[idx]._result.task_metadata.braketSchemaHeader.name; Tag property of task
            jobs.append(job)

            if saving and backend_name not in ['qasm_simulator', 'statevector_simulator', 'unitary_simulator']:
                job_ID = job.job_id()
                dict_to_save = saving_IDs_dictionary['dictionary_to_save']
                dict_to_save['job_ID'] = job_ID

                io.save(dictionary_to_save=dict_to_save,
                         directory=saving_IDs_dictionary['directory'],
                         custom_name=saving_IDs_dictionary['file_name'] + '_job%s' % iterations_done,
                         overwrie = False)
            while job.status() == JobStatus.INITIALIZING:
                print(job.status())
                time.sleep(2)
            qprint("Program sent for execution to: ", backend_name + '.')

        except BaseException as ex:
            print('There was an error in the circuit!. Error = {}'.format(ex))
            print(f'Waiting {wait_time_in_minutes} minute(s) before next try.')
            time.sleep(wait_time_in_minutes * 60)
            continue

        print()
        iterations_done += 1
    return jobs
