# -*- coding: utf-8 -*-
#   Copyright 2017, 2021 ProjectQ-Framework (www.projectq.ch)
#
#   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.
"""
Register the Z-Y decomposition for an arbitrary one qubit gate.

See paper "Elementary gates for quantum computing" by Adriano Barenco et al.,
arXiv:quant-ph/9503016v1. (Note: They use different gate definitions!)
Or see theorem 4.1 in Nielsen and Chuang.

Decompose an arbitrary one qubit gate U into
U = e^(i alpha) Rz(beta) Ry(gamma) Rz(delta). If a gate V is element of SU(2),
i.e., determinant == 1, then
V = Rz(beta) Ry(gamma) Rz(delta)

"""

import cmath
import itertools
import math

import numpy

from projectq.cengines import DecompositionRule
from projectq.meta import Control, get_control_count
from projectq.ops import BasicGate, Ph, Ry, Rz

TOLERANCE = 1e-12


def _recognize_arb1qubit(cmd):
    """
    Recognize an arbitrary one qubit gate which has a matrix property.

    It does not allow gates which have control qubits as otherwise the
    AutoReplacer might go into an infinite loop. Use
    carb1qubit2cnotrzandry instead.
    """
    try:
        return len(cmd.gate.matrix) == 2 and get_control_count(cmd) == 0
    except AttributeError:
        return False


def _test_parameters(matrix, a, b_half, c_half, d_half):  # pylint: disable=invalid-name
    """
    Build matrix U with parameters (a, b/2, c/2, d/2) and compares against matrix.

    U = [[exp(j*(a-b/2-d/2))*cos(c/2), -exp(j*(a-b/2+d/2))*sin(c/2)],
         [exp(j*(a+b/2-d/2))*sin(c/2), exp(j*(a+b/2+d/2))*cos(c/2)]]

    Args:
        matrix(list): 2x2 matrix
        a: parameter of U
        b_half: b/2. parameter of U
        c_half: c/2. parameter of U
        d_half: d/2. parameter of U

    Returns:
        True if matrix elements of U and `matrix` are TOLERANCE close.
    """
    unitary = [
        [
            cmath.exp(1j * (a - b_half - d_half)) * math.cos(c_half),
            -cmath.exp(1j * (a - b_half + d_half)) * math.sin(c_half),
        ],
        [
            cmath.exp(1j * (a + b_half - d_half)) * math.sin(c_half),
            cmath.exp(1j * (a + b_half + d_half)) * math.cos(c_half),
        ],
    ]
    return numpy.allclose(unitary, matrix, rtol=10 * TOLERANCE, atol=TOLERANCE)


def _find_parameters(matrix):  # pylint: disable=too-many-branches,too-many-statements
    """
    Find decomposition parameters.

    Given a 2x2 unitary matrix, find the parameters a, b/2, c/2, and d/2 such that
    matrix == [[exp(j*(a-b/2-d/2))*cos(c/2), -exp(j*(a-b/2+d/2))*sin(c/2)],
               [exp(j*(a+b/2-d/2))*sin(c/2), exp(j*(a+b/2+d/2))*cos(c/2)]]

    Note:
        If the matrix is element of SU(2) (determinant == 1), then we can choose a = 0.

    Args:
        matrix(list): 2x2 unitary matrix

    Returns:
        parameters of the matrix: (a, b/2, c/2, d/2)
    """
    # Determine a, b/2, c/2 and d/2 (3 different cases).
    # Note: everything is modulo 2pi.
    # Case 1: sin(c/2) == 0:
    if abs(matrix[0][1]) < TOLERANCE:
        two_a = cmath.phase(matrix[0][0] * matrix[1][1]) % (2 * math.pi)
        if abs(two_a) < TOLERANCE or abs(two_a) > 2 * math.pi - TOLERANCE:
            # from 2a==0 (mod 2pi), it follows that a==0 or a==pi,
            # w.l.g. we can choose a==0 because (see U above)
            # c/2 -> c/2 + pi would have the same effect as as a==0 -> a==pi.
            a = 0  # pylint: disable=invalid-name
        else:
            a = two_a / 2.0  # pylint: disable=invalid-name
        d_half = 0  # w.l.g
        b = cmath.phase(matrix[1][1]) - cmath.phase(matrix[0][0])  # pylint: disable=invalid-name
        possible_b_half = [
            (b / 2.0) % (2 * math.pi),
            (b / 2.0 + math.pi) % (2 * math.pi),
        ]
        # As we have fixed a, we need to find correct sign for cos(c/2)
        possible_c_half = [0.0, math.pi]
        found = False
        for b_half, c_half in itertools.product(possible_b_half, possible_c_half):
            if _test_parameters(matrix, a, b_half, c_half, d_half):
                found = True
                break
        if not found:
            raise Exception(
                "Couldn't find parameters for matrix ",
                matrix,
                "This shouldn't happen. Maybe the matrix is " + "not unitary?",
            )
    # Case 2: cos(c/2) == 0:
    elif abs(matrix[0][0]) < TOLERANCE: