"""Submodule for single-qubit gate calibration.

"""
from __future__ import annotations

import logging
from typing import Callable, Dict, List, Tuple

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from IPython.display import clear_output
from lmfit import Parameters
from numpy.typing import ArrayLike, NDArray

from qcal.circuit import Barrier, Circuit, CircuitSet, Cycle
from qcal.config import Config
from qcal.fitting.fit import (
    FitAbsoluteValue,
    FitCosine,
    FitDecayingCosine,
    FitLinear,
    FitParabola,
)
from qcal.fitting.utils import est_freq_fft
from qcal.gate.single_qubit import X90, Idle, Rz, X
from qcal.math.utils import round_to_order_error, wrap_phase
from qcal.plotting.utils import calculate_nrows_ncols
from qcal.qpu.qpu import QPU
from qcal.settings import Settings
from qcal.units import MHz, us

from .calibration import Calibration
from .utils import find_pulse_index, in_range

logger = logging.getLogger(__name__)


def Amplitude(
    qpu:          QPU,
    config:       Config,
    qubits:       List | Tuple,
    amplitudes:   ArrayLike | NDArray | Dict[ArrayLike | NDArray],
    gate:         str = 'X90',
    method:       str = 'Rabi',
    subspace:     str = 'GE',
    n_gates:      int = 1,
    relative_amp: bool = False,
    params:       Dict | None = None,
    **kwargs
) -> Callable:
    """Amplitude calibration for single-qubit gates.

    Basic example useage for initial calibration:

    ```
    amps = np.linspace(0, 1.0, 21)
    cal = Amplitude(
        CustomQPU,
        config,
        qubits=[0, 1, 2],
        amplitudes=amps)
    cal.run()
    ```

    Args:
        qpu (QPU): custom QPU object.
        config (Config): qcal Config object.
        qubits (List | Tuple): qubits to calibrate.
        amplitudes (ArrayLike | NDArray | Dict[ArrayLike | NDArray]): array of
            gate amplitudes to sweep over for calibrating the single-qubit gate
            amplitude. If calibrating multiple qubits at the same time, this
            should be a dictionary mapping an array to a qubit label.
        gate (str, optional): native gate to calibrate. Defaults to 'X90'.
        method (str, optional): method used to calibrate a gate. Defaults to
            'Rabi'. Other option is 'RAP'.
        subspace (str, optional): qubit subspace for the defined gate.
            Defaults to 'GE'.
        n_gates (int, optional): number of gates for pulse repetition.
            Defaults to 1.
        relative_amp (bool, optional): whether or not the amplitudes argument
            is defined relative to the existing pulse amplitude. Defaults to
            False. If True, the amplitudes are swept over the current amplitude
            times the amplitudes argument.

    Returns:
        Callable: Amplitude calibration class.
    """

    class Amplitude(qpu, Calibration):
        """Amplitude calibration class.

        This class inherits a custom QPU from the Amplitude calibration
        function.
        """

        def __init__(
            self,
            config:       Config,
            qubits:       List | Tuple,
            amplitudes:   ArrayLike | NDArray | Dict,
            gate:         str = 'X90',
            method:       str = 'Rabi',
            subspace:     str = 'GE',
            n_gates:      int = 1,
            relative_amp: bool = False,
            params:       Dict | None = None,
            **kwargs
        ) -> None:
            """Initialize the Amplitude calibration class within the function.
            """
            n_levels = 3 if subspace == 'EF' else 2
            qpu.__init__(self, config=config, n_levels=n_levels, **kwargs)
            Calibration.__init__(self, config)

            self._qubits = qubits

            assert gate in ('X90', 'X'), (
                "'gate' must be one of 'X90' or 'X'!"
            )
            self._gate = gate

            assert method.upper() in ('RABI', 'RAP'), (
                "'method' must be one of 'Rabi' or 'RAP'!"
            )
            self._method = method

            assert subspace in ('GE', 'EF'), (
                "'subspace' must be one of 'GE' or 'EF'!"
            )
            self._subspace = subspace

            if not isinstance(amplitudes, dict):
                self._amplitudes = dict.fromkeys(qubits, amplitudes)
            else:
                self._amplitudes = amplitudes
            self._param_sweep = self._amplitudes

            if params:
                self._params = params
            else:
                self._params = {}
                for q in qubits:
                    param = f'single_qubit/{q}/{subspace}/{gate}/pulse'
                    param += f'/{find_pulse_index(self._config, param)}/'
                    param += 'kwargs/amp'
                    self._params[q] = param

            if relative_amp:
                for q in qubits:
                    self._amplitudes[q] = (
                        self._config[self._params[q]] * self._amplitudes[q]
                    )
            self._relative_amp = relative_amp

            if n_gates > 1 and gate == 'X90' and method.upper() == 'RABI':
                assert n_gates % 4 == 0, 'n_gates must be a multiple of 4!'
            elif n_gates > 1 and gate == 'X' and method.upper() == 'RABI':
                assert n_gates % 2 == 0, 'n_gates must be a multiple of 2!'
            self._n_gates = n_gates

            if n_gates == 1:
                self._fit = {q: FitCosine() for q in qubits}
            elif n_gates > 1:
                self._fit = {q: FitParabola() for q in qubits}

        @property
        def amplitudes(self) -> Dict:
            """Amplitude sweep for each qubit.

            Returns:
                Dict: qubit to array map.
            """
            return self._amplitudes

        def generate_circuits(self):
            """Generate all amplitude calibration circuits."""
            logger.info(' Generating circuits...')

            gate = {'X90': X90, 'X': X}
            level = {2: 'GE', 3: 'EF'}
            circuit = Circuit()

            # Prepulse for EF calibration
            if self._subspace == 'EF':

                if self._gate == 'X90':
                    circuit.extend([
                        Cycle({
                            gate[self._gate](q, subspace='GE')
                            for q in self._qubits
                        }),
                        Barrier(self._qubits),
                        Cycle({
                            gate[self._gate](q, subspace='GE')
                            for q in self._qubits
                        }),
                        Barrier(self._qubits)
                    ])

                elif self._gate == 'X':
                    circuit.extend([
                        Cycle({
                            gate[self._gate](q, subspace='GE')
                            for q in self._qubits
                        }),
                        Barrier(self._qubits)
                    ])

            for _ in range(self._n_gates):
                circuit.extend([
                    Cycle({
                        gate[self._gate](q, subspace=level[self._n_levels])
                        for q in self._qubits
                    }),
                    Barrier(self._qubits)
                ])
            circuit.measure()

            circuits = []
            for _ in range(self._amplitudes[self._qubits[0]].size):
                circuits.append(circuit.copy())

            self._circuits = CircuitSet(circuits=circuits)
            for q in self._qubits:
                self._circuits[
                        f'param: {self._params[q]}'
                    ] = self._amplitudes[q]

        def analyze(self) -> None:
            """Analyze the data."""
            logger.info(' Analyzing the data...')

            # Fit the probability of being in 0 from the gate amplitude sweep
            # to a cosine or parabola
            level = {'GE': '0', 'EF': '1'}
            for i, q in enumerate(self._qubits):
                prob = []
                for circuit in self._circuits:
                    prob.append(
                        circuit.results.marginalize(i).populations[
                                level[self._subspace]
                            ]
                    )
                self._sweep_results[q] = prob
                self._circuits[f'Q{q}: Prob({level[self._subspace]})'] = prob

                # Cosine fit
                if self._n_gates == 1 and self._method.upper() != 'RAP':
                    est_freq = est_freq_fft(self._amplitudes[q], prob)
                    params = Parameters()
                    params.add('amp', value=0.5, min=0., max=1.)
                    params.add('freq', value=est_freq)
                    params.add('phase', value=0., min=-np.pi, max=np.pi)
                    params.add('offset', value=0.5, min=0., max=1.)
                    self._fit[q].fit(self._amplitudes[q], prob, params=params)
                # Quadratic fit
                elif self._n_gates > 1 and self._method.upper() != 'RAP':
                    params = Parameters()
                    params.add('a', value=-1.)
                    params.add('b', value=1.)
                    params.add('c', value=1.)
                    self._fit[q].fit(self._amplitudes[q], prob, params=params)

                # If the fit was successful, write find the new amp
                if self._fit[q].fit_success:

                    period_frac = {'X90': 0.25, 'X': 0.5}
                    if self._n_gates == 1:  # Cosine fit
                        freq = self._fit[q].fit_params['freq'].value
                        # _, freq, _, _ = self._fit[q].fit_params
                        newvalue = 1 / freq * period_frac[self._gate]
                        if not in_range(newvalue, self._amplitudes[q]):
                            logger.warning(
                              f'Fit failed for qubit {q} (out of range)!'
                            )
                            self._fit[q]._fit_success = False
                        else:
                            self._cal_values[q] = newvalue

                    elif self._n_gates > 1:  # Quadratic fit
                        a = self._fit[q].fit_params['a'].value
                        b = self._fit[q].fit_params['b'].value
                        # a, b, _ = self._fit[q].fit_params
                        newvalue = -b / (2 * a)  # Assume c = 0
                        if a > 0:
                            logger.warning(
                              f'Fit failed for qubit {q} (positive curvature)!'
                            )
                            self._fit[q]._fit_success = False
                            self._cal_values[q] = self._param_sweep[q][
                                np.array(prob).argmax()
                            ]
                        elif not in_range(newvalue, self._amplitudes[q]):
                            logger.warning(
                              f'Fit failed for qubit {q} (out of range)!'
                            )
                            self._fit[q]._fit_success = False
                            self._cal_values[q] = self._param_sweep[q][
                                np.array(prob).argmax()
                            ]
                        else:
                            self._cal_values[q] = newvalue

                else:

                    if self._n_gates == 1:
                        self._cal_values[q] = self._param_sweep[q][
                            np.array(prob).argmin()
                        ]

                    elif self._n_gates > 1:
                        self._cal_values[q] = self._param_sweep[q][
                            np.array(prob).argmax()
                        ]

        def save(self):
            """Save all circuits and data."""
            qpu.save(self)
            self._data_manager.save_to_csv(
                 pd.DataFrame([self._cal_values]), 'calibrated_values'
            )

        def run(self):
            """Run all experimental methods and analyze results."""
            self.generate_circuits()
            qpu.run(self, self._circuits, save=False)
            self.analyze()
            clear_output(wait=True)
            self._data_manager._exp_id += (
                f'_amp_cal_{"".join("Q" + str(q) for q in self._qubits)}'
            )
            if Settings.save_data:
                self.save()
            self.plot(
                xlabel='Amplitude (a.u.)',
                ylabel=(
                    r'$|0\rangle$ Population' if self._subspace == 'GE' else
                    r'$|1\rangle$ Population'
                ),
                save_path=self._data_manager._save_path
            )
            self.final()
            print(f"\nRuntime: {repr(self._runtime)[8:]}\n")

    return Amplitude(
        config=config,
        qubits=qubits,
        amplitudes=amplitudes,
        gate=gate,
        method=method,
        subspace=subspace,
        n_gates=n_gates,
        relative_amp=relative_amp,
        params=params,
        **kwargs
    )


def Frequency(
    qpu:        QPU,
    config:     Config,
    qubits:     List | Tuple,
    t_max:      float = 2*us,
    detunings:  NDArray = np.array([-2, -1, 1, 2])*MHz,  # noqa: B008
    subspace:   str = 'GE',
    n_elements: int = 30,
    params:     Dict | None = None,
    **kwargs
) -> Callable:
    """Frequency calibration for single qubits.

    Basic example useage:

    ```
    cal = Frequency(
        CustomQPU,
        config,
        qubits=[0],
        detunings=np.array([-2, -1, 1, 2])*MHz)
    cal.run()
    ```

    Args:
        qpu (QPU): custom QPU object.
        config (Config): qcal Config object.
        qubits (List | Tuple): qubits to measure.
        t_max (float, optional): maximum Ramsey time. Defaults to 2 us.
        detunings (float, optional): artificial detunings from the actual qubit
            frequency. Defaults to np.array([-2, -1, 1, 2])*MHz.
        subspace (str, optional): qubit subspace for frequency calibration.
            Defaults to 'GE'.
        n_elements (int, optional): number of delays starting from 0 to t_max.
            Defaults to 30.

    Returns:
        Callable: Frequency calibration class.
    """

    class Frequency(qpu, Calibration):
        """Frequency calibration class.

        This class inherits a custom QPU from the Frequency calibration
        function.
        """

        def __init__(
            self,
            config:     Config,
            qubits:     List | Tuple,
            t_max:      float = 2*us,
            detunings:  NDArray = np.array([-2, -1, 1, 2])*MHz,  # noqa: B008
            subspace:   str = 'GE',
            n_elements: int = 30,
            params:     Dict | None = None,
            **kwargs
        ) -> None:
            """Initialize the Frequency experiment class within the function.
            """
            # n_levels = 3 if subspace == 'EF' else 2
            qpu.__init__(self, config=config, **kwargs)
            Calibration.__init__(self, config)

            self._qubits = qubits
            self._detunings = detunings

            assert subspace in ('GE', 'EF'), (
                "'subspace' must be one of 'GE' or 'EF'!"
            )
            self._subspace = subspace

            self._times = {
                q: np.linspace(0., t_max, n_elements) for q in qubits
            }
            self._param_sweep = self._times

            if params:
                self._params = params
            else:
                self._params = {
                    q: f'single_qubit/{q}/{subspace}/freq' for q in qubits
                }
            self._sweep_results = {
                q: [] for q in qubits
            }
            self._fit = {
                q: FitAbsoluteValue() for q in qubits
            }
            self._freq_fit = {
                q: [FitDecayingCosine() for _ in range(detunings.size)]
                for q in qubits
            }

            self._freqs = {q: [] for q in qubits}
            self._fit_detunings = {q: [] for q in self._qubits}

        @property
        def detunings(self) -> NDArray:
            """Frequency detunings.

            Returns:
                NDArray: frequency detunings..
            """
            return self._detunings

        @property
        def times(self) -> Dict:
            """Time sweep for each qubit.

            Returns:
                Dict: qubit to time array map.
            """
            return self._times

        def generate_circuits(self):
            """Generate all amplitude calibration circuits."""
            logger.info(' Generating circuits...')

            circuits = []
            times = []
            detunings = []
            for detuning in self._detunings:
                for t in self._times[self._qubits[0]]:
                    times.append(t)
                    detunings.append(detuning)
                    phase = 2. * np.pi * detuning * t

                    circuit = Circuit()

                    # State prepration
                    circuit.extend([
                        Cycle({X90(q, subspace='GE') for q in self._qubits}),
                        Barrier(self._qubits)
                    ])
                    if self._subspace == 'EF':
                        circuit.extend([
                          Cycle({X90(q, subspace='GE') for q in self._qubits}),
                          Barrier(self._qubits),
                          Cycle({X90(q, subspace='EF') for q in self._qubits}),
                          Barrier(self._qubits)
                        ])

                    # Ramsey experiment
                    circuit.extend([
                        Cycle({Idle(q, duration=t) for q in self._qubits}),
                        Barrier(self._qubits),
                        Cycle({Rz(q, phase, subspace=self._subspace)
                              for q in self._qubits}),
                        Barrier(self._qubits),
                    ])

                    # Basis preparation
                    circuit.append(
                        Cycle({X90(q, subspace=self._subspace)
                            for q in self._qubits})
                    )

                    circuit.measure()
                    circuits.append(circuit)

            self._circuits = CircuitSet(circuits=circuits)
            self._circuits['detuning'] = detunings
            self._circuits['time'] = times

        def analyze(self) -> None:
            """Analyze the data."""
            logger.info(' Analyzing the data...')

            level = {'GE': '0', 'EF': '1'}
            # Fit the measured frequency for each detuning
            probs = {q: [] for q in self._qubits}
            for i, detuning in enumerate(self._detunings):
                for j, q in enumerate(self._qubits):
                    prob = []
                    for circuit in self._circuits[
                        self._circuits['detuning'] == detuning].circuit:
                        prob.append(
                            circuit.results.marginalize(j).populations[
                                level[self._subspace]
                            ]
                        )
                    probs[q] += prob
                    self._sweep_results[q].append(prob)

                    e = np.array(prob).min()
                    a = np.array(prob).max() - e
                    b = -np.mean( np.diff(prob) / np.diff(self._times[q]) ) / a
                    params = Parameters()
                    params.add('a', value=a)
                    params.add('b', value=b)
                    params.add('c', value=detuning)
                    params.add('d', value=0.)
                    params.add('e', value=e)
                    self._freq_fit[q][i].fit(
                        self._times[q], prob, params=params
                    )

                    # If the fit was successful, grab the frequency
                    if self._freq_fit[q][i].fit_success:
                        self._freqs[q].append(
                            abs(self._freq_fit[q][i].fit_params['c'].value)
                        )
                        self._fit_detunings[q].append(detuning)

            for q in self._qubits:
                self._circuits[f'Q{q}: Prob({level[self._subspace]})'] = (
                    probs[q]
                )

            # Fit the characterized frequencies to an absolute value curve
            for q in self._qubits:
                params = Parameters()
                params.add('a', value=1)
                params.add('b', value=0)
                params.add('c', value=0)
                self._fit[q].fit(self._detunings, self._freqs[q], params=params)

                if self._fit[q].fit_success:
                    a = self._fit[q].fit_params['a'].value
                    newval, err = round_to_order_error(
                        self._fit[q].fit_params['b'].value,
                        self._fit[q].fit_params['a'].stderr,
                        2
                    )

                    if a < 0:
                        logger.warning(
                            f'Fit failed for qubit {q} (negative curvature)!'
                        )
                        self._fit[q]._fit_success = False
                    if not in_range(newval, self._detunings):
                        logger.warning(
                            f'Fit failed for qubit {q} (out of range)!'
                        )
                        self._fit[q]._fit_success = False
                    else:
                        self._cal_values[q] = (
                            self._config[self._params[q]] + newval
                        )
                        self._errors[q] = err

        def save(self) -> None:
            """Save all circuits and data."""
            qpu.save(self)
            self._data_manager.save_to_csv(
                 pd.DataFrame([self._cal_values]), 'freq_values'
            )
            self._data_manager.save_to_csv(
                 pd.DataFrame([self._errors]), 'freq_fit_errors'
            )

        def plot(self) -> None:
            """Plot the frequency sweep and fit results."""
            nrows = len(self._qubits)
            figsize = (10, 4 * nrows)
            fig, axes = plt.subplots(
                nrows, 2, figsize=figsize, layout='constrained'
            )
            colors = plt.get_cmap('viridis', self._detunings.size)
            level = {'GE': '0', 'EF': '1'}

            for i, q in enumerate(self._qubits):
                if len(self._qubits) == 1:
                    ax = axes
                else:
                    ax = axes[i]

                ax[0].set_xlabel(r'Time ($\mu$s)', fontsize=15)
                ax[1].set_xlabel('Detuning (MHz)', fontsize=15)
                ax[0].set_ylabel(
                    rf'$|{level[self._subspace]}\rangle$ Population',
                    fontsize=15
                )
                ax[1].set_ylabel('Measured Detuning (MHz)', fontsize=15)
                ax[0].tick_params(axis='both', which='major', labelsize=12)
                ax[1].tick_params(axis='both', which='major', labelsize=12)
                ax[0].grid(True)
                ax[1].grid(True)

                # Plot the Ramsey sweeps
                for j, det in enumerate(self._detunings):
                    ax[0].plot(
                        self._times[q],
                        np.array(self._sweep_results[q][j]) + j,
                        'o',
                        c=colors(j),
                        label=f'{round(det / MHz, 1)} MHz'
                    )

                    if self._freq_fit[q][j].fit_success:
                        x = np.linspace(
                            self._times[q][0], self._times[q][-1], 100
                        )
                        ax[0].plot(x,
                            self._freq_fit[q][j].predict(x) + j,
                            ls='-',
                            c=colors(j)
                        )

                # Plot the frequency fit
                ax[1].plot(
                    self._fit_detunings[q],
                    self._freqs[q],
                    'o',
                    c='blue',
                    label=f'Meas, Q{q}'
                )
                if self._fit[q].fit_success:
                    x = np.linspace(
                            self._detunings[0],
                            self._detunings[-1],
                            100
                        )
                    ax[1].plot(x,
                        self._fit[q].predict(x),
                        ls='-',
                        c='orange',
                        label='Fit'
                    )
                    df = round(self._fit[q].fit_params['b'].value / MHz, 3)
                    ax[1].axvline(
                        round(self._fit[q].fit_params['b'].value, 2),
                        ls='--', c='k', label=rf'$\Delta f$ = {df} MHz'
                    )

                ax[0].legend(loc=0, fontsize=12)
                ax[1].legend(loc=0, fontsize=12)
                ax[0].xaxis.set_major_formatter(
                        lambda x, pos: round(x / 1e-6, 1)
                    )
                ax[1].xaxis.set_major_formatter(
                        lambda x, pos: round(x / MHz, 1)
                    )
                ax[1].yaxis.set_major_formatter(
                        lambda x, pos: round(x / MHz, 1)
                    )

            fig.set_tight_layout(True)
            if Settings.save_data:
                fig.savefig(
                    self._data_manager._save_path + 'freq_calibration.png',
                    dpi=300
                )
                fig.savefig(
                    self._data_manager._save_path + 'freq_calibration.pdf'
                )
                fig.savefig(
                    self._data_manager._save_path + 'freq_calibration.svg'
                )
            plt.show()

        def run(self):
            """Run all experimental methods and analyze results."""
            self.generate_circuits()
            qpu.run(self, self._circuits, save=False)
            self.analyze()
            clear_output(wait=True)
            self._data_manager._exp_id += (
                f'_freq_cal_{"".join("Q" + str(q) for q in self._qubits)}'
            )
            if Settings.save_data:
                self.save()
            self.plot()
            self.final()
            print(f"\nRuntime: {repr(self._runtime)[8:]}\n")

    return Frequency(
        config=config,
        qubits=qubits,
        t_max=t_max,
        detunings=detunings,
        subspace=subspace,
        n_elements=n_elements,
        params=params,
        **kwargs
    )


def Phase(
    qpu:            QPU,
    config:         Config,
    qubits:         List | Tuple,
    phases:         ArrayLike | NDArray | Dict[ArrayLike | NDArray],
    gate:           str = 'X90',
    subspace:       str = 'GE',
    relative_phase: bool = False,
    params:         Dict | None = None,
    **kwargs
) -> Callable:
    """Phase calibration for single-qubit gates.

    This calibration corrects for phases errors that occur due to stark shifts
    during single-qubit gates.

    Basic example useage for initial calibration:

    ```
    phases = np.pi * np.linspace(-1, 1, 21)
    cal = Phase(
        CustomQPU,
        config,
        qubits=[0, 1, 2],
        phases=phases)
    cal.run()
    ```

    Args:
        qpu (QPU): custom QPU object.
        config (Config): qcal Config object.
        qubits (List | Tuple): qubits to calibrate.
        phases (ArrayLike | NDArray | Dict[ArrayLike | NDArray]): array of
            phases to sweep over for calibrating the single-qubit gate
            phases. If calibrating multiple qubits at the same time, this
            should be a dictionary mapping an array to a qubit label.
        gate (str, optional): native gate to calibrate. Defaults to 'X90'.
        subspace (str, optional): qubit subspace for the defined gate.
            Defaults to 'GE'.
        relative_phase (bool, optional): whether phase sweep is relative to the
            current phase value. Defaults to False. If true, the phase sweep
            is added to the current value.
        params (Dict, None): dictionary mapping the qubit labels to the config
            parameter to sweep over.

    Returns:
        Callable: Phases calibration class.
    """

    class Phase(qpu, Calibration):
        """Phase calibration class.

        This class inherits a custom QPU from the Phase calibration
        function.
        """

        def __init__(
            self,
            config:         Config,
            qubits:         List | Tuple,
            phases:         ArrayLike | NDArray | Dict,
            gate:           str = 'X90',
            subspace:       str = 'GE',
            relative_phase: bool = False,
            params:         Dict | None = None,
            **kwargs
        ) -> None:
            """Initialize the Phase calibration class within the function."""
            n_levels = 3 if subspace == 'EF' else 2
            qpu.__init__(self, config=config, n_levels=n_levels, **kwargs)
            Calibration.__init__(self, config)

            self._qubits = qubits

            assert gate in ('X90', 'X'), (
                "'gate' must be one of 'X90' or 'X'!"
            )
            self._gate = gate

            assert subspace in ('GE', 'EF'), (
                "'subspace' must be one of 'GE' or 'EF'!"
            )
            self._subspace = subspace

            if not isinstance(phases, dict):
                self._phases = dict.fromkeys(qubits, phases)
            else:
                self._phases = phases

            if params:
                self._params = params
            else:
                self._params = {}
                for q in qubits:
                    pulse = f'single_qubit/{q}/{subspace}/{gate}/pulse/'
                    if gate == 'X90':
                        self._params[q] = [
                            pulse + '0/kwargs/phase', pulse + '2/kwargs/phase'
                        ]
                    elif gate == 'X':
                        self._params[q] = pulse + '1/kwargs/phase'

            if relative_phase:
                for q in qubits:
                    self._phases[q] = (
                        self._phases[q] + config[self._params[q][0]]
                    )
            self._param_sweep = self._phases

            if gate == 'X90':
                self._fit = {q: [FitLinear(), FitLinear()] for q in qubits}
            elif gate == 'X':
                self._fit = {q: FitCosine() for q in qubits}

        @property
        def phases(self) -> Dict:
            """Phase sweep for each qubit.

            Returns:
                Dict: qubit to array map.
            """
            return self._phases

        def generate_circuits(self):
            """Generate all amplitude calibration circuits."""
            logger.info(' Generating circuits...')

            level = {2: 'GE', 3: 'EF'}
            circuit_Y180_X90 = Circuit()
            circuit_X180_Y90 = Circuit()
            circuit_X90_X_X90 = Circuit()

            # Prepulse for EF calibration
            if self._subspace == 'EF':

                circuit_Y180_X90.extend([
                    Cycle({X90(q, subspace='GE') for q in self._qubits}),
                    Barrier(self._qubits),
                    Cycle({X90(q, subspace='GE') for q in self._qubits}),
                    Barrier(self._qubits)
                ])

                circuit_X180_Y90.extend([
                    Cycle({X90(q, subspace='GE') for q in self._qubits}),
                    Barrier(self._qubits),
                    Cycle({X90(q, subspace='GE') for q in self._qubits}),
                    Barrier(self._qubits)
                ])

                circuit_X90_X_X90.extend([
                    Cycle({X(q, subspace='GE') for q in self._qubits}),
                ])

            # Y180_X90
            circuit_Y180_X90.extend([
                Cycle({
                    Rz(q, np.pi/2, subspace=level[self._n_levels])
                    for q in self._qubits
                }),
                Cycle({
                    X90(q, subspace=level[self._n_levels])
                    for q in self._qubits
                }),
                Barrier(self._qubits),
                Cycle({
                    X90(q, subspace=level[self._n_levels])
                    for q in self._qubits
                }),
                Cycle({
                    Rz(q, -np.pi/2, subspace=level[self._n_levels])
                    for q in self._qubits
                }),
                Barrier(self._qubits),
                Cycle({
                    X90(q, subspace=level[self._n_levels])
                    for q in self._qubits
                })
            ])
            circuit_Y180_X90.measure()

            # X180_Y90
            circuit_X180_Y90.extend([
                Cycle({
                    X90(q, subspace=level[self._n_levels])
                    for q in self._qubits
                }),
                Barrier(self._qubits),
                Cycle({
                    X90(q, subspace=level[self._n_levels])
                    for q in self._qubits
                }),
                Barrier(self._qubits),
                Cycle({
                    Rz(q, np.pi/2, subspace=level[self._n_levels])
                    for q in self._qubits
                }),
                Cycle({
                    X90(q, subspace=level[self._n_levels])
                    for q in self._qubits
                }),
                Cycle({
                    Rz(q, -np.pi/2, subspace=level[self._n_levels])
                    for q in self._qubits
                }),
            ])
            circuit_X180_Y90.measure()

            # X90_X_X90
            circuit_X90_X_X90.extend([
                Cycle({
                    X90(q, subspace=level[self._n_levels])
                    for q in self._qubits
                }),
                Barrier(self._qubits),
                Cycle({
                    X(q, subspace=level[self._n_levels])
                    for q in self._qubits
                }),
                Barrier(self._qubits),
                Cycle({
                    X90(q, subspace=level[self._n_levels])
                    for q in self._qubits
                }),
            ])
            circuit_X90_X_X90.measure()

            circuits = []
            if self._gate == 'X90':
                for _ in range(self._phases[self._qubits[0]].size):
                    circuits.append(circuit_Y180_X90.copy())
                    circuits.append(circuit_X180_Y90.copy())

                self._circuits = CircuitSet(circuits=circuits)
                self._circuits['sequence'] = [
                        'Y180_X90', 'X180_Y90'
                    ] * int(self._circuits.n_circuits / 2)
                for q in self._qubits:
                    _phases = []
                    for p in self._phases[q]:
                        _phases.extend([p, p])
                    self._circuits[
                            f'param: {self._params[q][0]}'
                        ] = _phases
                    self._circuits[
                            f'param: {self._params[q][1]}'
                        ] = _phases

            elif self._gate == 'X':
                for _ in range(self._phases[self._qubits[0]].size):
                    circuits.append(circuit_X90_X_X90.copy())

                self._circuits = CircuitSet(circuits=circuits)
                self._circuits['sequence'] = (
                    ['X90_X_X90'] * int(self._circuits.n_circuits)
                )
                for q in self._qubits:
                    self._circuits[
                        f'param: {self._params[q]}'
                    ] = self._phases[q]

        def analyze(self) -> None:
            """Analyze the data."""
            logger.info(' Analyzing the data...')

            level = {'GE': '1', 'EF': '2'}
            for i, q in enumerate(self._qubits):
                if self._gate == 'X90':
                    pop0 = []  # circuit_Y180_X90
                    pop1 = []  # circuit_X180_Y90
                    pops = []
                    for j, circuit in enumerate(self._circuits):
                        pop = circuit.results.marginalize(i).populations[
                            level[self._subspace]
                        ]
                        pops.append(pop)

                        # Y180_X90
                        if j % 2 == 0:
                            pop0.append(pop)
                        # X180_Y90
                        elif j % 2 == 1:
                            pop1.append(pop)

                    self._circuits[
                        f'Q{q}: Prob({level[self._subspace]})'
                    ] = pops

                    self._sweep_results[q] = (
                        np.array(pop0) - np.array(pop1)
                    )**2

                    fit_df = {
                        "data": [pop0, pop1],
                        "params": [Parameters(), Parameters()],
                    }
                    for fit_idx in range(2): # For each of the two linear fits
                        params = fit_df["params"][fit_idx]
                        params.add('m', value=(-1.)**(fit_idx))
                        params.add('b', value=0.5)
                        self._fit[q][fit_idx].fit(
                            self._phases[q],
                            fit_df["data"][fit_idx],
                            params=params
                        )

                    # If the fit was successful, write find the new phase
                    if (
                        self._fit[q][0].fit_success and
                        self._fit[q][1].fit_success
                    ):
                        m0 = self._fit[q][0].fit_params['m'].value
                        b0 = self._fit[q][0].fit_params['b'].value
                        m1 = self._fit[q][1].fit_params['m'].value
                        b1 = self._fit[q][1].fit_params['b'].value

                        # Compute x value where the two lines cross
                        newvalue = (b1 - b0) / (m0 - m1)
                        if not in_range(newvalue, self._phases[q]):
                            logger.warning(
                                f'Fit failed for qubit {q} (out of range)!'
                            )
                            self._fit[q][0]._fit_success = False
                            self._fit[q][1]._fit_success = False
                            self._cal_values[q] = [
                                self._phases[q][
                                    np.argmin(self._sweep_results[q])
                                ]
                            ] * 2
                        elif (
                            self._fit[q][0].error > 10 or
                            self._fit[q][1].error > 10
                        ):
                            logger.warning(
                                f'Fit failed for qubit {q} '
                                '(underfitting based on reduced chi2>>1)!'
                            )
                            self._cal_values[q] = [
                                self._phases[q][
                                    np.argmin(self._sweep_results[q])
                                ]
                            ] * 2
                        else:
                            self._cal_values[q] = [newvalue] * 2

                    else:
                        self._cal_values[q] = [
                            self._phases[q][
                                np.argmin(self._sweep_results[q])
                            ]
                        ] * 2

                elif self._gate == 'X':
                    pop0 = []
                    pop1 = []
                    for circuit in self._circuits:
                        pop0.append(
                            circuit.results.marginalize(i).populations[
                                str(int(level[self._subspace]) - 1)
                            ]
                        )
                        pop1.append(
                            circuit.results.marginalize(i).populations[
                                level[self._subspace]
                            ]
                        )

                    self._circuits[
                        f'Q{q}: Prob({int(level[self._subspace]) - 1})'
                    ] = pop0
                    self._sweep_results[q] = np.array(pop0) - np.array(pop1)

                    est_freq = est_freq_fft(
                        self._phases[q], self._sweep_results[q]
                    )
                    params = Parameters()
                    params.add('amp', value=1.0, min=0.75, max=1.)
                    params.add('freq', value=est_freq, min=0.)
                    params.add('phase', value=0., min=-np.pi, max=np.pi)
                    params.add('offset', value=0., min=-0.1, max=0.1)
                    self._fit[q].fit(
                        self._phases[q], self._sweep_results[q], params=params
                    )

                    if self._fit[q].fit_success:
                        freq = self._fit[q].fit_params['freq'].value
                        phase = self._fit[q].fit_params['phase'].value
                        self._cal_values[q] = wrap_phase(
                            -phase / (2 * np.pi * freq)
                        )
                    else:
                        self._cal_values[q] = self._phases[q][
                            np.argmax(self._sweep_results[q])
                        ]

        def save(self):
            """Save all circuits and data."""
            clear_output(wait=True)
            self._data_manager._exp_id += (
                f'_phase_cal{"".join("Q" + str(q) for q in self._qubits)}'
            )
            if Settings.save_data:
                qpu.save(self)
                self._data_manager.save_to_csv(
                    pd.DataFrame([self._sweep_results]), 'sweep_results'
                )
                self._data_manager.save_to_csv(
                    pd.DataFrame([self._cal_values]), 'calibrated_values'
                )

        def plot(self) -> None:
            """Plot the phase sweep and fit results."""
            level = {'GE': '1', 'EF': '2'}

            nrows, ncols = calculate_nrows_ncols(len(self._qubits))
            figsize = (5 * ncols, 4 * nrows)
            fig, axes = plt.subplots(
                nrows, ncols, figsize=figsize, layout='constrained'
            )

            k = -1
            for i in range(nrows):
                for j in range(ncols):
                    k += 1

                    if len(self._qubits) == 1:
                        ax = axes
                    elif axes.ndim == 1:
                        ax = axes[j]
                    elif axes.ndim == 2:
                        ax = axes[i,j]

                    if k < len(self._qubits):
                        q = self._qubits[k]

                        ax.set_xlabel('Phase (rad.)', fontsize=15)
                        if self._gate == 'X90':
                            ax.set_ylabel(
                                r'$|2\rangle$ Population' if
                                self._subspace == 'EF' else
                                r'$|1\rangle$ Population',
                                fontsize=15
                            )
                        elif self._gate == 'X':
                            ax.set_ylabel(r'$\langle Z \rangle$', fontsize=15)
                        ax.tick_params(
                            axis='both', which='major', labelsize=12
                        )
                        ax.grid(True)

                        # Plot data
                        if self._gate == 'X90':
                            ax.plot(
                                self._param_sweep[q],
                                self._circuits[
                                    self._circuits['sequence'] == 'Y180_X90'
                                ][f'Q{q}: Prob({level[self._subspace]})'],
                                '-o', c='blue', label='Y180_X90'
                            )
                            ax.plot(
                                self._param_sweep[q],
                                self._circuits[
                                    self._circuits['sequence'] == 'X180_Y90'
                                ][f'Q{q}: Prob({level[self._subspace]})'],
                                '-o', c='blueviolet', label='X180_Y90'
                            )
                        ax.plot(
                            self._param_sweep[q], self._sweep_results[q],
                            'o', c='k', label=f'Meas, Q{q}'
                        )

                        # Plot fitted curves
                        x = np.linspace(
                            self._param_sweep[q][0],
                            self._param_sweep[q][-1],
                            100
                        )
                        if self._gate == 'X' and self._fit[q].fit_success:
                           ax.plot(
                                x, self._fit[q].predict(x),
                                '-', c='orange', label='Fit'
                            )

                           if self._cal_values[q]:
                                ax.axvline(
                                    self._cal_values[q],
                                    ls='--', c='k', label='Fit value'
                                )

                        elif self._gate == 'X90':
                            for fit_idx in range(2):
                                if self._fit[q][fit_idx].fit_success:
                                    ax.plot(
                                    x, self._fit[q][fit_idx].predict(x),
                                    '-', c=["blue", "blueviolet"][fit_idx]
                                )

                            if (
                                self._fit[q][0].fit_success and
                                self._fit[q][1].fit_success
                            ):
                                ax.plot(
                                    x,
                                    np.subtract(
                                        self._fit[q][0].predict(x),
                                        self._fit[q][1].predict(x)
                                    )**2,
                                    '-', c='orange'
                                )

                            if self._cal_values[q]:
                                ax.axvline(
                                    self._cal_values[q][0],
                                    ls='--', c='k', label='Fit value'
                                )

                        ax.legend(loc=0, fontsize=12)

                    else:
                        ax.axis('off')

            fig.set_tight_layout(True)
            if Settings.save_data:
                fig.savefig(
                    self._data_manager._save_path + 'calibration_results.png',
                    dpi=300
                )
                fig.savefig(
                    self._data_manager._save_path + 'calibration_results.pdf'
                )
                fig.savefig(
                    self._data_manager._save_path + 'calibration_results.svg'
                )
            plt.show()

        def final(self) -> None:
            """Final experimental method."""
            Calibration.final(self)
            print(f"\nRuntime: {repr(self._runtime)[8:]}\n")

        def run(self):
            """Run all experimental methods and analyze results."""
            self.generate_circuits()
            qpu.run(self, self._circuits, save=False)
            self.analyze()
            self.save()
            self.plot()
            self.final()

    return Phase(
        config=config,
        qubits=qubits,
        phases=phases,
        gate=gate,
        subspace=subspace,
        relative_phase=relative_phase,
        params=params,
        **kwargs
    )
