# Copyright 2026 Q-CTRL. All rights reserved.
#
# Licensed under the Q-CTRL Terms of service (the "License"). Unauthorized
# copying or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. You may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
#    https://q-ctrl.com/terms
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS. See the
# License for the specific language.

from __future__ import annotations

from functools import partial

__all__ = ("SyncQbloxSystem",)

import contextlib
import logging
import weakref
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

import numpy as np
from boulderopalscaleupsdk.device.controller.qblox import (
    OutputSequencerAcquisitions,
    PreparedProgram,
    SequencerAddr,
    process_sequencer_output,
)
from boulderopalscaleupsdk.errors import ScaleUpClientError

from boulderopalscaleup.qblox import arm_sequencers, execute_armed_sequencers
from boulderopalscaleup.runtime._sync.executable import ExecUnit, NoResourceProxy
from boulderopalscaleup.runtime.abstract import ExecutionMode
from boulderopalscaleup.runtime.systems.qblox import (
    BaseQbloxSystem,
    ConcatenateData,
    PreparedProgramUnitData,
    QbloxExecutionContext,
    QbloxRoutines,
    ReshapeData,
)

if TYPE_CHECKING:
    from collections.abc import Iterator

    from qblox_instruments import Cluster

LOGGER = logging.getLogger(__name__)

RawT = dict[SequencerAddr, OutputSequencerAcquisitions]


@dataclass
class _ClusterProxy:
    # Since the executable will contained bound units that hold onto these proxies, we want to
    # ensure a user cannot unintentionally extend the lifetime of the parent system by holding onto
    # the executable!
    _sys: weakref.ReferenceType[SyncQbloxSystem]

    @contextlib.contextmanager
    def get(self) -> Iterator[dict[str, Cluster]]:
        sys = self._sys()
        if sys is None:  # pragma: no cover
            raise RuntimeError("Resource proxy for QBLOX cluster is out-of-bounds.")
        if sys._stack is None:  # system will unreference the stack when it is closed
            raise RuntimeError("QBLOX system has closed.")
        yield sys._stack


class SyncQbloxSystem(BaseQbloxSystem[ExecUnit], execution_mode=ExecutionMode.SYNC):
    def get_resource_proxy(self) -> _ClusterProxy:
        return _ClusterProxy(weakref.ref(self))

    def make_unit(self, node_id: str, data: Any, routine: QbloxRoutines) -> ExecUnit:
        partial_unit = partial(ExecUnit, node_id, data, self.get_resource_proxy(), self._context)
        match routine:
            case QbloxRoutines.RUN_PROGRAM_RAW:
                return partial_unit(_exec_program_raw, tracks_duration=True)
            case QbloxRoutines.RUN_PROGRAM_PROCESSED:
                return partial_unit(_exec_program_processed, tracks_duration=True)
            case QbloxRoutines.CONCAT:
                return ExecUnit(
                    node_id,
                    data,
                    NoResourceProxy(),
                    self._context,
                    _run_concatenate,
                )
            case QbloxRoutines.RESHAPE:
                return ExecUnit(
                    node_id,
                    data,
                    NoResourceProxy(),
                    self._context,
                    _run_reshape,
                )
            case QbloxRoutines.MIXER_CALIBRATION:
                return partial_unit(lambda data, _, ctx: {})  # noqa: ARG005
            case _:  # pragma: no cover
                raise ValueError(f"Unsupported routine '{routine}' for QBLOX system.")


def _exec_program_raw(
    unit_data: PreparedProgramUnitData,
    cluster: dict[str, Cluster],
    ctx: QbloxExecutionContext,
) -> dict[SequencerAddr, OutputSequencerAcquisitions]:
    for ch, psp in unit_data.prepared_program.sequence_programs.items():
        LOGGER.debug(
            "Running program for '%s' on ch_out=%s\n %s",
            psp.ch_out,
            ch,
            psp.sequence_program.program,
        )

    # TODO: We could refactor the lo calibration and sideband calibration to be separate nodes
    # for better modularity. Preserving the current structure for simplicity
    armed = arm_sequencers(
        unit_data.prepared_program,
        cluster,
        calibrate_elements=unit_data.calibrate_elements,
    )
    return execute_armed_sequencers(
        armed,
        timeout=ctx.timeout,
        timeout_poll_res=ctx.timeout_poll_res,
    )


def _run_concatenate(
    data: ConcatenateData,
    _: None,
    ctx: QbloxExecutionContext,  # noqa: ARG001
) -> dict[str, list[float]]:
    """Concatenate raw results from dependency units.

    Results are injected via set_results() by SyncExecutable before execution.
    """
    return _concatenate_raw_results(
        {kk: vv.get() for kk, vv in data.dep_proxies.items()},
        data.dep_programs,
    )


def _run_reshape(
    data: ReshapeData,
    _: None,
    ctx: QbloxExecutionContext,  # noqa: ARG001
) -> dict[str, list[float]]:
    """Reshape raw results from dependency units into the target result_shape."""
    flat = _concatenate_raw_results(
        {kk: vv.get() for kk, vv in data.dep_proxies.items()},
        data.dep_programs,
    )
    return {key: np.reshape(values, data.result_shape).tolist() for key, values in flat.items()}


def _concatenate_raw_results(
    raw_results: dict[str, RawT],
    programs: dict[str, PreparedProgram],
) -> dict[str, list[float]]:
    combined_bins: dict[str, list[np.ndarray]] = {}

    result_types = []
    # QBLOXController._flatten_sequencer_results
    for node_id, raw in raw_results.items():
        program = programs[node_id]
        for seq_addr, seq_output in raw.items():
            seq_prog = program.get_sequencer_program(seq_addr)
            seq_results = process_sequencer_output(
                seq_prog,
                seq_output,
                data_type=int if seq_prog.result_type == "int" else complex,
            )
            result_types.append(seq_results.data_type)
            for result_key, result_bins in seq_results.bins.items():
                if result_key not in combined_bins:
                    combined_bins[result_key] = []
                combined_bins[result_key].append(result_bins)

    if len(set(result_types)) != 1:
        raise ScaleUpClientError(
            "Different acquisition modes present on different sequencers. Not supported",
        )

    data_type = result_types[0]
    # QBLOXController._results_post_process
    ret: dict[str, list[float]] = {}
    if data_type is int:
        for result_key, bins_list in combined_bins.items():
            ret[result_key] = np.concatenate(bins_list).tolist()
    else:
        for result_key, bins_list in combined_bins.items():
            concatenated = np.concatenate(bins_list)
            ret[f"{result_key}_i"] = np.real(concatenated).tolist()
            ret[f"{result_key}_q"] = np.imag(concatenated).tolist()

    return ret


def _exec_program_processed(
    unit_data: PreparedProgramUnitData,
    cluster: dict[str, Cluster],
    ctx: QbloxExecutionContext,
) -> dict[str, list[float]]:
    raw = _exec_program_raw(unit_data, cluster, ctx)
    bins: dict[str, np.ndarray] = {}

    result_types = []

    for seq_addr, seq_output in raw.items():
        seq_prog = unit_data.prepared_program.get_sequencer_program(seq_addr)
        seq_results = process_sequencer_output(
            seq_prog,
            seq_output,
            data_type=int if seq_prog.result_type == "int" else complex,
        )
        result_types.append(seq_results.data_type)

        bins |= seq_results.bins

    ret = {}

    if len(set(result_types)) != 1:
        raise ScaleUpClientError(
            "Different acquisition modes present on different sequencers. Not supported",
        )

    data_type = result_types[0]

    if data_type is int:
        for result_key, result_bins in bins.items():
            ret[result_key] = result_bins.tolist()
    else:
        for result_key, result_bins in bins.items():
            ret[f"{result_key}_i"] = np.real(result_bins).tolist()
            ret[f"{result_key}_q"] = np.imag(result_bins).tolist()

    return ret
