from collections import UserList

import numpy as np
import scipy.linalg
from pycce.utilities import expand


class Pulse:
    """
    Class containing properties of each control pulse, applied to the system.

    Args:
        axis (str): Axis of rotation of the central spin. Can be 'x', 'y', or 'z'. Default is None.
        angle (float or str): Angle of rotation of central spin. Can be provided in rad, or as a string, containing
            fraction of pi: ``'pi'``, ``'pi/2'``, ``'2*pi'`` etc. Default is None.
        delay (float or ndarray): Delay before the pulse or array of delays with the same shape as time points.
            Default is None.
        bath_names (str or array-like of str): Name or array of names of bath spin types, impacted by the bath pulse.
            Default is None.
        bath_axes (str or array-like of str): Axis of rotation or array of axes of the bath spins.
            Default is None. If ``bath_names`` is provided, but ``bath_axes`` and ``bath_angles`` are not,
            assumes the same axis and angle as the one of the central spin
        bath_angles (float or str or array-like): Angle of rotation or array of axes of rotations of the bath spins.

    Examples:

        >>> Pulse('x', 'pi')
        Pulse((x, 3.14))
        >>> Pulse('x', 'pi', bath_names=['13C', '14C'])
        Pulse((x, 3.14), {13C: (x, 3.14), 14C: (x, 3.14)})
        >>> import numpy as np
        >>> Pulse('x', 'pi', delay=np.linspace(0, 1, 5), bath_names=['13C', '14C'], bath_axes='x', bath_angles='pi/2')
        Pulse((x, 3.14), {13C: (x, 1.57), 14C: (x, 1.57)}, t = [0.   0.25 0.5  0.75 1.  ])

    """

    def __init__(self, axis=None, angle=None, delay=None, bath_names=None, bath_axes=None, bath_angles=None):
        pi = np.pi

        self.axis = axis
        """str: Axis of rotation of the central spin"""
        if isinstance(angle, str):
            angle = eval(angle)

        self.angle = angle
        """float: Angle of rotation of central spin in rad."""
        self.flip = False
        """bool: True if the angle == pi."""
        if angle is not None and np.isclose(angle, np.pi):
            self.flip = True

        self._has_delay = False
        self.delay = delay

        if bath_names is not None:
            bath_names = np.asarray(bath_names)

        self.bath_names = bath_names
        """ndarray: Array of names of bath spin types, impacted by the bath pulse."""
        if bath_axes is not None:
            bath_axes = np.asarray(bath_axes)
        elif bath_names is not None:
            bath_axes = np.asarray(axis)

        self.bath_axes = bath_axes
        """ndarray: Array of axes of rotation of the bath spins."""

        if bath_angles is not None:
            bath_angles = np.asarray(bath_angles)
            if bath_angles.dtype.char in 'SU':
                if bath_angles.shape:
                    bath_angles = np.asarray([eval(a) for a in bath_angles])
                else:
                    bath_angles = np.asarray(eval(bath_angles[()]))

        elif bath_names is not None:
            bath_angles = np.asarray(angle)

        self.bath_angles = bath_angles
        """ndarray: Array of angles of rotation of the bath spins."""

        self.rotation = None
        """ndarray: Matrix representation of the pulse for the given cluster. Generated by ``Sequence`` object."""

    @property
    def delay(self):
        """float or ndarray:  Delay before the pulse or array of delays with the same shape as time points."""
        return self._delay

    @delay.setter
    def delay(self, value):
        self._delay = value
        if value is not None:
            self._delay = np.asarray(value)
            self._has_delay = True
        else:
            self._has_delay = False
    def __repr__(self):

        w = f'Pulse('
        inner_message = ''
        if self.axis is not None:
            try:
                inner_message += f'({self.axis}, {self.angle:.2f})'
            except TypeError:
                inner_message += f'({self.axis}, {self.angle})'

        if self.bath_names is not None:
            if inner_message:
                inner_message += ', {'
            properties = np.broadcast(self.bath_names, self.bath_axes, self.bath_angles)
            for name, axis, angle in properties:
                try:
                    inner_message += f'{name}: ({axis}, {angle:.2f}), '
                except TypeError:
                    inner_message += f'{name}: ({axis}, {angle}), '
            inner_message = inner_message[:-2] + '}'

        if self.delay is not None:
            if inner_message:
                inner_message += ', '
            inner_message += f't = {self.delay}'
        w += inner_message + ')'
        return w


_rot = {'x': 0, 'y': 1, 'z': 2}


def _get_pulse(value):
    if not isinstance(value, Pulse):
        try:
            value = Pulse(**value)
        except TypeError:
            value = Pulse(*value)
    return value


class Sequence(UserList):
    """
    List-like object, which contains the sequence of the pulses.

    Each element is a ``Pulse`` instance, which can be generated from either the tuple with positional arguments
    or from the dictionary, or set manually.

    If delay is not provided in **all** pulses in the sequence, assume equispaced pulse sequence:

        t - pulse - 2t - pulse - 2t - ... - pulse - t

    If only **some** delays are provided, assumes 0 delay in the pulses without delay provided.


    Examples:
        >>> import numpy as np
        >>> Sequence([('x', np.pi, 0),
        >>>           {'axis': 'y', 'angle': 'pi', 'delay': np.linspace(0, 1, 3), 'bath_names': '13C'},
        >>>           Pulse('x', 'pi', 1)])
        [Pulse((x, 3.14), t = 0), Pulse((y, 3.14), {13C: (y, 3.14)}, t = [0.  0.5 1. ]), Pulse((x, 3.14), t = 1)]
    """

    def __init__(self, t=None):
        if t is not None:
            t = [_get_pulse(x) for x in t]

        super().__init__(t)

        self.small_sigma = None
        """dict: Dictionary with Pauli matrices of the central spin"""
        self.delays = None
        """list or None: List with delays before each pulse or None if equispaced.
        Generated by ``.generate_pulses`` method."""
        self.rotations = None
        """list: List with matrix representations of the rotation from each pulse.
        Generated by ``.generate_pulses`` method."""

    def __setitem__(self, key, value):
        self.data[key] = _get_pulse(value)

    def append(self, item):
        self.data.append(_get_pulse(item))

    def set_central_spin(self, alpha, beta):
        r"""
        Set Pauli matrices of the central spin.

        Args:
            alpha (ndarray): :math:`\ket{0}` state of the qubit in :math:`S_z` basis.
            beta (ndarray): :math:`\ket{1}` state of the qubit in :math:`S_z` basis.

        """
        alpha_x_alpha = np.tensordot(alpha, alpha, axes=0)
        beta_x_beta = np.tensordot(beta, beta, axes=0)
        alpha_x_beta = np.tensordot(alpha, beta, axes=0)
        beta_x_alpha = np.tensordot(beta, alpha, axes=0)

        self.small_sigma = {'x': alpha_x_beta + beta_x_alpha,
                            'y': -1j * alpha_x_beta + 1j * beta_x_alpha,
                            'z': alpha_x_alpha - beta_x_beta}

    def generate_pulses(self, dimensions=None,
                        bath=None, vectors=None, central_spin=True):
        """
        Generate list of matrix representations of the rotations, induced by the sequence of the pulses.

        The rotations are stored in the ``.rotation`` attribute of the each ``Pulse`` object
        and in ``Sequence.rotations``.

        Args:
            dimensions (ndarray with shape (N,)): Array of spin dimensions in the system.
            bath (BathArray with shape (n,)): Array of bath spins in the system.
            vectors (ndarray  with shape (N, 3, prod(dimensions), prod(dimensions))):
                 Array with vectors of spin matrices for each spin in the system.

            central_spin (bool): True if generate the rotations including central spin rotations. Default is True.

        Returns:
            tuple: *tuple* containing:

            * **list** or **None**: List with delays before each pulse or None if equispaced.

            * **list**: List with matrix representations of the rotation from each pulse.
        """
        self.delays = None
        self.rotations = None

        equispaced = not any(p._has_delay for p in self.data)

        if equispaced:
            delays = None

        else:
            delays = [p.delay if p.delay is not None else 0 for p in self.data]

        rots = []
        sigma = {}

        for p in self.data:

            initial_rotation = rotation = 1

            if p.angle and central_spin:

                if p.axis not in sigma:
                    sigma[p.axis] = expand(self.small_sigma[p.axis], len(dimensions) - 1, dimensions)

                if p.angle == np.pi:
                    rotation = -1j * sigma[p.axis]
                else:
                    rotation = scipy.linalg.expm(-1j * sigma[p.axis] * p.angle / 2)

            if p.bath_names is not None:
                if vectors.shape != bath.shape:
                    vectors = vectors[:bath.shape[0]]
                    # print(vectors.shape)
                properties = np.broadcast(p.bath_names, p.bath_axes, p.bath_angles)

                for name, axis, angle in properties:
                    if angle:
                        which = (bath.N == name)

                        if any(which):
                            vecs = vectors[which]
                            rotation = np.dot(bath_rotation(vecs, axis, angle), rotation)

            if initial_rotation is rotation:
                rotation = None

            p.rotation = rotation

            rots.append(rotation)

        self.delays = delays
        self.rotations = rots
        return delays, rots


def bath_rotation(vectors, axis, angle):
    """
    Generate rotation of the bath spins with given spin vectors.

    Args:
        vectors (ndarray with shape (n, 3, x, x)): Array of *n* bath spin vectors.
        axis (str): Axis of rotation.
        angle (float): Angle of rotation.

    Returns:
        ndarray with shape (x, x): Matrix representation of the spin rotation.

    """
    ax = _rot[axis]  # name -> index
    if (angle == np.pi) and (vectors[0][0, 0, 0] < 1):  # only works for spin-1/2
        rotation = -1j * 2 * vectors[0][ax]  # 2 here is to transform into pauli matrices
        for v in vectors[1:]:
            np.matmul(rotation, -1j * 2 * v[ax], out=rotation)
    else:
        rotation = scipy.linalg.expm(-1j * vectors[0][ax] * angle)
        for v in vectors[1:]:
            add = scipy.linalg.expm(-1j * v[ax] * angle)
            np.matmul(rotation, add, out=rotation)

    return rotation
