import logging
from matplotlib.patches import Ellipse
import numpy as np
import typing

import autoarray as aa

logging.basicConfig()
logger = logging.getLogger(__name__)


class AbstractShearField:
    @property
    def ellipticities(self) -> aa.ValuesIrregular:
        """
        If we treat this vector field as a set of weak lensing shear measurements, the galaxy ellipticity each vector
        corresponds too.
        """
        return aa.ValuesIrregular(
            values=np.sqrt(self.slim[:, 0] ** 2 + self.slim[:, 1] ** 2.0)
        )

    @property
    def semi_major_axes(self) -> aa.ValuesIrregular:
        """
        If we treat this vector field as a set of weak lensing shear measurements, the semi-major axis of each
        galaxy ellipticity that each vector corresponds too.
        """
        return aa.ValuesIrregular(values=3 * (1 + self.ellipticities))

    @property
    def semi_minor_axes(self) -> aa.ValuesIrregular:
        """
        If we treat this vector field as a set of weak lensing shear measurements, the semi-minor axis of each
        galaxy ellipticity that each vector corresponds too.
        """
        return aa.ValuesIrregular(values=3 * (1 - self.ellipticities))

    @property
    def phis(self) -> aa.ValuesIrregular:
        """
        If we treat this vector field as a set of weak lensing shear measurements, the position angle defined
        counter clockwise from the positive x-axis of each galaxy ellipticity that each vector corresponds too.
        """
        return aa.ValuesIrregular(
            values=np.arctan2(self.slim[:, 0], self.slim[:, 1]) * 180.0 / np.pi / 2.0
        )

    @property
    def elliptical_patches(self) -> typing.List[Ellipse]:
        """
        If we treat this vector field as a set of weak lensing shear measurements, the elliptical patch representing
        each galaxy ellipticity. This patch is used for visualizing an ellipse of each galaxy in an image.
        """

        return [
            Ellipse(
                xy=(x, y), width=semi_major_axis, height=semi_minor_axis, angle=angle
            )
            for x, y, semi_major_axis, semi_minor_axis, angle in zip(
                self.grid.slim[:, 1],
                self.grid.slim[:, 0],
                self.semi_major_axes,
                self.semi_minor_axes,
                self.phis,
            )
        ]


class ShearYX2D(aa.VectorYX2D, AbstractShearField):
    """
    A regular shear field, which is collection of (y,x) vectors which are located on a uniform grid
    of (y,x) coordinates.

    The structure of this data structure is described in
    `autoarray.structures.vectors.uniform.VectorYX2D`

    This class extends `VectorYX2D` to include methods that are specific to a shear field, typically
    used for weak lensing calculations.

    Parameters
    ----------
    vectors
        The 2D (y,x) vectors on an irregular grid that represent the vector-field.
    grid
        The irregular grid of (y,x) coordinates where each vector is located.
    """


class ShearYX2DIrregular(aa.VectorYX2DIrregular, AbstractShearField):

    """
    An irregular shear field, which is collection of (y,x) vectors which are located on an irregular grid
    of (y,x) coordinates.

    The structure of this data structure is described in
    `autoarray.structures.vectors.irregular.VectorYX2DIrregular`

    This class extends `VectorYX2DIrregular` to include methods that are specific to a shear field, typically
    used for weak lensing calculations.

    Parameters
    ----------
    vectors
        The 2D (y,x) vectors on an irregular grid that represent the vector-field.
    grid
        The irregular grid of (y,x) coordinates where each vector is located.
    """

import logging
from matplotlib.patches import Ellipse
import numpy as np
import typing

import autoarray as aa

logging.basicConfig()
logger = logging.getLogger(__name__)


class AbstractShearField:
    @property
    def ellipticities(self) -> aa.ValuesIrregular:
        """
        If we treat this vector field as a set of weak lensing shear measurements, the galaxy ellipticity each vector
        corresponds too.
        """
        return aa.ValuesIrregular(
            values=np.sqrt(self.slim[:, 0] ** 2 + self.slim[:, 1] ** 2.0)
        )

    @property
    def semi_major_axes(self) -> aa.ValuesIrregular:
        """
        If we treat this vector field as a set of weak lensing shear measurements, the semi-major axis of each
        galaxy ellipticity that each vector corresponds too.
        """
        return aa.ValuesIrregular(values=3 * (1 + self.ellipticities))

    @property
    def semi_minor_axes(self) -> aa.ValuesIrregular:
        """
        If we treat this vector field as a set of weak lensing shear measurements, the semi-minor axis of each
        galaxy ellipticity that each vector corresponds too.
        """
        return aa.ValuesIrregular(values=3 * (1 - self.ellipticities))

    @property
    def phis(self) -> aa.ValuesIrregular:
        """
        If we treat this vector field as a set of weak lensing shear measurements, the position angle defined
        counter clockwise from the positive x-axis of each galaxy ellipticity that each vector corresponds too.
        """
        return aa.ValuesIrregular(
            values=np.arctan2(self.slim[:, 0], self.slim[:, 1]) * 180.0 / np.pi / 2.0
        )

    @property
    def elliptical_patches(self) -> typing.List[Ellipse]:
        """
        If we treat this vector field as a set of weak lensing shear measurements, the elliptical patch representing
        each galaxy ellipticity. This patch is used for visualizing an ellipse of each galaxy in an image.
        """

        return [
            Ellipse(
                xy=(x, y), width=semi_major_axis, height=semi_minor_axis, angle=angle
            )
            for x, y, semi_major_axis, semi_minor_axis, angle in zip(
                self.grid.slim[:, 1],
                self.grid.slim[:, 0],
                self.semi_major_axes,
                self.semi_minor_axes,
                self.phis,
            )
        ]


class ShearYX2D(aa.VectorYX2D, AbstractShearField):
    """
    A regular shear field, which is collection of (y,x) vectors which are located on a uniform grid
    of (y,x) coordinates.

    The structure of this data structure is described in
    `autoarray.structures.vectors.uniform.VectorYX2D`

    This class extends `VectorYX2D` to include methods that are specific to a shear field, typically
    used for weak lensing calculations.

    Parameters
    ----------
    vectors
        The 2D (y,x) vectors on an irregular grid that represent the vector-field.
    grid
        The irregular grid of (y,x) coordinates where each vector is located.
    """


class ShearYX2DIrregular(aa.VectorYX2DIrregular, AbstractShearField):

    """
    An irregular shear field, which is collection of (y,x) vectors which are located on an irregular grid
    of (y,x) coordinates.

    The structure of this data structure is described in
    `autoarray.structures.vectors.irregular.VectorYX2DIrregular`

    This class extends `VectorYX2DIrregular` to include methods that are specific to a shear field, typically
    used for weak lensing calculations.

    Parameters
    ----------
    vectors
        The 2D (y,x) vectors on an irregular grid that represent the vector-field.
    grid
        The irregular grid of (y,x) coordinates where each vector is located.
    """

If we treat this vector field as a set of weak lensing shear measurements, the galaxy ellipticity each vector
         corresponds too. If we treat this vector field as a set of weak lensing shear measurements, the semi-major axis of each
         galaxy ellipticity that each vector corresponds too. If we treat this vector field as a set of weak lensing shear measurements, the semi-minor axis of each
         galaxy ellipticity that each vector corresponds too. If we treat this vector field as a set of weak lensing shear measurements, the position angle defined
         counter clockwise from the positive x-axis of each galaxy ellipticity that each vector corresponds too. If we treat this vector field as a set of weak lensing shear measurements, the elliptical patch representing
         each galaxy ellipticity. This patch is used for visualizing an ellipse of each galaxy in an image. A regular shear field, which is collection of (y,x) vectors which are located on a uniform grid
     of (y,x) coordinates.
 
     The structure of this data structure is described in
     `autoarray.structures.vectors.uniform.VectorYX2D`
 
     This class extends `VectorYX2D` to include methods that are specific to a shear field, typically
     used for weak lensing calculations.
 
     Parameters
     ----------
     vectors
         The 2D (y,x) vectors on an irregular grid that represent the vector-field.
     grid
         The irregular grid of (y,x) coordinates where each vector is located. An irregular shear field, which is collection of (y,x) vectors which are located on an irregular grid
     of (y,x) coordinates.
 
     The structure of this data structure is described in
     `autoarray.structures.vectors.irregular.VectorYX2DIrregular`
 
     This class extends `VectorYX2DIrregular` to include methods that are specific to a shear field, typically
     used for weak lensing calculations.
 
     Parameters
     ----------
     vectors
         The 2D (y,x) vectors on an irregular grid that represent the vector-field.
     grid
         The irregular grid of (y,x) coordinates where each vector is located.