Complete experiments.
Built to evolve.

from __future__ import annotations

import typing as t
import holoviews as hv
import hvplot.xarray
import numpy as np
import numpy.typing as npt
from qm import qua

from data_analyzer.data_formatting.parameter_converter import unit
from data_analyzer.data_formatting.xarray_formating import without_timestamps
from data_analyzer.models.basic import Oscillation
from opx_control.experiment_frames.experiment_base import (
    ExperimentBase,
    OneOrIterable,
    log_status,
    mirror_doc,
    to_1d,
)
from opx_control.qua_macros.as_qua_program import as_qua_program
from opx_control.quam.components.qubits import FluxLine, ScQubit

if t.TYPE_CHECKING:
    from opx_control.fetching.streams import StreamManager
    from qm.program import Program
    import xarray as xr

__all__ = ["RabiAmplitude"]
u = unit(coerce_to_integer=True)


@log_status(status="in progress", auther="Malthe Marciniak", version="1.0.0")
class RabiAmplitude(ExperimentBase):
    models: dict[str, Oscillation]
    qubits: list[ScQubit]

    def program(qubits, amplitudes = np.linspace(0.05 * u.V, 0.4 * u.V, 101)) -> Program:
        sts, axt = StreamManager(), AxisTracker()
      
        n, a = qua.declare(int), qua.declare(qua.fixed)

        with axt.for_(n, 0, n < number_of_averages, n + 1, name="averages"):
            with axt.for_each_(a, amplitudes, name="amplitude", long_name="Amplitude Scale"):
                for qubit in self.qubits:

                    # Reset the qubit.
                    qubit.reset()

                    # Play the pi pulse at the current amplitude scale.
                    qubit.xy.play("x180", amplitude_scale=a)
                    qubit.align()

                    # Measure and save both readout quadratures.
                    I, Q = qubit.measure("readout")
                    sts[qubit.name].save(I, Q)

    def set_temporary_changes_for_experiment(
        self, qubits: OneOrIterable[ScQubit], *args, **kwargs
    ):
        self.qubits = to_1d(qubits)
        for qubit in self.qubits:
            qubit.xy.operations["x180"].set_at_reference(
                "amplitude", 0.5, allow_non_reference=True
            )

    def analyze(qubits, data: xr.Dataset) -> dict[str, xr.Dataset]:
        return {
            qubit.name: Oscillation().fit(data=data[qubit.name]) 
            for qubit in qubits
            }

    def plot(qubits, data: xr.Dataset, models, *args, **kwargs) -> hv.Layout:
        plots = []
        for qubit in qubits:

            # Plot raw data
            plot = data[qubit.name].hvplot()

            # Overlay fit
            x = data[qubit.name]["amplitude"]
            y = models[qubit.name].eval(x)
            plot *= hv.Curve((x, y))

            # Add to plot list
            plots.append(plot)

        return hv.Layout(plots).cols(1)

    def update_machine(qubits, models):
        for qubit in qubits:

            # Get the fitted results
            fit_results = models[qubit.name].xarray_results()
            period = fit_results["period"]

            # Update the qubit's XY drive amplitudes.
            for pulse_90 in ("x90", "y90"):
                qubit.xy.operations[pulse_90].amplitude = "amplitude", 0.5 * period / 4

            for pulse_180 in ("x180", "y180"):
                qubit.xy.operations[pulse_180].amplitude = 0.5 * period / 2

    @mirror_doc(program)
    def run(self, *args, **kwargs):
        return super().run(*args, **kwargs)