import numpy as np
import cmath

from treams import misc

class AcousticMaterial:
    r"""Material definition.

    The material properties are defined in the frequency domain by scalar values
    for density :math:`\rho`, the speed of longitudinal elastic waves :math:`c`, and
    the speed of transverse elastic waves :math:`c_t`. Materials are, thus, assumed to be linear, 
    time-invariant, homogeneous, isotropic, local and to have no gain.

    Args:
        rho (complex, optional): Mass density (kg/m^3). Defaults to 1.3 kg/m^3.
        c (complex, optional): Longitudinal speed of sound (m/s). Defaults to 343 m/s.
        ct (complex, optional): Transverse speed of sound (m/s). Defaults to 0 m/s.
    """

    def __init__(self, rho=1.3, c=343., ct=0.):
        """Initialization."""
        if isinstance(rho, AcousticMaterial):
            rho, c, ct = rho()
        elif isinstance(rho, (tuple, list, np.ndarray)):
            if len(rho) == 0:
               rho = 1.3
            elif len(rho) == 1:
                rho = rho[0]
            elif len(rho) == 2:
                rho, c = rho
            elif len(rho) == 3:
                rho, c, ct = rho
            else:
                raise ValueError("invalid material definition")
        self._rho = rho
        self._c = c
        self._ct = ct

    @property
    def rho(self):
        """Mass density.

        Returns:
            float or complex
        """
        return self._rho

    @property
    def c(self):
        """Speed of longitudinal elastic waves.

        Returns:
            float or complex
        """
        return self._c
    
    @property
    def ct(self):
        """Speed of transverse elastic waves.

        Returns:
            float or complex
        """
        return self._ct
    
    def __iter__(self):
        """Iterator for a tuple containing the material parameters.

        Useful for unpacking the material parameters into a function that takes these
        parameters separately, e.g. ``foo(*material)``.
        """
        return iter((self.rho, self.c, self.ct))
    
    @classmethod
    def from_params(cls, rho=1.3, params=(151767.21, 0.)):
        r"""Create acoustic material from Lamé parameters.

        This function calculates the longitudinal and transverse speeds of sound with
        :math:`c = \sqrt{\frac{\lambda+2\mu}{\rho}}` and :math:`c_t = \sqrt{\frac{\mu}{\rho}}`. 

        Args:
            rho (complex, optional): Mass density (kg/m^3). Defaults to 1.3 kg/m^3.
            params (complex, optional): Lamé parameters (Pa). Defaults to (151767.21, 0.) Pa.

        Returns:
            AcousticMaterial
        """
        c = cmath.sqrt((params[0] + 2 * params[1]) / rho)
        ct = cmath.sqrt(params[1] / rho)
        return cls(rho, c, ct)
    
    @classmethod
    def from_pratio(cls, rho=1.3, c=343, pratio=0.5):
        r"""Create acoustic material from Poisson's ratio.

        This function calculates the transverse speed of sound with
        :math:`c_t = c \sqrt{\frac{1 - 2\sigma}{2 - 2\sigma}}`. 

        Args:
            rho (complex, optional): Mass density (kg/m^3). Defaults to 1.3 kg/m^3.
            c (complex, optional): Longitudinal speed of sound (m/s). Defaults to 343 m/s.
            pratio (complex, optional): Poisson's ratio. Defaults to 0.5.

        Returns:
            AcousticMaterial
        """
        ct = c * cmath.sqrt((1 - 2 * pratio) / (2 - 2 * pratio))
        return cls(rho, c, ct)
    
    @property
    def impedance(self):
        r"""Relative impedance for longitudinal waves (Pa*s/m^3).

        The relative impedance is defined by :math:`Z = \rho c`.

        Returns:
            complex
        """
        return self.rho * self.c
    
    @property
    def impedancet(self):
        r"""Relative impedance for transverse waves (Pa*s/m^3).

        The relative impedance is defined by :math:`Z = \rho c_t`.

        Returns:
            complex
        """
        return self.rho * self.ct
    
    def __call__(self):
        """Return a tuple containing all material parameters.

        Returns:
            tuple
        """
        return self.rho, self.c, self.ct

    def __eq__(self, other):
        """Compare material parameters.

        Materials are considered equal if all material parameters are equal.
        Moreover, compares objects that contain at most three values.

        Returns:
            bool
        """
        if other is None:
            return False
        if not isinstance(other, AcousticMaterial):
            other = AcousticMaterial(*other)
        return (
            self.rho == other.rho
            and self.c == other.c
            and self.ct == other.ct
        )
    
    @property
    def isreal(self):
        """Test if the material has purely real-valued parameters.

        Returns:
            bool
        """
        return all(i.imag == 0 for i in self)
    
    @property
    def isshear(self):
        """Test if the material sustains shear waves.

        Returns:
            bool
        """
        return self.ct != 0

    def __str__(self):
        """All three material parameters.

        Returns:
            str
        """
        return "(" + ", ".join([str(i) for i in self()]) + ")"

    def __repr__(self):
        """Representation that allows recreating the object.

        Returns:
            str
        """
        return self.__class__.__name__ + str(self)

    def ks(self, k0):
        """Return the wavenumber of longitudinal waves in the medium.

        Args:
            k0 (float): Angular wavenumber in the air.

        Returns:
            tuple
        """
        if isinstance(self.c, (tuple, list, np.ndarray)):
            return tuple(k0 * AcousticMaterial().c / c for c in self.c)
        else:
            return k0 * AcousticMaterial().c / self.c

        
    def kst(self, k0):
        """Return the wavenumber of transverse waves in the medium.

        Args:
            k0 (float): Angular wavenumber in the air.

        Returns:
            tuple
        """
        if isinstance(self.ct, (tuple, list, np.ndarray)):
            return tuple(k0 * AcousticMaterial().c / ct for ct in self.ct)
        else:
            return k0 * AcousticMaterial().c / self.ct
    
    def krhos(self, k0, kz):
        r"""The (cylindrically) radial part of the wave vector for longitudinal waves .

        The cylindrically radial part is defined by :math:`k_\rho = \sqrt(k^2 - k_z^2)`.
        The returned values have non-negative imaginary parts.

        Args:
            k0 (float): Angular wavenumber in the air.
            kz (float or array_like): Z-component of the wave vector. Has units of k0.

        Returns:
            complex or array_like
        """
        ks = self.ks(k0)
        return misc.wave_vec_z(kz, 0, ks)

    def kzs(self, k0, kx, ky):
        r"""The z-component of the wave vector for longitudinal waves.

        The z-component of the wave vector is defined by
        :math:`k_z = \sqrt(k^2 - k_x^2 - k_y^2)`. The returned values have
        non-negative imaginary parts.

        Args:
            k0 (float): Angular wavenumber in the air.
            kx (float or array_like): X-component of the wave vector. Has units of k0.
            ky (float or array_like): Y-component of the wave vector. Has units of k0.

        Returns:
            complex or array_like
        """
        ks = self.ks(k0)
        return misc.wave_vec_z(kx, ky, ks)   
    
    def krhost(self, k0, kz):
        r"""The (cylindrically) radial part of the wave vector for transverse waves.

        The cylindrically radial part is defined by :math:`k_\rho = \sqrt(k^2 - k_z^2)`.
        The returned values have non-negative imaginary parts.

        Args:
            k0 (float): Angular wavenumber in the air.
            kz (float or array_like): Z-component of the wave vector. Has units of k0.

        Returns:
            complex or array_like
        """
        ks = self.kst(k0)
        return misc.wave_vec_z(kz, 0, ks)

    def kzst(self, k0, kx, ky):
        r"""The z-component of the wave vector for transverse waves.

        The z-component of the wave vector is defined by
        :math:`k_z = \sqrt(k^2 - k_x^2 - k_y^2)`. The returned values have
        non-negative imaginary parts.

        Args:
            k0 (float): Angular wavenumber in the air.
            kx (float or array_like): X-component of the wave vector. Has units of k0.
            ky (float or array_like): Y-component of the wave vector. Has units of k0.

        Returns:
            complex or array_like
        """
        ks = self.kst(k0)
        return misc.wave_vec_z(kx, ky, ks) 