from qcware_transpile.matching import (
    TranslationSet,
    trivial_rules,
    untranslatable_instructions,
    untranslated_gates,
    simple_translate,
)
from qcware_transpile.dialects import quasar as quasar_dialect, braket as braket_dialect
from qcware_transpile import TranslationException
from pyrsistent import pset
from icontract.errors import ViolationError
import braket.circuits
import quasar
from toolz.functoolz import thread_first
from typing import Dict, Any


def half_angle(theta):
    return theta / 2.0


def translation_set():
    """
    Creates a translation set from braket to quasar
    """
    trivial_gates = {
        ("I", "I"),
        ("H", "H"),
        ("X", "X"),
        ("Y", "Y"),
        ("Z", "Z"),
        ("S", "S"),
        ("Si", "ST"),
        ("T", "T"),
        ("Ti", "TT"),
        ("CNot", "CX"),
        ("CY", "CY"),
        ("CZ", "CZ"),
        ("CCNot", "CCX"),
        ("PhaseShift", "u1"),
        ("Swap", "SWAP"),
        ("CSwap", "CSWAP"),
        ("Rx", "Rx", half_angle),
        ("Ry", "Ry", half_angle),
        ("Rz", "Rz", half_angle),
    }

    quasar_d = quasar_dialect.dialect()
    braket_d = braket_dialect.dialect()

    rules = pset().union(trivial_rules(braket_d, quasar_d, trivial_gates))
    return TranslationSet.from_trivial_rules(
        from_dialect=braket_d, to_dialect=quasar_d, rules=rules
    )


target_gatenames = sorted([x.name for x in translation_set().to_dialect.gate_defs])
untranslated = sorted([x.name for x in untranslated_gates(translation_set())])


def is_centered(c: braket.circuits.Circuit) -> bool:
    # determine if circuit has leading or following "empty" qubits
    if len(c.qubits):
        min_qubit = c.qubits[0]
        max_qubit = c.qubits[-1]
        return min_qubit == 0 and max_qubit == len(c.qubits) - 1
    return False


def audit(c: braket.circuits.Circuit) -> Dict:
    ir_audit = braket_dialect.audit(c)
    if len(ir_audit.keys()) > 0:
        return ir_audit

    irc = braket_dialect.native_to_ir(c)
    untranslatable = untranslatable_instructions(irc, translation_set())

    result: Dict[str, Any] = {}
    if len(untranslatable) > 0:
        result["untranslatable_instructions"] = untranslatable
    # currently there is no way to express leading or following "empty" qubits in quasar
    if not is_centered(c):
        result["circuit_not_centered"] = True
    return result


def native_is_translatable(c: braket.circuits.Circuit):
    return len(audit(c)) == 0


def translate(c: braket.circuits.Circuit) -> quasar.Circuit:
    """
    Native-to-native translation
    """
    if not native_is_translatable(c):
        raise TranslationException(audit(c))
    try:
        return thread_first(
            c,
            braket_dialect.native_to_ir,
            lambda x: simple_translate(translation_set(), x),
            quasar_dialect.ir_to_native,
        )
    except ViolationError:
        raise TranslationException(audit(c))
