"""Class definition of the Generator Drive model."""

import jax.numpy as jnp

from paraqeet.model.drive import Drive
from paraqeet.quantity import Array, Quantity
from paraqeet.signal.generator import Generator


class DriveOperator(Drive):
    r"""Create a generator drive model.

    Transversal (a^\\dagger + a or \\sigma_x type) or longitudinal (a^\\dagger a or \\sigma_z type) drive with a
    time-dependent scalar coefficient that is generated by a Generator object.

    Parameters
    ----------
    generator : Generator
        Signal generator stack.
    is_longitudinal : bool
        Generator is longitudinal or transversal depending on this boolean.

    """

    _generator: Generator
    _is_longitudinal: bool

    def __init__(self, generator: Generator, is_longitudinal: bool):
        self._generator = generator
        self._is_longitudinal = is_longitudinal

    @property
    def generator(self) -> Generator:
        """Get the signal generator from the system.

        Returns
        -------
        paraqeet.signal.generator.Generator
            Returns the signal generator object from the system.

        """
        return self._generator

    def get_parameters(self) -> list[Quantity]:
        """Get a list of parameters of the system.

        Returns
        -------
        list[Quantity]
            List of optimizable parameters of the system.

        """
        return self._generator.get_parameters()

    def _compute_matrix(self, a: Array) -> Array:
        """Return the operator for the longitudinal or transverse drive.

        Parameters
        ----------
        a: Array
            Operator for computation.

        Returns
        -------
        Array
            Returns the operator for the longitudinal or transverse drive.

        """
        return (jnp.conjugate(a.T) @ a) if self._is_longitudinal else (jnp.conjugate(a.T) + a)

    def get_value_at_timestep(self, a: Array, t: float) -> Array:
        """Get the one-time matrix of the system.

        Fetches the coefficient from the drive and transforms it
        into the correct shape for the Hamiltonian.

        Parameters
        ----------
        a: Array
            Operator for longitudinal or transverse drive.
        t: Array
            One-dimensional vector of timestamps.

        Returns
        -------
        Array
            Returns the shape-shifted coefficient from the drive.

        """
        # TODO: generator.get_value expects an array, even for one time point.
        # Is the naming of the method correct then?
        signal = self._generator.get_value(t)
        matrix = self._compute_matrix(a)
        return signal * matrix

    def get_gradient_at_timestep(self, a: Array, timestep: float) -> Array:
        """Get the one-time gradient of the system.

        Fetches the gradient from the drive and transforms it into the
        correct shape for the Hamiltonian.

        Parameters
        ----------
        a: Array
            Operator for longitudinal or transverse drive.
        timestep: Array
            One-dimensional vector of timestamps.

        Returns
        -------
        Array
            Returns the shape-shifted gradient from the drive.

        """
        signal_grad = self._generator.get_gradient_at_timestep(timestep).reshape((-1, 1, 1))
        matrix = Drive._repeat(self._compute_matrix(a), signal_grad.shape[0])
        return signal_grad * matrix
