import warnings

import numpy as np

import treams._operators as op
import treams.special as sc
from treams import config
from treams._core import CylindricalWaveBasis as CWB
from treams._core import PhysicsArray
from treams._core import PlaneWaveBasisByComp as PWBC
from treams._core import PlaneWaveBasisByUnitVector as PWBUV
from treams._core import SphericalWaveBasis as SWB
from treams._material import Material
from treams.coeffs import mie, mie_cyl
from treams.util import AnnotationError


class _Interaction:
    def __init__(self):
        self._obj = self._objtype = None

    def __get__(self, obj, objtype=None):
        self._obj = obj
        self._objtype = objtype
        return self

    def __call__(self):
        basis = self._obj.basis
        return (
            np.eye(self._obj.shape[-1]) - self._obj @ op.Expand(basis, "singular").inv
        )

    def solve(self):
        return np.linalg.solve(self(), self._obj)


class _LatticeInteraction:
    def __init__(self):
        self._obj = self._objtype = None

    def __get__(self, obj, objtype=None):
        self._obj = obj
        self._objtype = objtype
        return self

    def __call__(self, lattice, kpar):
        return np.eye(self._obj.shape[-1]) - self._obj @ op.ExpandLattice(
            lattice=lattice, kpar=kpar
        )

    def solve(self, lattice, kpar):
        return np.linalg.solve(self(lattice, kpar), self._obj)


class TMatrix(PhysicsArray):
    """T-matrix with a spherical basis.

    The T-matrix is square relating incident (regular) fields
    :func:`treams.special.vsw_rA` (helical polarizations) or
    :func:`treams.special.vsw_rN` and :func:`treams.special.vsw_rM` (parity
    polarizations) to the corresponding scattered fields :func:`treams.special.vsw_A` or
    :func:`treams.special.vsw_N` and :func:`treams.special.vsw_M`. The modes themselves
    are defined in :attr:`basis`, the polarization type in :attr:`poltype`. Also, the
    wave number :attr:`k0` and, if not vacuum, the material :attr:`material` are
    specified.

    Args:
        arr (float or complex, array-like): T-matrix itself.
        k0 (float): Wave number in vacuum.
        basis (SphericalWaveBasis, optional): Basis definition.
        material (Material, optional): Embedding material, defaults to vacuum.
        poltype (str, optional): Polarization type (:ref:`params:Polarizations`).
        lattice (Lattice, optional): Lattice definition. If specified the T-Matrix is
            assumed to be periodically repeated in the defined lattice.
        kpar (list, optional): Phase factor for the periodic T-Matrix.
    """

    interaction = _Interaction()
    latticeinteraction = _LatticeInteraction()

    def _check(self):
        """Fill in default values or raise errors for missing attributes."""
        super()._check()
        shape = np.shape(self)
        if len(shape) != 2 or shape[0] != shape[1]:
            raise AnnotationError(f"invalid shape: '{shape}'")
        if not isinstance(self.k0, (int, float, np.floating, np.integer)):
            raise AnnotationError("invalid k0")
        if self.poltype is None:
            self.poltype = config.POLTYPE
        if self.poltype not in ("parity", "helicity"):
            raise AnnotationError("invalid poltype")
        modetype = self.modetype
        if modetype is None or (
            modetype[0] in (None, "singular") and modetype[1] in (None, "regular")
        ):
            self.modetype = ("singular", "regular")
        else:
            raise AnnotationError("invalid modetype")
        if self.basis is None:
            self.basis = SWB.default(SWB.defaultlmax(shape[0]))
        if self.material is None:
            self.material = Material()

    @property
    def ks(self):
        """Wave numbers (in medium).

        The wave numbers for both polarizations.
        """
        return self.material.ks(self.k0)

    @classmethod
    def sphere(cls, lmax, k0, radii, materials):
        """T-Matrix of a (multi-layered) sphere.

        Construct the T-matrix of the given order and material for a sphere. The object
        can also consist of multiple concentric spherical shells with an arbitrary
        number of layers. The calculation is always done in helicity basis.

        Args:
            lmax (int): Positive integer for the maximum degree of the T-matrix.
            k0 (float): Wave number in vacuum.
            radii (float or array): Radii from inside to outside of the sphere. For a
                simple sphere the radius can be given as a single number, for a multi-
                layered sphere it is a list of increasing radii for all shells.
            material (list[Material]): The material parameters from the inside to the
                outside. The last material in the list specifies the embedding medium.

        Returns:
            TMatrix
        """
        materials = [Material(m) for m in materials]
        radii = np.atleast_1d(radii)
        if radii.size != len(materials) - 1:
            raise ValueError("incompatible lengths of radii and materials")
        dim = SWB.defaultdim(lmax)
        tmat = np.zeros((dim, dim), np.complex128)
        for l in range(1, lmax + 1):  # noqa: E741
            miecoeffs = mie(l, k0 * radii, *zip(*materials))
            pos = SWB.defaultdim(l - 1)
            for i in range(2 * l + 1):
                tmat[
                    pos + 2 * i : pos + 2 * i + 2, pos + 2 * i : pos + 2 * i + 2
                ] = miecoeffs[::-1, ::-1]
        return cls(tmat, k0=k0, basis=SWB.default(lmax), material=materials[-1])

    @classmethod
    def cluster(cls, tmats, positions):
        r"""Block-diagonal T-matrix of multiple objects.

        Construct the initial block-diagonal T-matrix for a cluster of objects. The
        T-matrices in the list are placed together into a block-diagonal matrix and the
        complete (local) basis is defined based on the individual T-matrices and their
        bases together with the defined positions. In mathematical terms the matrix

        .. math::

            \begin{pmatrix}
                T_0 & 0 & \dots & 0 \\
                0 & T_1 & \ddots & \vdots \\
                \vdots & \ddots & \ddots & 0 \\
                0 & \dots & 0 & T_{N-1} \\
            \end{pmatrix}

        is created from the list of T-matrices :math:`(T_0, \dots, T_{N-1})`. Only
        T-matrices of the same wave number, embedding material, and polarization type
        can be combined.

        Args:
            tmats (Sequence): List of T-matrices.
            positions (array): The positions of all individual objects in the cluster.

        Returns:
            TMatrix
        """
        for tm in tmats:
            if not tm.basis.isglobal:
                raise ValueError("global basis required")
        positions = np.array(positions)
        if len(tmats) < positions.shape[0]:
            warnings.warn("specified more positions than T-matrices")
        elif len(tmats) > positions.shape[0]:
            raise ValueError(
                f"'{len(tmats)}' T-matrices "
                f"but only '{positions.shape[0]}' positions given"
            )
        mat = tmats[0].material
        k0 = tmats[0].k0
        poltype = tmats[0].poltype
        modes = [], [], []
        pidx = []
        dim = sum(tmat.shape[0] for tmat in tmats)
        tres = np.zeros((dim, dim), complex)
        i = 0
        for j, tm in enumerate(tmats):
            if tm.material != mat:
                raise ValueError(f"incompatible materials: '{mat}' and '{tm.material}'")
            if tm.k0 != k0:
                raise ValueError(f"incompatible k0: '{k0}' and '{tm.k0}'")
            if tm.poltype != poltype:
                raise ValueError(f"incompatible modetypes: '{poltype}', '{tm.poltype}'")
            dim = tm.shape[0]
            for m, n in zip(modes, tm.basis.lms):
                m.extend(list(n))
            pidx += [j] * dim
            tres[i : i + dim, i : i + dim] = tm
            i += dim
        basis = SWB(zip(pidx, *modes), positions)
        return cls(tres, k0=k0, material=mat, basis=basis, poltype=poltype)

    @property
    def isglobal(self):
        """Test if a T-matrix is global.

        A T-matrix is considered global, when its basis refers to only a single point
        and it is not placed periodically in a lattice.
        """
        return self.basis.isglobal and self.lattice is None and self.kpar is None

    @property
    def xs_ext_avg(self):
        r"""Rotation and polarization averaged extinction cross section.

        The average is calculated as

        .. math::

            \langle \sigma_\mathrm{ext} \rangle
            = -2 \pi \sum_{slm} \frac{\Re(T_{slm,slm})}{k_s^2}

        where :math:`k_s` is the wave number in the embedding medium for the
        polarization :math:`s`. It is only implemented for global T-matrices.

        Returns:
            float or complex
        """
        if not self.isglobal or not self.material.isreal:
            raise NotImplementedError
        if not self.material.ischiral:
            k = self.ks[0]
            res = -2 * np.pi * self.trace().real / (k * k)
        else:
            res = 0
            modetype = self.modetype
            del self.modetype
            diag = np.diag(self)
            self.modetype = modetype
            for pol in [0, 1]:
                choice = self.basis.pol == pol
                k = self.ks[pol]
                res += -2 * np.pi * diag[choice].sum().real / (k * k)
        if res.imag == 0:
            return res.real
        return res

    @property
    def xs_sca_avg(self):
        r"""Rotation and polarization averaged scattering cross section.

        The average is calculated as

        .. math::

            \langle \sigma_\mathrm{sca} \rangle
            = 2 \pi \sum_{slm} \sum_{s'l'm'}
            \frac{|T_{slm,s'l'm'}|^2}{k_s^2}

        where :math:`k_s` is the wave number in the embedding medium for the
        polarization :math:`s`. It is only implemented for global T-matrices.

        Returns:
            float or complex
        """
        if not self.isglobal or not self.material.isreal:
            raise NotImplementedError
        if not self.material.ischiral:
            ks = self.ks[0]
        else:
            ks = self.ks[self.basis.pol, None]
        re, im = self.real, self.imag
        res = 2 * np.pi * np.sum((re * re + im * im) / (ks * ks))
        return res.real

    @property
    def cd(self):
        r"""Circular dichroism (CD).

        The CD is calculated as

        .. math::

            CD
            = \frac{\langle \sigma_\mathrm{abs} \rangle_+
            - \langle \sigma_\mathrm{abs} \rangle_-}
            {\langle \sigma_\mathrm{abs} \rangle_+
            + \langle \sigma_\mathrm{abs} \rangle_-}

        where :math:`\langle \sigma \rangle_s` is the rotationally averaged absorption
        cross section under the illumination with the polarization :math:`s`.
        It is only implemented for global T-matrices.

        Returns:
            float or complex
        """
        if not (self.isglobal and self.poltype == "helicity" and self.material.isreal):
            raise NotImplementedError
        sel = np.array(self.basis.pol, bool)
        re, im = self.real, self.imag
        plus = -np.sum(re[sel[:, None] & sel]) / (self.ks[1] * self.ks[1])
        re_part = re[:, sel] / self.ks[self.basis.pol, None]
        im_part = im[:, sel] / self.ks[self.basis.pol, None]
        plus -= np.sum(re_part * re_part + im_part * im_part)
        sel = ~sel
        minus = -np.sum(re[sel[:, None] & sel]) / (self.ks[0] * self.ks[0])
        re_part = re[:, sel] / self.ks[self.basis.pol, None]
        im_part = im[:, sel] / self.ks[self.basis.pol, None]
        minus -= np.sum(re_part * re_part + im_part * im_part)
        return np.real((plus - minus) / (plus + minus))

    @property
    def chi(self):
        """Electromagnetic chirality.

        Only implemented for global T-matrices.

        Returns:
            float
        """
        if not (self.isglobal and self.poltype == "helicity"):
            raise NotImplementedError
        sel = self.basis.pol == 0, self.basis.pol == 1
        spp = np.linalg.svd(self[np.ix_(sel[1], sel[1])], compute_uv=False)
        spm = np.linalg.svd(self[np.ix_(sel[1], sel[0])], compute_uv=False)
        smp = np.linalg.svd(self[np.ix_(sel[0], sel[1])], compute_uv=False)
        smm = np.linalg.svd(self[np.ix_(sel[0], sel[0])], compute_uv=False)
        plus = np.concatenate((np.asarray(spp), np.asarray(spm)))
        minus = np.concatenate((np.asarray(smm), np.asarray(smp)))
        return np.linalg.norm(plus - minus) / np.sqrt(np.sum(np.power(np.abs(self), 2)))

    @property
    def db(self):
        """Duality breaking.

        Only implemented for global T-matrices.

        Returns:
            float
        """
        if not (self.isglobal and self.poltype == "helicity"):
            raise NotImplementedError
        sel = self.basis.pol == 0, self.basis.pol == 1
        tpm = np.asarray(self[np.ix_(sel[1], sel[0])])
        tmp = np.asarray(self[np.ix_(sel[0], sel[1])])
        return np.sum(
            tpm.real * tpm.real
            + tpm.imag * tpm.imag
            + tmp.real * tmp.real
            + tmp.imag * tmp.imag
        ) / (np.sum(np.power(np.abs(self), 2)))

    def xs(self, illu, flux=0.5):
        r"""Scattering and extinction cross section.

        Possible for all T-matrices (global and local) in non-absorbing embedding. The
        values are calculated by

        .. math::

            \sigma_\mathrm{sca}
            = \frac{1}{2 I}
            a_{slm}^\ast T_{s'l'm',slm}^\ast k_{s'}^{-2} C_{s'l'm',s''l''m''}^{(1)}
            T_{s''l''m'',s'''l'''m'''} a_{s'''l'''m'''} \\
            \sigma_\mathrm{ext}
            = \frac{1}{2 I}
            a_{slm}^\ast k_s^{-2} T_{slm,s'l'm'} a_{s'l'm'}

        where :math:`a_{slm}` are the expansion coefficients of the illumination,
        :math:`T` is the T-matrix, :math:`C^{(1)}` is the (regular) translation
        matrix and :math:`k_s` are the wave numbers in the medium. All repeated indices
        are summed over. The incoming flux is :math:`I`.

        Args:
            illu (complex, array): Illumination coefficients
            flux (optional): Ingoing flux corresponding to the illumination. Used for
                the result's normalization. The flux is given in units of
                :math:`\frac{\text{V}^2}{{l^2}} \frac{1}{Z_0 Z}` where :math:`l` is the
                unit of length used in the wave number (and positions). A plane wave
                has the flux `0.5` in this normalization, which is used as default.

        Returns:
            tuple[float]
        """
        if not self.material.isreal:
            raise NotImplementedError
        illu = PhysicsArray(illu)
        illu_basis = illu.basis
        illu_basis = illu_basis[-2] if isinstance(illu_basis, tuple) else illu_basis
        if not isinstance(illu_basis, SWB):
            illu = illu.expand(self.basis)
        p = self @ illu
        p_invksq = p * np.power(self.ks[self.basis.pol], -2)
        del illu.modetype
        return (
            0.5 * np.real(p.conjugate().T @ p_invksq.expand(p.basis)) / flux,
            -0.5 * np.real(illu.conjugate().T @ p_invksq) / flux,
        )

    def valid_points(self, grid, radii):
        """Points on the grid where the expansion is valid.

        The expansion of the electromagnetic field is valid outside of the
        circumscribing spheres of each object. From a given set of coordinates mark
        those that are outside of the given radii.

        Args:
            grid (array-like): Points to assess. The last dimension needs length three
                and corresponds to the Cartesian coordinates.
            radii (Sequence[float]): Radii of the circumscribing spheres. Each radius
                corresponds to a position of the basis.

        Returns:
            array
        """
        grid = np.asarray(grid)
        if grid.shape[-1] != 3:
            raise ValueError("invalid grid")
        if len(radii) != len(self.basis.positions):
            raise ValueError("invalid length of 'radii'")
        res = np.ones(grid.shape[:-1], bool)
        for r, p in zip(radii, self.basis.positions):
            res &= np.sum(np.power(grid - p, 2), axis=-1) > r * r
        return res

    def __getitem__(self, key):
        if isinstance(key, SWB):
            key = np.array([self.basis.index(i) for i in key])
            key = (key[:, None], key)
        return super().__getitem__(key)


class TMatrixC(PhysicsArray):
    """T-matrix with a cylindrical basis.

    The T-matrix is square relating incident (regular) fields
    :func:`treams.special.vcw_rA` (helical polarizations) or
    :func:`treams.special.vcw_rN` and :func:`treams.special.vcw_rM` (parity
    polarizations) to the corresponding scattered fields :func:`treams.special.vcw_A` or
    :func:`treams.special.vcw_N` and :func:`treams.special.vcw_M`. The modes themselves
    are defined in :attr:`basis`, the polarization type in :attr:`poltype`. Also, the
    wave number :attr:`k0` and, if not vacuum, the material :attr:`material` are
    specified.

    Args:
        arr (float or complex, array-like): T-matrix itself.
        k0 (float): Wave number in vacuum.
        basis (SphericalWaveBasis, optional): Basis definition.
        material (Material, optional): Embedding material, defaults to vacuum.
        poltype (str, optional): Polarization type (:ref:`params:Polarizations`).
        lattice (Lattice, optional): Lattice definition. If specified the T-Matrix is
            assumed to be periodically repeated in the defined lattice.
        kpar (list, optional): Phase factor for the periodic T-Matrix.
    """

    interaction = _Interaction()
    latticeinteraction = _LatticeInteraction()

    def _check(self):
        """Fill in default values or raise errors for missing attributes."""
        super()._check()
        shape = np.shape(self)
        if len(shape) != 2 or shape[0] != shape[1]:
            raise AnnotationError(f"invalid shape: '{shape}'")
        if not isinstance(self.k0, (int, float, np.floating, np.integer)):
            raise AnnotationError("invalid k0")
        if self.poltype is None:
            self.poltype = config.POLTYPE
        if self.poltype not in ("parity", "helicity"):
            raise AnnotationError("invalid poltype")
        modetype = self.modetype
        if modetype is None or (
            modetype[0] in (None, "singular") and modetype[1] in (None, "regular")
        ):
            self.modetype = ("singular", "regular")
        else:
            raise AnnotationError("invalid modetype")
        if self.basis is None:
            self.basis = CWB.default([0], CWB.defaultmmax(shape[0]))
        if self.material is None:
            self.material = Material()

    @property
    def ks(self):
        """Wave numbers (in medium).

        The wave numbers for both polarizations.
        """
        return self.material.ks(self.k0)

    @property
    def krhos(self):
        r"""Radial part of the wave.

        Calculate :math:`\sqrt{k^2 - k_z^2}`, where :math:`k` is the wave number in the
        medium for each illumination

        Returns:
            Sequence[complex]
        """
        return self.material.krhos(self.k0, self.basis.kz, self.basis.pol)

    @classmethod
    def cylinder(cls, kzs, mmax, k0, radii, materials):
        """T-Matrix of a (multi-layered) cylinder.

        Construct the T-matrix of the given order and material for an infinitely
        extended cylinder. The object can also consist of multiple concentric
        cylindrical shells with an arbitrary number of layers. The calculation is always
        done in helicity basis.

        Args:
            kzs (float, array_like): Z component of the cylindrical wave.
            mmax (int): Positive integer for the maximum order of the T-matrix.
            k0 (float): Wave number in vacuum.
            radii (float or array): Radii from inside to outside of the cylinder. For a
                simple cylinder the radius can be given as a single number, for a multi-
                layered cylinder it is a list of increasing radii for all shells.
            material (list[Material]): The material parameters from the inside to the
                outside. The last material in the list specifies the embedding medium.

        Returns:
            T-Matrix object
        """
        materials = [Material(m) for m in materials]
        kzs = np.atleast_1d(kzs)
        radii = np.atleast_1d(radii)
        if radii.size != len(materials) - 1:
            raise ValueError("incompatible lengths of radii and materials")
        dim = CWB.defaultdim(len(kzs), mmax)
        tmat = np.zeros((dim, dim), np.complex128)
        idx = 0
        for kz in kzs:
            for m in range(-mmax, mmax + 1):
                miecoeffs = mie_cyl(kz, m, k0, radii, *zip(*materials))
                tmat[idx : idx + 2, idx : idx + 2] = miecoeffs[::-1, ::-1]
                idx += 2
        return cls(tmat, k0=k0, basis=CWB.default(kzs, mmax), material=materials[-1])

    @classmethod
    def cluster(cls, tmats, positions):
        r"""Block-diagonal T-matrix of multiple objects.

        Construct the initial block-diagonal T-matrix for a cluster of objects. The
        T-matrices in the list are placed together into a block-diagonal matrix and the
        complete (local) basis is defined based on the individual T-matrices and their
        bases together with the defined positions. In mathematical terms the matrix

        .. math::

            \begin{pmatrix}
                T_0 & 0 & \dots & 0 \\
                0 & T_1 & \ddots & \vdots \\
                \vdots & \ddots & \ddots & 0 \\
                0 & \dots & 0 & T_{N-1} \\
            \end{pmatrix}

        is created from the list of T-matrices :math:`(T_0, \dots, T_{N-1})`. Only
        T-matrices of the same wave number, embedding material, and polarization type
        can be combined.

        Args:
            tmats (Sequence): List of T-matrices.
            positions (array): The positions of all individual objects in the cluster.

        Returns:
            TMatrix
        """
        for tm in tmats:
            if not tm.basis.isglobal:
                raise ValueError("global basis required")
        positions = np.array(positions)
        if len(tmats) < positions.shape[0]:
            warnings.warn("specified more positions than T-matrices")
        elif len(tmats) > positions.shape[0]:
            raise ValueError(
                f"'{len(tmats)}' T-matrices "
                f"but only '{positions.shape[0]}' positions given"
            )
        mat = tmats[0].material
        k0 = tmats[0].k0
        poltype = tmats[0].poltype
        modes = [], [], []
        pidx = []
        dim = sum(tmat.shape[0] for tmat in tmats)
        tres = np.zeros((dim, dim), complex)
        i = 0
        for j, tm in enumerate(tmats):
            if tm.material != mat:
                raise ValueError(f"incompatible materials: '{mat}' and '{tm.material}'")
            if tm.k0 != k0:
                raise ValueError(f"incompatible k0: '{k0}' and '{tm.k0}'")
            if tm.poltype != poltype:
                raise ValueError(f"incompatible modetypes: '{poltype}', '{tm.poltype}'")
            dim = tm.shape[0]
            for m, n in zip(modes, tm.basis.zms):
                m.extend(list(n))
            pidx += [j] * dim
            tres[i : i + dim, i : i + dim] = tm
            i += dim
        basis = CWB(zip(pidx, *modes), positions)
        return cls(tres, k0=k0, material=mat, basis=basis, poltype=poltype)

    @classmethod
    def from_array(cls, tm, basis, *, eta=0):
        """1d array of spherical T-matrices."""
        return cls(
            (tm @ op.Expand(basis).inv).expandlattice(basis=basis, eta=eta),
            lattice=tm.lattice,
            kpar=tm.kpar,
        )

    @property
    def xw_ext_avg(self):
        r"""Rotation and polarization averaged extinction cross width.

        The average is calculated as

        .. math::

            \langle \lambda_\mathrm{ext} \rangle
            = -\frac{2 \pi}{n_{k_z}} \sum_{sk_zm} \frac{\Re(T_{sk_zm,sk_zm})}{k_s}

        where :math:`k_s` is the wave number in the embedding medium for the
        polarization :math:`s` and :math:`n_{k_z}` is the number of wave components
        :math:`k_z` included in the T-matrix. The average is taken over all given
        z-components of the wave vector and rotations around the z-axis. It is only
        implemented for global T-matrices.

        Returns:
            float or complex
        """
        if not self.isglobal or not self.material.isreal:
            raise NotImplementedError
        nk = np.unique(self.basis.kz).size
        if not self.material.ischiral:
            res = -2 * np.real(np.trace(self)) / (self.ks[0] * nk)
        else:
            res = 0
            modetype = self.modetype
            del self.modetype
            diag = np.diag(self)
            self.modetype = modetype
            for pol in [0, 1]:
                choice = self.basis.pol == pol
                res += -2 * np.real(diag[choice].sum()) / (self.ks[pol] * nk)
        if res.imag == 0:
            return res.real
        return res

    @property
    def xw_sca_avg(self):
        r"""Rotation and polarization averaged scattering cross width.

        The average is calculated as

        .. math::

            \langle \lambda_\mathrm{sca} \rangle
            = \frac{2 \pi}{n_{k_z}} \sum_{sk_zm} \sum_{s'{k_z}'m'}
            \frac{|T_{sk_zm,s'{k_z}'m'}|^2}{k_s}

        where :math:`k_s` is the wave number in the embedding medium for the
        polarization :math:`s`. and :math:`n_{k_z}` is the number of wave components
        :math:`k_z` included in the T-matrix. The average is taken over all given
        z-components of the wave vector and rotations around the z-axis. It is only
        implemented for global T-matrices.

        Returns:
            float or complex
        """
        if not self.isglobal or not self.material.isreal:
            raise NotImplementedError
        if not self.material.ischiral:
            ks = self.ks[0]
        else:
            ks = self.ks[self.basis.pol, None]
        re, im = self.real, self.imag
        nk = np.unique(self.basis.kz).size
        res = 2 * np.sum((re * re + im * im) / (ks * nk))
        return res.real

    @property
    def isglobal(self):
        """Test if a T-matrix is global.

        A T-matrix is considered global, when its basis refers to only a single point
        and it is not placed periodically in a lattice.
        """
        return self.basis.isglobal and self.lattice is None and self.kpar is None

    def xw(self, illu, flux=0.5):
        r"""Scattering and extinction cross width.

        Possible for all T-matrices (global and local) in non-absorbing embedding. The
        values are calculated by

        .. math::

            \lambda_\mathrm{sca}
            = \frac{1}{2 I}
            a_{sk_zm}^\ast T_{s'{k_z}'m',sk_zm}^\ast k_{s'}^{-2}
            C_{s'l'm',s''l''m''}^{(1)}
            T_{s''l''m'',s'''l'''m'''} a_{s'''l'''m'''} \\
            \sigma_\mathrm{ext}
            = \frac{1}{2 I}
            a_{slm}^\ast k_s^{-2} T_{slm,s'l'm'} a_{s'l'm'}

        where :math:`a_{slm}` are the expansion coefficients of the illumination,
        :math:`T` is the T-matrix, :math:`C^{(1)}` is the (regular) translation
        matrix and :math:`k_s` are the wave numbers in the medium. All repeated indices
        are summed over. The incoming flux is :math:`I`.

        Args:
            illu (complex, array): Illumination coefficients
            flux (optional): Ingoing flux corresponding to the illumination. Used for
                the result's normalization. The flux is given in units of
                :math:`\frac{\text{V}^2}{{l^2}} \frac{1}{Z_0 Z}` where :math:`l` is the
                unit of length used in the wave number (and positions). A plane wave
                has the flux `0.5` in this normalization, which is used as default.

        Returns:
            tuple[float]
        """
        if not self.material.isreal:
            raise NotImplementedError
        illu = PhysicsArray(illu)
        illu_basis = illu.basis
        illu_basis = illu_basis[-2] if isinstance(illu_basis, tuple) else illu_basis
        if not isinstance(illu_basis, CWB):
            illu = illu.expand(self.basis) @ illu
        p = self @ illu
        m = self.expand() / self.ks[self.basis.pol]
        del illu.modetype
        return (
            2 * np.real(p.conjugate().T @ (m @ p)) / flux,
            -2 * np.real(illu.conjugate().T @ (p / self.ks[self.basis.pol])) / flux,
        )

    def valid_points(self, grid, radii):
        grid = np.asarray(grid)
        if grid.shape[-1] not in (2, 3):
            raise ValueError("invalid grid")
        if len(radii) != len(self.basis.positions):
            raise ValueError("invalid length of 'radii'")
        res = np.ones(grid.shape[:-1], bool)
        for r, p in zip(radii, self.basis.positions):
            res &= np.sum(np.power(grid[..., :2] - p[:2], 2), axis=-1) > r * r
        return res

    def __getitem__(self, key):
        if isinstance(key, CWB):
            key = np.array([self.basis.index(i) for i in key])
            key = (key[:, None], key)
        return super().__getitem__(key)


def _plane_wave_partial(
    kpar, pol, *, k0=None, basis=None, material=None, modetype=None, poltype=None,
):
    if basis is None:
        basis = PWBC.default([kpar])
    if pol in (0, -1):
        pol = [1, 0]
    elif pol == 1:
        pol = [0, 1]
    elif len(pol) == 3:
        modetype = "up" if modetype is None else modetype
        if modetype not in ("up", "down"):
            raise ValueError(f"invalid 'modetype': {modetype}")
        kzs = Material(material).kzs(k0, *kpar, [0, 1])
        poltype = config.POLTYPE if poltype is None else poltype
        if poltype == "parity":
            pol = [
                -sc.vpw_M(*kpar, kzs[0], 0, 0, 0) @ pol,
                sc.vpw_N(*kpar, kzs[1], 0, 0, 0) @ pol,
            ]
        elif poltype == "helicity":
            pol = sc.vpw_A(*kpar, kzs[::-1], 0, 0, 0, [1, 0]) @ pol
        else:
            raise ValueError(f"invalid 'poltype': {poltype}")
    res = [pol[x[2]] * (np.abs(np.array(kpar) - x[:2]) < 1e-14).all() for x in basis]
    return PhysicsArray(
        res, basis=basis, k0=k0, material=material, modetype=modetype, poltype=poltype,
    )


def _plane_wave(
    kvec, pol, *, k0=None, basis=None, material=None, modetype=None, poltype=None,
):
    if basis is None:
        basis = PWBUV.default([kvec])
    norm = np.sqrt(np.sum(np.power(kvec, 2)))
    qvec = kvec / norm
    if pol in (0, -1):
        pol = [1, 0]
    elif pol == 1:
        pol = [0, 1]
    elif len(pol) == 3:
        if None not in (k0, material):
            kvec = Material(material).ks(k0) * qvec[:, None]
        else:
            kvec = qvec
        poltype = config.POLTYPE if poltype is None else poltype
        if poltype == "parity":
            pol = [
                -sc.vpw_M(*kvec[:, 0], 0, 0, 0) @ pol,
                sc.vpw_N(*kvec[:, 1], 0, 0, 0) @ pol,
            ]
        elif poltype == "helicity":
            pol = sc.vpw_A(*kvec, 0, 0, 0, [1, 0]) @ pol
        else:
            raise ValueError(f"invalid 'poltype': {poltype}")
    res = [pol[x[3]] * (np.abs(qvec - x[:3]) < 1e-14).all() for x in basis]
    return PhysicsArray(
        res, basis=basis, k0=k0, material=material, modetype=modetype, poltype=poltype,
    )


def plane_wave(
    kvec, pol, *, k0=None, basis=None, material=None, modetype=None, poltype=None,
):
    """Array describing a plane wave.

    Args:
        kvec (Sequence): Wave vector.
        pol (int or Sequence): Polarization index (see
            :ref:`params:Polarizations`) to have a unit amplitude wave of the
            corresponding wave, two values to specify the amplitude for each
            polarization, or three values in a sequence to specify the Cartesian
            electric field components. In the latter case, if the electric field has
            longitudinal components they are neglected.
        basis (PlaneWaveBasis, optional): Basis definition.
        k0 (float, optional): Wave number in vacuum.
        material (Material, optional): Material definition.
        modetype (str, optional): Mode type (see :ref:`params:Mode types`).
        poltype (str, optional): Polarization type (see
            :ref:`polarization:Polarizations`).
    """
    if len(kvec) == 2:
        return _plane_wave_partial(
            kvec,
            pol,
            k0=k0,
            basis=basis,
            material=material,
            modetype=modetype,
            poltype=poltype,
        )
    if len(kvec) == 3:
        return _plane_wave(
            kvec,
            pol,
            k0=k0,
            basis=basis,
            material=material,
            modetype=modetype,
            poltype=poltype,
        )
    raise ValueError(f"invalid length of 'kvec': {len(kvec)}")


def plane_wave_angle(theta, phi, pol, **kwargs):
    qvec = [np.sin(theta) * np.cos(phi), np.sin(theta) * np.sin(phi), np.cos(theta)]
    return plane_wave(qvec, pol, **kwargs)


def spherical_wave(
    l,  # noqa: E741
    m,
    pol,
    *,
    k0=None,
    basis=None,
    material=None,
    modetype=None,
    poltype=None,
):
    if basis is None:
        basis = SWB.default(l)
    if not basis.isglobal:
        raise ValueError("basis must be global")
    res = [0] * len(basis)
    res[basis.index((0, l, m, pol))] = 1
    return PhysicsArray(
        res, basis=basis, k0=k0, material=material, modetype=modetype, poltype=poltype
    )


def cylindrical_wave(
    kz, m, pol, *, k0=None, basis=None, material=None, modetype=None, poltype=None,
):
    if basis is None:
        basis = CWB.default([kz], abs(m))
    if not basis.isglobal:
        raise ValueError("basis must be global")
    res = [0] * len(basis)
    res[basis.index((0, kz, m, pol))] = 1
    return PhysicsArray(
        res, basis=basis, k0=k0, material=material, modetype=modetype, poltype=poltype
    )
