#  =========================================================================  #
#   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.
#  =========================================================================  #

"""
Simple error generator meant to demonstrate a basic error generator that produces errors.
"""
from .class_errors_circuit import ErrorCircuits
from ..circuits.quantum_circuit import QuantumCircuit
from .parent_class_error_gen import ParentErrorGen


class GatewiseGen(ParentErrorGen):
    """
    A simple error generator for the depolarizing model.

    This error generator does not allow much modification of the error model.
    """

    measurements = {'measure X', 'measure Y', 'measure Z'}
    inits = {'init |0>', 'init |1>', 'init |+>', 'init |->', 'init |+i>', 'init |-i>'}
    two_qubits = {'CNOT', 'CZ', 'SWAP', 'G'}
    one_qubits = {'I', 'X', 'Y', 'Z', 'Q', 'Qd', 'R', 'Rd', 'S', 'Sd', 'H', 'H1', 'H2', 'H3', 'H4', 'H5', 'H6',
                  'H+z+x', 'H-z-x', 'H+y-z', 'H-y-z', 'H-x+y', 'H-x-y',
                  'F1', 'F1d', 'F2', 'F2d', 'F3', 'F3d', 'F4', 'F4d'
                  }

    inits_z = {'init |0>', 'init |1>'}
    inits_x = {'init |+>', 'init |->'}
    inits_y = {'init |+i>', 'init |-i>'}

    def __init__(self):
        """

        """

        super().__init__()
        self.gen = self.generator_class()
        self.gen.set_gate_group('measurements', self.measurements)
        self.gen.set_gate_group('inits', self.inits)
        self.gen.set_gate_group('two_qubits', self.two_qubits)
        self.gen.set_gate_group('one_qubits', self.one_qubits)

        # Inherit gen methods
        self.set_gate_error = self.gen.set_gate_error
        self.set_group_error = self.gen.set_group_error
        self.set_default_error = self.gen.set_default_error
        self.set_gate_group = self.gen.set_gate_group

        self.gate_groups = self.gen.gate_groups

        # Inherit error function classes
        self.ErrorStaticSymbol = self.gen.ErrorStaticSymbol
        self.ErrorSet = self.gen.ErrorSet
        self.ErrorSetTwoQuditTensorProduct = self.gen.ErrorSetTwoQuditTensorProduct

    def start(self, logical_circuit, error_params):
        """
        Start up at the beginning of a circuit simulation.

        Args:
            logical_circuit:
            error_params:

        Returns:

        """

        self.error_circuits = ErrorCircuits()
        self.logical_circuit = logical_circuit
        self.error_params = error_params

        return self.error_circuits

    def generate(self, logical_gate, logical_coord, tick_index):
        """
        Returns before errors, after errors, and replaced locations for the given key (args).

        Returns:

        """

        # Simple model where for each gate there is a probability "p" for an X, Y, or Z error to occur.

        before = QuantumCircuit()
        after = QuantumCircuit()
        replace = set([])

        # instr_index = logical_coord[1].instr_index
        instr_index = logical_coord[1]

        # Need to iterate over the tick.
        instr_circuit = logical_gate.circuits[instr_index]

        # Data errors
        # -----------
        if tick_index == 0 and 'data' in self.gen.error_func_dict:
            data_qudits = logical_gate.qecc.data_qudit_set
            self.gen.create_errors(self, 'data', data_qudits, after, before, replace)

        # unitary and measurement errors
        # ------------------------------
        for symbol, gate_locations in instr_circuit.items(tick_index, params=False):

            self.gen.create_errors(self, symbol, gate_locations, after, before, replace)

        # idle errors
        # -----------
        if 'idle' in self.gen.error_func_dict:
            inactive_qudits = logical_gate.qecc.qudit_set - instr_circuit.active_qudits[tick_index]

            if tick_index == 0 and 'data' in self.gen.error_func_dict:
                inactive_qudits -= logical_gate.qecc.data_qudit_set

            self.gen.create_errors(self, 'idle', inactive_qudits, after, before, replace)

        self.error_circuits.add_circuits(key=(logical_coord, tick_index), before_errors=before, after_errors=after)

        # return {'before': before, 'after': after}

    def get_gate_error(self, symbol, gate_locations, error_params):

        self.error_params = error_params

        before = QuantumCircuit()
        after = QuantumCircuit()
        replace = set([])
        self.gen.create_errors(self, symbol, gate_locations, after, before, replace)
        return after, before, replace
