# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file 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.

from __future__ import annotations

from copy import deepcopy
from typing import Any, Dict, List, Set, Union

from openpulse import ast
from oqpy import BitVar, Program
from oqpy.timing import OQDurationLiteral

from braket.parametric.free_parameter import FreeParameter
from braket.parametric.free_parameter_expression import FreeParameterExpression
from braket.parametric.parameterizable import Parameterizable
from braket.pulse.ast.approximation_parser import _ApproximationParser
from braket.pulse.ast.free_parameters import (
    _FreeParameterExpressionIdentifier,
    _FreeParameterTransformer,
)
from braket.pulse.ast.qasm_parser import ast_to_qasm
from braket.pulse.ast.qasm_transformer import _IRQASMTransformer
from braket.pulse.frame import Frame
from braket.pulse.pulse_sequence_trace import PulseSequenceTrace
from braket.pulse.waveforms import Waveform


class PulseSequence:
    """
    A representation of a collection of instructions to be performed on a quantum device
    and the requested results.
    """

    def __init__(self):
        self._capture_v0_count = 0
        self._program = Program()
        self._frames = {}
        self._waveforms = {}
        self._free_parameters = set()

    def to_time_trace(self) -> PulseSequenceTrace:
        """Generate an approximate trace of the amplitude, frequency, phase for each frame
        contained in the PulseSequence, under the action of the instructions contained in
        the pulse sequence.

        Returns:
            PulseSequenceTrace: The approximation information with each attribute
            (amplitude, frequency and phase) mapping a str (frame id) to a TimeSeries
            (containing the time evolution of that attribute).
        """
        parser = _ApproximationParser(deepcopy(self._program), self._frames)
        return PulseSequenceTrace(
            amplitudes=parser.amplitudes, frequencies=parser.frequencies, phases=parser.phases
        )

    @property
    def parameters(self) -> Set[FreeParameter]:
        """Returns the set of `FreeParameter` s in the PulseSequence."""
        return self._free_parameters.copy()

    def set_frequency(
        self, frame: Frame, frequency: Union[float, FreeParameterExpression]
    ) -> PulseSequence:
        """
        Adds an instruction to set the frequency of the frame to the specified `frequency` value.

        Args:
            frame (Frame): Frame for which the frequency needs to be set.
            frequency (Union[float, FreeParameterExpression]): frequency value to set
                for the specified frame.

        Returns:
            PulseSequence: self, with the instruction added.
        """

        _validate_uniqueness(self._frames, frame)
        self._program.set_frequency(frame=frame, freq=self._format_parameter_ast(frequency))
        self._frames[frame.id] = frame
        return self

    def shift_frequency(
        self, frame: Frame, frequency: Union[float, FreeParameterExpression]
    ) -> PulseSequence:
        """
        Adds an instruction to shift the frequency of the frame by the specified `frequency` value.

        Args:
            frame (Frame): Frame for which the frequency needs to be shifted.
            frequency (Union[float, FreeParameterExpression]): frequency value by which to shift
                the frequency for the specified frame.

        Returns:
            PulseSequence: self, with the instruction added.
        """
        _validate_uniqueness(self._frames, frame)
        self._program.shift_frequency(frame=frame, freq=self._format_parameter_ast(frequency))
        self._frames[frame.id] = frame
        return self

    def set_phase(
        self, frame: Frame, phase: Union[float, FreeParameterExpression]
    ) -> PulseSequence:
        """
        Adds an instruction to set the phase of the frame to the specified `phase` value.

        Args:
            frame (Frame): Frame for which the frequency needs to be set.
            phase (Union[float, FreeParameterExpression]): phase value to set
                for the specified frame.

        Returns:
            PulseSequence: self, with the instruction added.
        """
        _validate_uniqueness(self._frames, frame)
        self._program.set_phase(frame=frame, phase=self._format_parameter_ast(phase))
        self._frames[frame.id] = frame
        return self

    def shift_phase(
        self, frame: Frame, phase: Union[float, FreeParameterExpression]
    ) -> PulseSequence:
        """
        Adds an instruction to shift the phase of the frame by the specified `phase` value.

        Args:
            frame (Frame): Frame for which the phase needs to be shifted.
            phase (Union[float, FreeParameterExpression]): phase value by which to shift
                the phase for the specified frame.

        Returns:
            PulseSequence: self, with the instruction added.
        """
        _validate_uniqueness(self._frames, frame)
        self._program.shift_phase(frame=frame, phase=self._format_parameter_ast(phase))
        self._frames[frame.id] = frame
        return self

    def set_scale(
        self, frame: Frame, scale: Union[float, FreeParameterExpression]
    ) -> PulseSequence:
        """
        Adds an instruction to set the scale on the frame to the specified `scale` value.

        Args:
            frame (Frame): Frame for which the scale needs to be set.
            scale (Union[float, FreeParameterExpression]): scale value to set
                on the specified frame.

        Returns:
            PulseSequence: self, with the instruction added.
        """
        _validate_uniqueness(self._frames, frame)
        self._program.set_scale(frame=frame, scale=self._format_parameter_ast(scale))
        self._frames[frame.id] = frame
        return self

    def delay(
        self, frames: Union[Frame, List[Frame]], duration: Union[float, FreeParameterExpression]
    ) -> PulseSequence:
        """
        Adds an instruction to advance the frame clock by the specified `duration` value.

        Args:
            frames (Union[Frame, List[Frame]]): Frame(s) on which the delay needs to be introduced.
            duration (Union[float, FreeParameterExpression]): value (in seconds) defining
                the duration of the delay.

        Returns:
            PulseSequence: self, with the instruction added.
        """
        if not isinstance(frames, list):
            frames = [frames]
        if isinstance(duration, FreeParameterExpression):
            for p in duration.expression.free_symbols:
                self._free_parameters.add(FreeParameter(p.name))
            duration = OQDurationLiteral(duration)
        _validate_uniqueness(self._frames, frames)
        self._program.delay(time=duration, qubits_or_frames=frames)
        for frame in frames:
            self._frames[frame.id] = frame
        return self

    def barrier(self, frames: List[Frame]) -> PulseSequence:
        """
        Adds an instruction to align the frame clocks to the latest time across all the specified
        frames.

        Args:
            frames (List[Frame]): Frames across which the frame clocks need to be aligned.

        Returns:
            PulseSequence: self, with the instruction added.
        """
        _validate_uniqueness(self._frames, frames)
        self._program.barrier(qubits_or_frames=frames)
        for frame in frames:
            self._frames[frame.id] = frame
        return self

    def play(self, frame: Frame, waveform: Waveform) -> PulseSequence:
        """Adds an instruction to play the specified waveform on the supplied frame.

        Args:
            frame (Frame): Frame on which the specified waveform signal would be output.
            waveform (Waveform): Waveform envelope specifying the signal to output on the
                specified frame.

        Returns:
            PulseSequence: returns self.
        """
        _validate_uniqueness(self._frames, frame)
        _validate_uniqueness(self._waveforms, waveform)
        self._program.play(frame=frame, waveform=waveform)
        if isinstance(waveform, Parameterizable):
            for param in waveform.parameters:
                if isinstance(param, FreeParameterExpression):
                    for p in param.expression.free_symbols:
                        self._free_parameters.add(FreeParameter(p.name))
        self._frames[frame.id] = frame
        self._waveforms[waveform.id] = waveform
        return self

    def capture_v0(self, frame: Frame) -> PulseSequence:
        """
        Adds an instruction to capture the bit output from measuring the specified frame.

        Args:
            frame (Frame): Frame on which the capture operation needs
                to be performed.

        Returns:
            PulseSequence: self, with the instruction added.
        """
        _validate_uniqueness(self._frames, frame)
        self._program.function_call("capture_v0", [frame])
        self._capture_v0_count += 1
        self._frames[frame.id] = frame
        return self

    def make_bound_pulse_sequence(self, param_values: Dict[str, float]) -> PulseSequence:
        """
        Binds FreeParameters based upon their name and values passed in. If parameters
        share the same name, all the parameters of that name will be set to the mapped value.

        Args:
            param_values (Dict[str, float]):  A mapping of FreeParameter names
                to a value to assign to them.

        Returns:
            PulseSequence: Returns a PulseSequence with all present parameters fixed to
            their respective values.
        """
        program = deepcopy(self._program)
        tree: ast.Program = program.to_ast(include_externs=False, ignore_needs_declaration=True)
        new_tree: ast.Program = _FreeParameterTransformer(param_values).visit(tree)

        new_program = Program()
        new_program.declared_vars = program.declared_vars
        new_program.undeclared_vars = program.undeclared_vars
        for x in new_tree.statements:
            new_program._add_statement(x)

        new_pulse_sequence = PulseSequence()
        new_pulse_sequence._program = new_program
        new_pulse_sequence._frames = deepcopy(self._frames)
        new_pulse_sequence._waveforms = {
            wf.id: wf.bind_values(**param_values) if isinstance(wf, Parameterizable) else wf
            for wf in deepcopy(self._waveforms).values()
        }

        # Update waveforms to bind values
        for v in new_program.undeclared_vars:
            if v in self._waveforms:
                new_program.undeclared_vars[v] = new_pulse_sequence._waveforms[
                    v
                ]._to_oqpy_expression()

        new_pulse_sequence._capture_v0_count = self._capture_v0_count
        new_pulse_sequence._free_parameters = set(
            [p for p in self._free_parameters if p.name not in param_values]
        )

        return new_pulse_sequence

    def to_ir(self) -> str:
        """Converts this OpenPulse problem into IR representation.

        Returns:
            str: a str representing the OpenPulse program encoding the PulseSequence.
        """
        program = deepcopy(self._program)
        if self._capture_v0_count:
            register_identifier = "psb"
            program.declare(
                BitVar[self._capture_v0_count](name=register_identifier), to_beginning=True
            )
            tree = program.to_ast(encal=True, include_externs=False)
            tree = _IRQASMTransformer(register_identifier).visit(tree)
        else:
            tree = program.to_ast(encal=True, include_externs=False)
        return ast_to_qasm(tree)

    def _format_parameter_ast(
        self, parameter: Union[float, FreeParameterExpression]
    ) -> Union[float, _FreeParameterExpressionIdentifier]:
        if isinstance(parameter, FreeParameterExpression):
            for p in parameter.expression.free_symbols:
                self._free_parameters.add(FreeParameter(p.name))
            return _FreeParameterExpressionIdentifier(parameter)
        return parameter

    def __call__(self, arg: Any = None, **kwargs) -> PulseSequence:
        """
        Implements the call function to easily make a bound PulseSequence.

        Args:
            arg (Any): A value to bind to all parameters. Defaults to None and
                can be overridden if the parameter is in kwargs.

        Returns:
            PulseSequence: A pulse sequence with the specified parameters bound.
        """
        param_values = dict()
        if arg is not None:
            for param in self.parameters:
                param_values[str(param)] = arg
        for key, val in kwargs.items():
            param_values[str(key)] = val
        return self.make_bound_pulse_sequence(param_values)


def _validate_uniqueness(
    mapping: Dict[str, Any], values: Union[Frame, Waveform, List[Frame], List[Waveform]]
) -> None:
    if not isinstance(values, list):
        values = [values]

    for value in values:
        if value.id in mapping and mapping[value.id] != value:
            raise ValueError(
                f"{value.id} has already been used for defining {mapping[value.id]} "
                f"which differs from {value}"
            )
