# This code is part of Tergite
#
# (C) Copyright Axel Andersson 2022
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.
#
# Alteration Notice
# -----------------
# This code was modified from the original by:
#
# - Martin Ahindura 2023
# - Stefan Hill 2023
"""Handles the creation of schedules for the devices of type TergiteBackend"""
from typing import TYPE_CHECKING, Iterable

import numpy as np
import qiskit.circuit as circuit
import qiskit.pulse as pulse
from qiskit.pulse.library import SymbolicPulse
from sympy import symbols

from .functions import delta_t_function_sympy

if TYPE_CHECKING:
    from ...types.backend import TergiteBackend
    from ..api_client.dtos import DeviceCalibration


def rx(
    backend: "TergiteBackend",
    qubits: Iterable,
    rx_theta: circuit.Parameter,
    device_properties: "DeviceCalibration",
) -> pulse.ScheduleBlock:
    """Creates a rotation around the x-axis on the Bloch sphere for the list of qubits in the backend.

    It returns a backend-specific schedule which implements a
    rotation on a set of qubits around the x-axis on the Bloch sphere.
    The schedule is parameterised by θ [rad], which specifies the angle.

    Args:
        backend: the backend for which the schedule is to be created.
        qubits: the set of qubits to be affected.
        rx_theta: the theta parameter for the circuit
        device_properties: the device parameters of the backend

    Returns:
        qiskit.pulse.ScheduleBlock: the schedule implementing the rotation
    """
    qubit = device_properties.qubits

    sched = pulse.ScheduleBlock(name=f"RX(θ, {qubits})")
    for q in qubits:
        sched += pulse.SetFrequency(
            qubit[q].frequency.value,
            channel=backend.drive_channel(q),
        )
        duration = round(qubit[q].pi_pulse_duration.value / backend.dt)
        calibration_sigma = qubit[q].pulse_sigma.value
        sigma = (
            round(calibration_sigma / backend.dt)
            if calibration_sigma
            else round(duration / 4)
        )
        amplitude = rx_theta / np.pi * qubit[q].pi_pulse_amplitude.value
        motzoi_params = device_properties.qubits[q].pi_pulse_motzoi

        motzoi = motzoi_params.value if motzoi_params else 0.0
        sched += pulse.Play(
            pulse.Drag(
                duration=duration,
                amp=amplitude,
                sigma=sigma,
                beta=motzoi,
                name=f"RX q{q}",
            ),
            channel=backend.drive_channel(q),
        )

    return sched


def rz(
    backend: "TergiteBackend", qubits: Iterable, rz_lambda: circuit.Parameter
) -> pulse.ScheduleBlock:
    """Creates a rotation around the z-axis on the Bloch sphere for the list of qubits in the backend.

    It returns a backend-specific schedule which implements a rotation on a set
    of qubits around the z-axis on the Block sphere. The schedule is
    parameterised by λ [rad], which specifies the angle.

    Args:
        backend: the backend for which the schedule is to be created.
        qubits: the set of qubits to be affected.
        rz_lambda: the lambda parameter for the circuit

    Returns:
        qiskit.pulse.ScheduleBlock: the schedule implementing the rotation
    """
    sched = pulse.ScheduleBlock(name=f"RZ(λ, {qubits})")
    for q in qubits:
        sched += pulse.ShiftPhase(
            rz_lambda, channel=backend.drive_channel(q), name=f"RZ q{q}"
        )
    return sched


def cz(
    backend: "TergiteBackend",
    control_qubit_idxs: Iterable,
    target_qubit_idxs: Iterable,
    device_properties: "DeviceCalibration",
) -> pulse.ScheduleBlock:
    """Two qubit CNOT gate building block. TODO: Add doc comment."""
    # FIXME: Why choose to pass control_qubits and target_qubits as Iterables yet
    #    when being called, only as single value is passed
    sched = pulse.ScheduleBlock(
        name=f"CZ(θ, control={control_qubit_idxs}, target={target_qubit_idxs})"
    )

    for control_qubit_idx, target_qubit_idx in zip(
        control_qubit_idxs, target_qubit_idxs
    ):
        # FIXME: there is a an assumption how we choose target and control qubits for local phases
        # This should be replaced with proper explicit labels for target and control qubit ids per coupler
        # The assumption: when we have two qubits q_x and q_y coupled by a coupler
        # the control qubit id from calibration config is even, target id is odd
        control_qubit_id_from_calibration = device_properties.qubits[
            control_qubit_idx
        ].id
        is_control_qubit_even = control_qubit_id_from_calibration % 2 == 0
        cz_control_qubit_idx = (
            control_qubit_idx if is_control_qubit_even else target_qubit_idx
        )
        cz_target_qubit_idx = (
            target_qubit_idx if is_control_qubit_even else control_qubit_idx
        )

        control_qubit = device_properties.qubits[cz_control_qubit_idx]
        target_qubit = device_properties.qubits[cz_target_qubit_idx]

        # Get control channels between control and target qubits

        control_channels = backend.control_channel((control_qubit.id, target_qubit.id))

        c_props = [
            x for x in device_properties.couplers if x.id == control_channels[0].index
        ]
        c_props = c_props[0] if c_props else None

        # TODO: raise error if c_props is none

        f1 = min(_val(target_qubit.frequency), _val(control_qubit.frequency))
        f0 = max(_val(target_qubit.frequency), _val(control_qubit.frequency))

        alpha0 = _val(c_props.anharmonicity, default=-0.17e9)

        f2 = _val(c_props.frequency)  # Coupler frequency
        detuning = _val(
            c_props.frequency_detuning, default=0
        )  # detuning in radial frequency

        args = {
            "delta_0": _val(c_props.cz_pulse_amplitude),  # Maximum delta value
            "Theta": _val(c_props.cz_pulse_dc_bias, default=0.0),  # DC bias term
            "omega_c0": 2 * np.pi * f2,  # Maximum frequency in Hz
            "omega_Phi": detuning
            + 2 * np.pi * (f1 - f0 - alpha0),  # Transition frequency in Hz
            "phi": _val(c_props.cz_pulse_phase_offset, default=0.0),  # Phase offset
            "t_w": _val(
                c_props.cz_pulse_duration_before, default=4e-9
            ),  # s, duration before pulse
            "t_rf": _val(c_props.cz_pulse_duration_rise, default=4e-9),  # s, rise time
            "t_p": _val(
                c_props.cz_pulse_duration_constant
            ),  # s, constant pulse duration
        }

        # we now set wait time before and after pulse explicitly for both sim and real backends
        t_gate = args["t_p"] + args["t_rf"]
        # required param for pulse, for display purposes delta_0 is equivalent to max_amplitude
        amp = args["delta_0"]
        # this is for display purposes, as we overriding frequency modulation in backend
        freq = int(f2)

        args["amp"] = amp
        args["freq"] = freq
        control_rz_lambda = c_props.control_rz_lambda.value
        target_rz_lambda = c_props.target_rz_lambda.value
        cz_gate_duration = round(t_gate / backend.dt)
        cz_wait_duration = round(args["t_w"] / backend.dt)

        cz_gate = wacqt_cz_gate(
            duration=cz_gate_duration, name="cz_pulse", numerical_args=args
        )
        sched += pulse.SetFrequency(
            c_props.frequency.value,
            channel=control_channels[0],
        )
        sched += pulse.Delay(channel=control_channels[0], duration=cz_wait_duration)
        sched += pulse.Play(cz_gate, control_channels[0])
        sched += pulse.Delay(channel=control_channels[0], duration=cz_wait_duration)

        if control_rz_lambda and target_rz_lambda:
            sched += pulse.ShiftPhase(
                control_rz_lambda,
                channel=backend.drive_channel(cz_control_qubit_idx),
                name=f"RZ q{cz_control_qubit_idx}",
            )
            sched += pulse.ShiftPhase(
                target_rz_lambda,
                channel=backend.drive_channel(cz_target_qubit_idx),
                name=f"RZ q{cz_target_qubit_idx}",
            )

    return sched


def measure(
    backend: "TergiteBackend",
    qubits: Iterable,
    device_properties: "DeviceCalibration",
) -> pulse.ScheduleBlock:
    """Creates a measurement on the list of qubits in the given backend.

    It returns a backend-specific schedule which implements
    a measurement on a set of qubits.

    Args:
        backend: the backend for which the schedule is to be created.
        qubits: the set of qubits to be affected.
        device_properties: the device parameters of the backend

    Returns:
        qiskit.pulse.ScheduleBlock: the schedule implementing the measurement
    """
    readout_resonator_props = device_properties.resonators

    sched = pulse.ScheduleBlock(name=f"Measure({qubits})")
    for q in qubits:
        readout_resonator = readout_resonator_props[q]
        sched += pulse.SetFrequency(
            readout_resonator.frequency.value, channel=backend.measure_channel(q)
        )
        sched += pulse.Delay(
            channel=backend.measure_channel(q), duration=round(4e-9 / backend.dt)
        )
        sched += pulse.Delay(
            channel=backend.acquire_channel(q), duration=round(4e-9 / backend.dt)
        )
        sched += pulse.Play(
            pulse.Constant(
                amp=readout_resonator.pulse_amplitude.value,
                duration=round(readout_resonator.pulse_duration.value / backend.dt),
                name=f"Readout q{q}",
            ),
            channel=backend.measure_channel(q),
        )
        sched += pulse.Delay(
            duration=round(readout_resonator.acq_delay.value / backend.dt),
            channel=backend.acquire_channel(q),
            name=f"Time of flight q{q}",
        )
        sched += pulse.Acquire(
            duration=round(readout_resonator.acq_integration_time.value / backend.dt),
            channel=backend.acquire_channel(q),
            mem_slot=backend.memory_slot(q),
            name=f"Integration window q{q}",
        )
    return sched


def delay(
    backend: "TergiteBackend",
    qubits: Iterable,
    delay_tau: circuit.Parameter,
    delay_str: str = "Delay",
) -> pulse.ScheduleBlock:
    """Creates a delay on the set of qubits in the backend.

    It returns a backend-specific schedule which implements a delay
    operation on a set of qubits.
    The schedule is parameterised by τ [ns], which specifies the delay duration.

    Args:
        backend: the backend for which the schedule is to be created.
        qubits: the set of qubits to be affected.
        delay_tau: the tau parameter for the circuit
        delay_str: the delay string for the instruction set

    Returns:
        qiskit.pulse.ScheduleBlock: the schedule implementing the delay
    """
    sched = pulse.ScheduleBlock(name=f"{delay_str}({qubits}, τ)")
    for q in qubits:
        sched += pulse.Delay(
            duration=delay_tau,
            channel=backend.drive_channel(q),
            name=f"{delay_str} q{q}",
        )
        sched += pulse.Delay(
            duration=delay_tau,
            channel=backend.measure_channel(q),
            name=f"{delay_str} q{q}",
        )
        sched += pulse.Delay(
            duration=delay_tau,
            channel=backend.acquire_channel(q),
            name=f"{delay_str} q{q}",
        )
    return sched


def wacqt_cz_gate(duration, name, numerical_args):
    # define the time variable
    t = symbols("t", real=True)

    # define symbolic variables (can also be passed as parameters)
    symbolic_args = {
        "t_w": symbols("t_w", real=True),
        "t_rf": symbols("t_rf", real=True),
        "t_p": symbols("t_p", real=True),
        "delta_0": symbols("delta_0", real=True),
    }

    # create the symbolic expression
    envelope_expr = delta_t_function_sympy(t, symbolic_args)

    # substitute numerical values into the symbolic expression
    if numerical_args:
        envelope_expr = envelope_expr.subs(numerical_args)

    numerical_args["amp"] = [numerical_args["amp"], 0]
    # create the SymbolicPulse instance
    instance = SymbolicPulse(
        pulse_type="Wacqt_cz_gate_pulse",
        duration=duration,
        parameters=numerical_args,
        envelope=envelope_expr,
        name=name,
    )

    return instance


def _val(x, default=None):
    """Unwrap calibration parameters safely."""
    if x is None:
        return default
    v = getattr(x, "value", x)
    return default if v is None else v
