#  =========================================================================  #
#   Copyright 2018 National Technology & Engineering Solutions of Sandia,
#   LLC (NTESS). Under the terms of Contract DE-NA0003525 with NTESS,
#   the U.S. Government retains certain rights in this software.
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#  =========================================================================  #

import numpy as np
from .. import circuit_runners, circuits
from ..qeccs import Surface4444
from ..decoders import MWPM2D
from ..error_gens import XGen
from ..misc.threshold_curve import threshold_fit as default_fit, func as default_func


def threshold_code_capacity(qecc_class, error_gen, decoder_class, ps, ds, runs, verbose=False, mode=1,
                            threshold_fit=None, p0=None, func=None, circuit_runner=None, basis=None):
    """
    Function that generates p_logical values given a list of physical errors (ps) and distance (ds).
    Args:
        ps(list of float):
        ds(list of int):
        runs(int):
        error_gen:
        qecc_class:
        decoder_class:
        verbose:
        mode:
        threshold_fit:
        p0:
        func:
        circuit_runner:

    Returns:

    """

    if circuit_runner is None:
        circuit_runner = circuit_runners.Standard()

    if error_gen is None:
        error_gen = XGen(model_level='code_capacity')

    if qecc_class is None:
        qecc_class = Surface4444

    if decoder_class is None:
        decoder_class = MWPM2D

    if threshold_fit is None:
        threshold_fit = default_fit

    if func is None:
        func = default_func

        if p0 is None:
            p0 = (0.1, 1.5, 1, 1, 1)

    if basis not in [None, 'zero', 'plus', 'both']:
        raise Exception('`basis` can only be "None", "zero", "plus", "both"!')

    if mode == 1 and basis != 'both':
        determine_rate = codecapacity_logical_rate
    elif mode == 1 and basis == 'both':
        determine_rate = codecapacity_logical_rate2
    elif mode == 2:
        determine_rate = codecapacity_logical_rate3
    else:
        raise Exception('Mode "%s" is not handled!' % mode)

    plist = np.array(ps * len(ds))

    dlist = []
    for d in ds:
        for p in ps:
            dlist.append(d)
    dlist = np.array(dlist)

    plog = []
    for d in ds:

        qecc = qecc_class(distance=d)
        decoder = decoder_class(qecc)

        for p in ps:

            logical_error_rate, time = determine_rate(runs, qecc, d, error_gen, error_params={'p': p}, decoder=decoder,
                                                      verbose=verbose, circuit_runner=circuit_runner, basis=basis)
            if verbose:
                if time:
                    print('Runtime: %s s' % time)

                print('----')

            plog.append(logical_error_rate)

    plog = np.array(plog)

    # results = threshold_fit(plist, dlist, plog, func, p0)

    # return plist, dlist, plog, results
    # return {'distances': ds, 'ps': ps, 'plog': plog, 'opt': results[0], 'std': results[1]}
    return {'distances': ds, 'ps_physical': plist, 'p_logical': plog}


def threshold_code_capacity_calc(ps, ds, runs, error_gen=None, qecc_class=None, decoder_class=None, verbose=True,
                                 mode=1, threshold_fit=None, p0=None, func=None, circuit_runner=None):
    """
    Function that generates p_logical values given a list of physical errors (ps) and distance (ds).
    Args:
        ps(list of float):
        ds(list of int):
        runs(int):
        error_gen:
        qecc_class:
        decoder_class:
        verbose:
        mode:
        threshold_fit:
        p0:
        func:
        circuit_runner:

    Returns:

    """

    if circuit_runner is None:
        circuit_runner = circuit_runners.Standard()

    if error_gen is None:
        error_gen = XGen(model_level='code_capacity')

    if qecc_class is None:
        qecc_class = Surface4444

    if decoder_class is None:
        decoder_class = MWPM2D

    if threshold_fit is None:
        threshold_fit = default_fit

    if func is None:
        func = default_func

        if p0 is None:
            p0 = (0.1, 1.5, 1, 1, 1)

    if mode == 1:
        determine_rate = codecapacity_logical_rate
    elif mode == 2:
        determine_rate = codecapacity_logical_rate2
    elif mode == 3:
        determine_rate = codecapacity_logical_rate3
    else:
        raise Exception('Mode "%s" is not handled!' % mode)

    plist = np.array(ps * len(ds))

    dlist = []
    for d in ds:
        for p in ps:
            dlist.append(d)
    dlist = np.array(dlist)

    plog = []
    for d in ds:

        qecc = qecc_class(distance=d)
        decoder = decoder_class(qecc)

        for p in ps:

            logical_error_rate, time = determine_rate(runs, qecc, d, error_gen, error_params={'p': p}, decoder=decoder,
                                                      verbose=verbose, circuit_runner=circuit_runner)
            if verbose:
                if time:
                    print('Runtime: %s s' % time)

                print('----')

            plog.append(logical_error_rate)

    plog = np.array(plog)

    results = threshold_fit(plist, dlist, plog, func, p0)

    # return plist, dlist, plog, results
    return {'plist': plist, 'dlist': dlist, 'plog': plog, 'opt': results[0], 'std': results[1]}


def codecapacity_logical_rate(runs, qecc, distance, error_gen, error_params, decoder, seed=None, state_sim=None,
                              verbose=True, circuit_runner=None, basis=None):
    """
    A tool for determining the code-capacity logical-error rate for syndrome extraction.

    In this analysis only logical |0> is prepared and each run consists of an ideal logical |0> preparation followed by
    a single round of syndrome extraction. The error rate is determined by number of runs with logical failures divided
    by the total number of runs.

    Args:
        runs: Number of runs to evaluate the logical error rate.
        qecc:
        distance:
        error_gen:
        error_params:
        decoder:
        seed:
        state_sim:
        verbose:
        circuit_runner:

    Returns:

    """

    p = error_params['p']
    total_time = 0.0

    # Circuit simulator
    if circuit_runner is None:
        circuit_runner = circuit_runners.TimingRunner(seed=seed)
    # circ_sim = circuit_runners.Standard(seed=seed)

    # Syndrome extraction
    syn_extract = circuits.LogicalCircuit(supress_warning=True)
    syn_extract.append(qecc.gate('I', num_syn_extract=1))

    # Choosing basis
    if basis is None or basis == 'zero':
        basis = '|0>'
    elif basis == 'plus':
        basis = '|+>'
    else:
        raise Exception('Basis must be "zero", "plus", "None"!')

    # init circuit
    initzero = circuits.LogicalCircuit(supress_warning=True)
    instr_symbol = 'ideal init %s' % basis
    gate = qecc.gate(instr_symbol)
    initzero.append(gate)

    # print(gate)
    # print(initzero)

    logical_circ_dict = gate.final_instr().final_logical_ops
    logical_ops_sym = gate.final_instr().logical_stabilizers

    if len(logical_circ_dict) != 1:
        raise Exception('This tool expects a code that stores one logical qubit.')

    logical_circ = logical_circ_dict[0][logical_ops_sym[0]]

    num_failure = 0

    for _ in range(runs):
        # State
        state = circuit_runner.init(qecc.num_qudits, simulator=state_sim)

        # Create ideal logical |0>
        circuit_runner.run_logic(state, initzero)
        try:
            total_time += circuit_runner.total_time
        except AttributeError:
            pass

        output, _ = circuit_runner.run_logic(state, syn_extract, error_gen=error_gen, error_params=error_params)
        try:
            total_time += circuit_runner.total_time
        except AttributeError:
            pass

        # syn = output.simplified(True)

        if output:

            # Recovery operation
            recovery = decoder.decode(output)

            # Apply recovery operation
            circuit_runner.run_circuit(state, recovery)

        sign = state.logical_sign(logical_circ)

        num_failure += sign

    logical_rate = float(num_failure) / float(runs)

    if verbose:
        print('\ndistance = %s' % distance)
        print('p = %s' % p)
        print('runs = %s' % runs)

        print('\nlogical error rate: %s' % logical_rate)
        r = float(logical_rate)/float(p)
        print('\nplog/p = %s' % r)

    return logical_rate, total_time


def codecapacity_logical_rate2(runs, qecc, distance, error_gen, error_params, decoder, seed=None, state_sim=None,
                               verbose=True, circuit_runner=None, basis=None):
    """
    A tool for determining the code-capacity logical-error rate for syndrome extraction.

    In this analysis only logical |0> is prepared and each run consists of an ideal logical |0> preparation followed by
    a single round of syndrome extraction. The error rate is determined by number of runs with logical failures divided
    by the total number of runs.

    Args:
        runs: Number of runs to evaluate the logical error rate.
        qecc:
        distance:
        error_gen:
        error_params:
        decoder:
        seed:
        state_sim:
        verbose:
        circuit_runner:

    Returns:

    """

    p = error_params['p']
    total_time = 0.0

    # Circuit simulator
    if circuit_runner is None:
        circuit_runner = circuit_runners.TimingRunner(seed=seed)

    # Syndrome extraction
    syn_extract = circuits.LogicalCircuit(supress_warning=True)
    syn_extract.append(qecc.gate('I', num_syn_extract=1))

    # init logical |0> circuit
    initzero = circuits.LogicalCircuit(supress_warning=True)
    initzero.append(qecc.gate('ideal init |0>'))

    # init logical |+> circuit
    initplus = circuits.LogicalCircuit(supress_warning=True)
    initplus.append(qecc.gate('ideal init |+>'))

    logical_ops_zero = qecc.instruction('instr_init_zero').logical_stabs[0]['Z']
    logical_ops_plus = qecc.instruction('instr_init_plus').logical_stabs[0]['X']

    num_failure = 0

    for _ in range(runs):
        # States
        state0 = circuit_runner.init(qecc.num_qudits, simulator=state_sim)
        state1 = circuit_runner.init(qecc.num_qudits, simulator=state_sim)

        # Create ideal logical |0>
        circuit_runner.run_logic(state0, initzero)
        try:
            total_time += circuit_runner.total_time
        except AttributeError:
            pass

        # Create ideal logical |+>
        circuit_runner.run_logic(state1, initplus)
        try:
            total_time += circuit_runner.total_time
        except AttributeError:
            pass

        output, error_circuits = circuit_runner.run_logic(state0, syn_extract, error_gen=error_gen,
                                                          error_params=error_params)
        try:
            total_time += circuit_runner.total_time
        except AttributeError:
            pass

        # syn = output.simplified(True)

        circuit_runner.run_logic(state1, syn_extract, error_circuits=error_circuits)
        try:
            total_time += circuit_runner.total_time
        except AttributeError:
            pass

        if output:

            # Recovery operation
            recovery = decoder.decode(output)

            # Apply recovery operation

            circuit_runner.run_circuit(state0, recovery)
            try:
                total_time += circuit_runner.total_time
            except AttributeError:
                pass
            circuit_runner.run_circuit(state1, recovery)
            try:
                total_time += circuit_runner.total_time
            except AttributeError:
                pass

        sign0 = state0.logical_sign(logical_ops_zero)
        sign1 = state1.logical_sign(logical_ops_plus)

        if sign0 or sign1:

            num_failure += 1

    logical_rate = float(num_failure) / float(runs)

    if verbose:
        print('\ndistance = %s' % distance)
        print('p = %s' % p)
        print('runs = %s' % runs)

        print('\nlogical error rate: %s' % logical_rate)
        r = float(logical_rate)/float(p)
        print('\nplog/p = %s' % r)
        # print('----')

    return logical_rate, total_time


def codecapacity_logical_rate3(runs, qecc, distance, error_gen, error_params, decoder, seed=None, state_sim=None,
                               max_syn_extract=1e7, circuit_runner=None, verbose=True, init_circuit=None,
                               init_logical_ops=None, basis=None):
    """
    A tool for determining the code-capacity logical-error rate for syndrome extraction.

    In this analysis only logical |0> is prepared and each run consists of an ideal logical |0> preparation followed by
    a single round of syndrome extraction. The error rate is determined by number of runs with logical failures divided
    by the total number of runs.

    !!! This version determines logical threshold from 1/avg(duration)

    Args:
        runs: Number of runs to evaluate the logical error rate.
        qecc:
        distance:
        error_gen:
        error_params:
        decoder:
        seed:
        state_sim:
        max_syn_extract:
        circuit_runner:
        verbose:
        init_circuit:
        init_logical_ops:

    Returns:

    """

    p = error_params['p']
    total_time = 0.0

    # Circuit simulator
    if circuit_runner is None:
        circuit_runner = circuit_runners.TimingRunner(seed=seed)

    if init_circuit is None:
        # init circuit
        init_circuit = circuits.LogicalCircuit(supress_warning=True)

        # Choosing basis
        if basis is None or basis == 'zero':
            basis = '|0>'
        elif basis == 'plus':
            basis = '|+>'
        else:
            raise Exception('Basis must be "zero", "plus", "None"!')

        gate = qecc.gate('ideal init %s' % basis)
        init_circuit.append(gate)

    if init_logical_ops is None:
        # logical_ops = qecc.instruction('instr_syn_extract').final_logical_ops[0]
        ## logical_ops = qecc.instruction('instr_init_zero').logical_stabs[0]

        if init_circuit is None:

            gate = qecc.gate('ideal init %s' % basis)

            # if len(gate.final_logical_stabs()) != 1:
            #    raise Exception('This tool expects a code that stores one logical qubit.')

            # logical_ops = qecc.instruction('instr_syn_extract').final_logical_ops[0]
            # logical_ops = gate.final_logical_stabs()[0]
            # logical_ops_plus = qecc.instruction('instr_init_plus').logical_stabs[0]

            logical_circ_dict = gate.final_instr().final_logical_ops
            logical_ops_sym = gate.final_instr().logical_stabilizers

            if len(logical_circ_dict) != 1:
                raise Exception('This tool expects a code that stores one logical qubit.')

            logical_ops = logical_circ_dict[0][logical_ops_sym[0]]
        else:
            raise Exception('This case is not handled!')
    else:
        logical_ops = init_logical_ops

    # Syndrome extraction
    syn_extract = circuits.LogicalCircuit(supress_warning=True)
    syn_extract.append(qecc.gate('I', num_syn_extract=1))

    # logical_ops = qecc.instruction('instr_syn_extract').final_logical_ops[0]

    run_durations = []

    for _ in range(runs):

        # State
        state = circuit_runner.init(qecc.num_qudits, simulator=state_sim)

        # Create ideal logical |0>
        circuit_runner.run_logic(state, init_circuit)
        try:
            total_time += circuit_runner.total_time
        except AttributeError:
            pass

        for duration in range(max_syn_extract):

            # Run syndrome extraction
            output, _ = circuit_runner.run_logic(state, syn_extract, error_gen=error_gen, error_params=error_params)
            try:
                total_time += circuit_runner.total_time
            except AttributeError:
                pass

            # syn = output.simplified(True)

            if output:
                # Recovery operation
                recovery = decoder.decode(output)

                # Apply recovery operation
                circuit_runner.run_circuit(state, recovery)
                try:
                    total_time += circuit_runner.total_time
                except AttributeError:
                    pass

            sign = state.logical_sign(logical_ops)

            if sign:
                break

        else:
            raise Exception('Max syndrome extraction (%s) met.' % max_syn_extract)

        run_durations.append(duration+1)  # duration + 1 == number of syndrome extractions.

    if verbose:
        print('\nTotal number of runs: %s' % sum(run_durations))

    run_durations = np.array(run_durations)
    duration_mean = np.mean(run_durations)

    logical_rate = 1.0 / duration_mean

    if verbose:
        print('\ndistance = %s' % distance)
        print('p = %s' % p)
        print('Number of failures = %s' % runs)

        print('\nlogical error rate: %s' % logical_rate)
        r = float(logical_rate)/float(p)
        print('\nplog/p = %s' % r)

    return logical_rate, total_time
