# Licensed under a 3-clause BSD style license - see LICENSE.rst

"""Functions for performing aperture photometry on 2-D arrays."""
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)
import math
import abc
import numpy as np
import warnings
import astropy.units as u
from astropy.io import fits
from astropy.table import Table
from astropy.wcs import WCS
from astropy.coordinates import SkyCoord
from astropy.extern import six
from astropy.utils.misc import InheritDocstrings
from astropy.utils.exceptions import AstropyUserWarning
from .aperture_funcs import do_circular_photometry, do_elliptical_photometry
from .wcsutils import skycoord_to_pixel, skycoord_to_pixel_scale_angle, assert_angle_or_pixel, assert_angle

__all__ = ['Aperture', 'SkyAperture', 'PixelAperture',
           'SkyCircularAperture', 'CircularAperture',
           'SkyCircularAnnulus', 'CircularAnnulus',
           'SkyEllipticalAperture', 'EllipticalAperture',
           'SkyEllipticalAnnulus', 'EllipticalAnnulus',
           'RectangularAperture', 'aperture_photometry']

def _make_annulus_path(patch_inner, patch_outer):
    import matplotlib.path as mpath
    verts_inner = patch_inner.get_verts()
    verts_outer = patch_outer.get_verts()
    codes_inner = (np.ones(len(verts_inner), dtype=mpath.Path.code_type) *
                   mpath.Path.LINETO)
    codes_inner[0] = mpath.Path.MOVETO
    codes_outer = (np.ones(len(verts_outer), dtype=mpath.Path.code_type) *
                   mpath.Path.LINETO)
    codes_outer[0] = mpath.Path.MOVETO
    codes = np.concatenate((codes_inner, codes_outer))
    verts = np.concatenate((verts_inner, verts_outer[::-1]))
    return mpath.Path(verts, codes)

class _ABCMetaAndInheritDocstrings(InheritDocstrings, abc.ABCMeta):
    pass

@six.add_metaclass(_ABCMetaAndInheritDocstrings)
class Aperture(object):
    """
    Abstract base class for all apertures.
    """

@six.add_metaclass(_ABCMetaAndInheritDocstrings)
class SkyAperture(Aperture):
    """
    Abstract base class for 2-d apertures defined in celestial coordinates.
    """

@six.add_metaclass(_ABCMetaAndInheritDocstrings)
class PixelAperture(Aperture):
    """
    Abstract base class for 2-d apertures defined in pixel coordinates.

    Derived classes should contain whatever internal data is needed to
    define the aperture, and provide methods `do_photometry` and `extent`
    (and optionally, ``area``).
    """

    @abc.abstractmethod
    def plot(self, ax=None, fill=False, **kwargs):
        """
        Plot the aperture(s) on a matplotlib Axes instance.

        Parameters
        ----------
        ax : `matplotlib.axes.Axes` instance, optional
            If `None`, then the current ``Axes`` instance is used.

        fill : bool, optional
            Set whether to fill the aperture patch.  The default is
            `False`.

        kwargs
            Any keyword arguments accepted by `matplotlib.patches.Patch`.
        """

    @abc.abstractmethod
    def do_photometry(self, data, error=None, gain=None, pixelwise_error=True,
                      method='exact', subpixels=5):
        """Sum flux within aperture(s).

        Parameters
        ----------
        data : array_like
            The 2-d array on which to perform photometry.
        error : array_like, optional
            Error in each pixel, interpreted as Gaussian 1-sigma uncertainty.
            ``error`` has to have the same shape as ``data``.
        gain : array_like, optional
            Ratio of counts (e.g., electrons or photons) to units of the
            data (e.g., ADU), for the purpose of calculating Poisson error from
            the object itself. ``gain`` has to have the same shape as ``data``.
        pixelwise_error : bool, optional
            For error and/or gain arrays. If `True`, assume error and/or gain
            vary significantly within an aperture: sum contribution from each
            pixel. If `False`, assume error and gain do not vary significantly
            within an aperture.
        method : str, optional
            Method to threat masked pixels. For the currently supported methods
            see the documentation of `aperture_photometry`.
        subpixels : int, optional
            For the ``'subpixel'`` method, resample pixels by this factor (in
            each dimension). That is, each pixel is divided into
            ``subpixels ** 2`` subpixels.

        Returns
        -------
        flux : `~astropy.units.Quantity`
            Sum of the values withint the aperture(s). Unit is kept to be the
            unit of the input ``data``.

        fluxvar :  `~astropy.units.Quantity`
            Corresponting uncertainity in the ``flux`` values. Returned only
            if input ``error`` is not `None`.
        """

    def area():
        """
        Area of aperture.

        Returns
        -------
        area : float
            Area of aperture.
        """

class SkyCircularAperture(SkyAperture):
    """
    Circular aperture(s), defined in sky coordinates.

    Parameters
    ----------
    positions : `~astropy.coordinates.SkyCoord`
        Celestial coordinates of the aperture center(s). This can be either
        scalar coordinates or an array of coordinates.
    r : `~astropy.units.Quantity`
        The radius of the aperture(s), either in angular or pixel units.
    """

    def __init__(self, positions, r):

        if isinstance(positions, SkyCoord):
            self.positions = positions
        else:
            raise TypeError("positions should be a SkyCoord instance")

        assert_angle_or_pixel('r', r)
        self.r = r

    def to_pixel(self, wcs):
        """
        Return a CircularAperture instance in pixel coordinates.
        """

        if self.r.unit.physical_type == 'angle':
            x, y, scale, angle = skycoord_to_pixel_scale_angle(self.positions, wcs)
            # TODO: no need to use the mean once we support arrays of aperture values
            r = (np.mean(scale) * self.r).to(u.pixel).value
        else:  # pixel
            x, y = skycoord_to_pixel(self.positions, wcs)
            r = self.r.value

        pixel_positions = np.array([x, y]).transpose()

        return CircularAperture(pixel_positions, r)

class CircularAperture(PixelAperture):
    """
    Circular aperture(s), defined in pixel coordinates.

    Parameters
    ----------
    positions : tuple, list, array, or `~astropy.units.Quantity`
        Pixel coordinates of the aperture center(s), either as a single
        ``(x, y)`` tuple, a list of ``(x, y)`` tuples, an ``Nx2`` Numpy
        array, or an ``Nx2`` `~astropy.units.Quantity` in units of pixels.
    r : float
        The radius of the aperture(s), in pixels.

    Raises
    ------
    ValueError : `~.exceptions.ValueError`
        If the radius is negative.
    """

    def __init__(self, positions, r):

        try:
            self.r = float(r)
        except TypeError:
            raise TypeError('r must be numeric, received {0}'.format(type(r)))

        if r < 0:
            raise ValueError('r must be non-negative')

        if isinstance(positions, u.Quantity):
            if positions.unit is u.pixel:
                positions = positions.value
            else:
                raise u.UnitsError("positions should be in pixel units")

        self.positions = np.asarray(positions)

        if self.positions.ndim == 1 and len(self.positions) == 2:
            self.positions = np.atleast_2d(positions)
        elif self.positions.ndim != 2 or self.positions.shape[1] != 2:
            raise TypeError("Expected (x, y) tuple, a list of (x, y) "
                            "tuples, or an Nx2 array, got {0}".format(positions))

    def area(self):
        return math.pi * self.r ** 2

    def plot(self, ax=None, fill=False, source_id=None, **kwargs):

        import matplotlib.pyplot as plt
        import matplotlib.patches as mpatches

        kwargs['fill'] = fill

        if ax is None:
            ax = plt.gca()
        if source_id is None:
            positions = self.positions
        else:
            positions = self.positions[np.atleast_1d(source_id)]

        for position in positions:
            patch = mpatches.Circle(position, self.r, **kwargs)
            ax.add_patch(patch)

    def do_photometry(self, data, error=None, gain=None, pixelwise_error=True,
                      method='exact', subpixels=5):

        if method not in ('center', 'subpixel', 'exact'):
            raise ValueError('{0} method not supported for aperture class '
                             '{1}'.format(method, self.__class__.__name__))

        flux = do_circular_photometry(data, self.positions,
                                      self.r, error=error, gain=gain,
                                      pixelwise_error=pixelwise_error,
                                      method=method,
                                      subpixels=subpixels)
        return flux

class SkyCircularAnnulus(SkyAperture):
    """
    Circular annulus aperture(s), defined in sky coordinates.

    Parameters
    ----------

    positions : `~astropy.coordinates.SkyCoord`
        Celestial coordinates of the aperture center(s). This can be either
        scalar coordinates or an array of coordinates.
    r : `~astropy.units.Quantity`
        The radius of the aperture(s), either in angular or pixel units.
    """

    def __init__(self, positions, r_in, r_out):

        if isinstance(positions, SkyCoord):
            self.positions = positions
        else:
            raise TypeError("positions should be a SkyCoord instance")

        assert_angle_or_pixel('r_in', r_in)
        assert_angle_or_pixel('r_out', r_out)

        if r_in.unit.physical_type != r_out.unit.physical_type:
            raise ValueError("r_in and r_out should either both be angles or in pixels")

        self.r_in = r_in
        self.r_out = r_out

    def to_pixel(self, wcs):
        """
        Return a CircularAnnulus instance in pixel coordinates.
        """

        if self.r_in.unit.physical_type == 'angle':
            x, y, scale, angle = skycoord_to_pixel_scale_angle(self.positions, wcs)
            # TODO: no need to use the mean once we support arrays of aperture values
            r_in = (np.mean(scale) * self.r_in).to(u.pixel).value
            r_out = (np.mean(scale) * self.r_out).to(u.pixel).value
        else:  # pixel
            x, y = skycoord_to_pixel(self.positions, wcs)
            r_in = self.r_in.value
            r_out = self.r_out.value

        pixel_positions = np.array([x, y]).transpose()

        return CircularAnnulus(pixel_positions, r_in, r_out)

class CircularAnnulus(PixelAperture):
    """
    Circular annulus aperture(s), defined in pixel coordinates.

    Parameters
    ----------
    positions : tuple, list, array, or `~astropy.units.Quantity`
        Pixel coordinates of the aperture center(s), either as a single
        ``(x, y)`` tuple, a list of ``(x, y)`` tuples, an ``Nx2`` Numpy
        array, or an ``Nx2`` `~astropy.units.Quantity` in units of pixels.
    r_in : float
        The inner radius of the annulus.
    r_out : float
        The outer radius of the annulus.

    Raises
    ------
    ValueError : `~.exceptions.ValueError`
        If inner radius (``r_in``) is greater than outer radius (``r_out``).
    ValueError : `~.exceptions.ValueError`
        If inner radius is negative.
    """

    def __init__(self, positions, r_in, r_out):
        try:
            self.r_in = r_in
            self.r_out = r_out
        except TypeError:
            raise TypeError("'r_in' and 'r_out' must be numeric, received {0} "
                            "and {1}".format((type(r_in), type(r_out))))

        if not (r_out > r_in):
            raise ValueError('r_out must be greater than r_in')
        if r_in < 0:
            raise ValueError('r_in must be non-negative')

        if isinstance(positions, u.Quantity):
            positions = positions.value
        if isinstance(positions, (list, tuple, np.ndarray)):
            self.positions = np.atleast_2d(positions)
        else:
            raise TypeError("List or array of (x,y) pixel coordinates is "
                            "expected got '{0}'.".format(positions))

        if self.positions.ndim > 2:
            raise ValueError('{0}-d position array not supported. Only 2-d '
                             'arrays supported.'.format(self.positions.ndim))

    def area(self):
        return math.pi * (self.r_out ** 2 - self.r_in ** 2)

    def do_photometry(self, data, error=None, gain=None, pixelwise_error=True,
                      method='exact', subpixels=5):

        if method not in ('center', 'subpixel', 'exact'):
            raise ValueError('{0} method not supported for aperture class '
                             '{1}'.format(method, self.__class__.__name__))

        flux = do_circular_photometry(data, self.positions,
                                      self.r_out, error=error, gain=gain,
                                      pixelwise_error=pixelwise_error,
                                      method=method,
                                      subpixels=subpixels,
                                      r_in=self.r_in)

        return flux

    def plot(self, ax=None, fill=False, source_id=None, **kwargs):

        import matplotlib.pyplot as plt
        import matplotlib.patches as mpatches

        kwargs['fill'] = fill

        if ax is None:
            ax = plt.gca()

        if source_id is None:
            positions = self.positions
        else:
            positions = self.positions[np.atleast_1d(source_id)]

        resolution = 20

        for position in positions:
            patch_inner = mpatches.CirclePolygon(position, self.r_in,
                                                 resolution=resolution)
            patch_outer = mpatches.CirclePolygon(position, self.r_out,
                                                 resolution=resolution)
            path = _make_annulus_path(patch_inner, patch_outer)
            patch = mpatches.PathPatch(path, **kwargs)
            ax.add_patch(patch)

class SkyEllipticalAperture(SkyAperture):
    """
    Elliptical aperture(s), defined in sky coordinates.

    Parameters
    ----------
    positions : `~astropy.coordinates.SkyCoord`
        Celestial coordinates of the aperture center(s). This can be either
        scalar coordinates or an array of coordinates.
    a : `~astropy.units.Quantity`
        The semimajor axis, either in angular or pixel units.
    b : `~astropy.units.Quantity`
        The semiminor axis, either in angular or pixel units.
    theta : `~astropy.units.Quantity`
        The position angle of the semimajor axis (counterclockwise), either
        in angular or pixel units.
    """

    def __init__(self, positions, a, b, theta):

        if isinstance(positions, SkyCoord):
            self.positions = positions
        else:
            raise TypeError("positions should be a SkyCoord instance")

        assert_angle_or_pixel('a', a)
        assert_angle_or_pixel('b', b)
        assert_angle('theta', theta)

        if a.unit.physical_type != b.unit.physical_type:
            raise ValueError("a and b should either both be angles or in pixels")

        self.a = a
        self.b = b
        self.theta = theta

    def to_pixel(self, wcs):
        """
        Return a EllipticalAperture instance in pixel coordinates.
        """

        x, y, scale, angle = skycoord_to_pixel_scale_angle(self.positions, wcs)

        # TODO: no need to use the mean once we support arrays of aperture values
        if self.a.unit.physical_type == 'angle':
            a = (np.mean(scale) * self.a).to(u.pixel).value
            b = (np.mean(scale) * self.b).to(u.pixel).value
        else:  # pixel
            a = self.a.value
            b = self.b.value

        theta = (angle + self.theta).to(u.radian).value

        pixel_positions = np.array([x, y]).transpose()

        return EllipticalAperture(pixel_positions, a, b, theta)

class EllipticalAperture(PixelAperture):
    """
    Elliptical aperture(s), defined in pixel coordinates.

    Parameters
    ----------
    positions : tuple, list, array, or `~astropy.units.Quantity`
        Pixel coordinates of the aperture center(s), either as a single
        ``(x, y)`` tuple, a list of ``(x, y)`` tuples, an ``Nx2`` Numpy
        array, or an ``Nx2`` `~astropy.units.Quantity` in units of pixels.
    a : float
        The semimajor axis.
    b : float
        The semiminor axis.
    theta : float
        The position angle of the semimajor axis in radians
        (counterclockwise).

    Raises
    ------
    ValueError : `~.exceptions.ValueError`
        If either axis (``a`` or ``b``) is negative.

    """

    def __init__(self, positions, a, b, theta):
        try:
            self.a = float(a)
            self.b = float(b)
            self.theta = float(theta)
        except TypeError:
            raise TypeError("'a' and 'b' and 'theta' must be numeric, received"
                            "{0} and {1} and {2}."
                            .format((type(a), type(b), type(theta))))

        if a < 0 or b < 0:
            raise ValueError("'a' and 'b' must be non-negative.")

        if isinstance(positions, u.Quantity):
            positions = positions.value
        if isinstance(positions, (list, tuple, np.ndarray)):
            self.positions = np.atleast_2d(positions)
        else:
            raise TypeError("List or array of (x,y) pixel coordinates is "
                            "expected got '{0}'.".format(positions))

        if self.positions.ndim > 2:
            raise ValueError('{0}-d position array not supported. Only 2-d '
                             'arrays supported.'.format(self.positions.ndim))

    def area(self):
        return math.pi * self.a * self.b

    def do_photometry(self, data, error=None, gain=None, pixelwise_error=True,
                      method='exact', subpixels=5):

        if method not in ('center', 'subpixel', 'exact'):
            raise ValueError('{0} method not supported for aperture class '
                             '{1}'.format(method, self.__class__.__name__))

        flux = do_elliptical_photometry(data, self.positions,
                                        self.a, self.b, self.theta,
                                        error=error, gain=gain,
                                        pixelwise_error=pixelwise_error,
                                        method=method,
                                        subpixels=subpixels)
        return flux

    def plot(self, ax=None, fill=False, source_id=None, **kwargs):

        import matplotlib.pyplot as plt
        import matplotlib.patches as mpatches

        kwargs['fill'] = fill

        if ax is None:
            ax = plt.gca()

        if source_id is None:
            positions = self.positions
        else:
            positions = self.positions[np.atleast_1d(source_id)]

        theta_deg = self.theta * 180. / np.pi
        for position in positions:
            patch = mpatches.Ellipse(position, 2.*self.a, 2.*self.b,
                                     theta_deg, **kwargs)
            ax.add_patch(patch)

class SkyEllipticalAnnulus(SkyAperture):
    """
    Elliptical annulus aperture(s), defined in sky coordinates.

    Parameters
    ----------
    positions : `~astropy.coordinates.SkyCoord`
        Celestial coordinates of the aperture center(s). This can be either
        scalar coordinates or an array of coordinates.
    a_in : `~astropy.units.Quantity`
        The inner semimajor axis, either in angular or pixel units.
    a_out : `~astropy.units.Quantity`
        The outer semimajor axis, either in angular or pixel units.
    b_out : `~astropy.units.Quantity`
        The outer semiminor axis, either in angular or pixel units. The inner
        semiminor axis is determined by scaling by a_in/a_out.
    theta : `~astropy.units.Quantity`
        The position angle of the semimajor axis (counterclockwise), either
        in angular or pixel units.
    """

    def __init__(self, positions, a_in, a_out, b_out, theta):

        if isinstance(positions, SkyCoord):
            self.positions = positions
        else:
            raise TypeError("positions should be a SkyCoord instance")

        assert_angle_or_pixel('a_in', a_in)
        assert_angle_or_pixel('a_out', a_out)
        assert_angle_or_pixel('b_out', b_out)
        assert_angle('theta', theta)

        if a_in.unit.physical_type != a_out.unit.physical_type:
            raise ValueError("a_in and a_out should either both be angles or in pixels")

        if a_out.unit.physical_type != b_out.unit.physical_type:
            raise ValueError("a_out and b_out should either both be angles or in pixels")

        self.a_in = a_in
        self.a_out = a_out
        self.b_out = b_out
        self.theta = theta

    def to_pixel(self, wcs):
        """
        Return a EllipticalAnnulus instance in pixel coordinates.
        """

        x, y, scale, angle = skycoord_to_pixel_scale_angle(self.positions, wcs)

        # TODO: no need to use the mean once we support arrays of aperture values
        if self.a_in.unit.physical_type == 'angle':
            a_in = (np.mean(scale) * self.a_in).to(u.pixel).value
            a_out = (np.mean(scale) * self.a_out).to(u.pixel).value
            b_out = (np.mean(scale) * self.b_out).to(u.pixel).value
        else:
            a_in = self.a_in.value
            a_out = self.a_out.value
            b_out = self.b_out.value

        theta = (angle + self.theta).to(u.radian).value

        pixel_positions = np.array([x, y]).transpose()

        return EllipticalAnnulus(pixel_positions, a_in, a_out, b_out, theta)

class EllipticalAnnulus(PixelAperture):
    """
    Elliptical annulus aperture(s), defined in pixel coordinates.

    Parameters
    ----------
    positions : tuple, list, array, or `~astropy.units.Quantity`
        Pixel coordinates of the aperture center(s), either as a single
        ``(x, y)`` tuple, a list of ``(x, y)`` tuples, an ``Nx2`` Numpy
        array, or an ``Nx2`` `~astropy.units.Quantity` in units of pixels.
    a_in : float
        The inner semimajor axis.
    a_out : float
        The outer semimajor axis.
    b_out : float
        The outer semiminor axis. (The inner semiminor axis is determined
        by scaling by a_in/a_out.)
    theta : float
        The position angle of the semimajor axis in radians.
        (counterclockwise).

    Raises
    ------
    ValueError : `~.exceptions.ValueError`
        If inner semimajor axis (``a_in``) is greater than outer semimajor
        axis (``a_out``).
    ValueError : `~.exceptions.ValueError`
        If either the inner semimajor axis (``a_in``) or the outer semiminor
        axis (``b_out``) is negative.
    """

    def __init__(self, positions, a_in, a_out, b_out, theta):
        try:
            self.a_in = float(a_in)
            self.a_out = float(a_out)
            self.b_out = float(b_out)
            self.theta = float(theta)
        except TypeError:
            raise TypeError("'a_in' and 'a_out' and 'b_out' and 'theta' must "
                            "be numeric, received {0} and {1} and {2} and {3}."
                            .format((type(a_in), type(a_out),
                                     type(b_out), type(theta))))

        if not (a_out > a_in):
            raise ValueError("'a_out' must be greater than 'a_in'")
        if a_in < 0 or b_out < 0:
            raise ValueError("'a_in' and 'b_out' must be non-negative")

        self.b_in = a_in * b_out / a_out

        if isinstance(positions, u.Quantity):
            positions = positions.value
        if isinstance(positions, (list, tuple, np.ndarray)):
            self.positions = np.atleast_2d(positions)
        else:
            raise TypeError("List or array of (x,y) pixel coordinates is "
                            "expected got '{0}'.".format(positions))

        if self.positions.ndim > 2:
            raise ValueError('{0}-d position array not supported. Only 2-d '
                             'arrays supported.'.format(self.positions.ndim))

    def area(self):
        return math.pi * (self.a_out * self.b_out - self.a_in * self.b_in)

    def plot(self, ax=None, fill=False, source_id=None, **kwargs):

        import matplotlib.pyplot as plt
        import matplotlib.patches as mpatches

        kwargs['fill'] = fill

        if ax is None:
            ax = plt.gca()

        if source_id is None:
            positions = self.positions
        else:
            positions = self.positions[np.atleast_1d(source_id)]

        theta_deg = self.theta * 180. / np.pi
        for position in positions:
            patch_inner = mpatches.Ellipse(position, 2.*self.a_in,
                                           2.*self.b_in, theta_deg, **kwargs)
            patch_outer = mpatches.Ellipse(position, 2.*self.a_out,
                                           2.*self.b_out, theta_deg, **kwargs)
            path = _make_annulus_path(patch_inner, patch_outer)
            patch = mpatches.PathPatch(path, **kwargs)
            ax.add_patch(patch)

    def do_photometry(self, data, error=None, gain=None, pixelwise_error=True,
                      method='exact', subpixels=5):

        if method not in ('center', 'subpixel', 'exact'):
            raise ValueError('{0} method not supported for aperture class '
                             '{1}'.format(method, self.__class__.__name__))

        flux = do_elliptical_photometry(data, self.positions,
                                        self.a_out, self.b_out, self.theta,
                                        error=error, gain=gain,
                                        pixelwise_error=pixelwise_error,
                                        method=method,
                                        subpixels=subpixels,
                                        a_in=self.a_in)

        return flux

class RectangularAperture(PixelAperture):
    """
    A rectangular aperture.

    Parameters
    ----------
    positions : tuple, or list, or array
        Center coordinates of the apertures as list or array of (x, y)
        pixelcoordinates.
    w : float
        The full width of the aperture (at theta = 0, this is the "x" axis).
    h : float
        The full height of the aperture (at theta = 0, this is the "y" axis).
    theta : float
        The position angle of the semimajor axis in radians
        (counterclockwise).

    Raises
    ------
    ValueError : `~.exceptions.ValueError`
        If either width (``w``) or height (``h``) is negative.
    """

    def __init__(self, positions, w, h, theta):
        try:
            self.w = float(w)
            self.h = float(h)
            self.theta = float(theta)
        except TypeError:
            raise TypeError("'w' and 'h' and 'theta' must "
                            "be numeric, received {0} and {1} and {2}."
                            .format((type(w), type(h), type(theta))))
        if w < 0 or h < 0:
            raise ValueError("'w' and 'h' must be nonnegative.")

        if isinstance(positions, u.Quantity):
            positions = positions.value
        if isinstance(positions, (list, tuple, np.ndarray)):
            self.positions = np.atleast_2d(positions)
        else:
            raise TypeError("List or array of (x,y) pixel coordinates is "
                            "expected got '{0}'.".format(positions))

        if self.positions.ndim > 2:
            raise ValueError('{0}-d position array not supported. Only 2-d '
                             'arrays supported.'.format(self.positions.ndim))

    def extent(self):
        r = max(self.h, self.w) * 2 ** -0.5
        # this is an overestimate by up to sqrt(2) unless theta = 45 deg
        extents = []
        centers = []
        for x, y in self.positions:
            extents.append((int(x - r + 0.5), int(x + r + 1.5),
                            int(y - r + 0.5), int(y + r + 1.5)))
            centers.append((x, x, y, y))

        self._centers = np.array(centers)
        return np.array(extents)

    def area(self):
        return self.w * self.h

    def plot(self, ax=None, fill=False, source_id=None, **kwargs):

        import matplotlib.pyplot as plt
        import matplotlib.patches as mpatches

        kwargs['fill'] = fill

        if ax is None:
            ax = plt.gca()

        if source_id is None:
            positions = self.positions
        else:
            positions = self.positions[np.atleast_1d(source_id)]

        hw = self.w / 2.
        hh = self.h / 2.
        sint = math.sin(self.theta)
        cost = math.cos(self.theta)
        dx = (hh * sint) - (hw * cost)
        dy = -(hh * cost) - (hw * sint)
        positions = positions + np.array([dx, dy])
        theta_deg = self.theta * 180. / np.pi
        for position in positions:
            patch = mpatches.Rectangle(position, self.w, self.h, theta_deg,
                                       **kwargs)
            ax.add_patch(patch)

    def do_photometry(self, data, error=None, gain=None, pixelwise_error=True,
                      method='subpixel', subpixels=5):

        if method == 'exact':
            warnings.warn("'exact' method is not implemented, defaults to "
                          "'subpixel' instead", AstropyUserWarning)
            method = 'subpixel'

        from .aperture_funcs import get_phot_extents
        ood_filter, pixel_extent, phot_extent = get_phot_extents(data, self.positions, self.extent())

        if method not in ('center', 'subpixel', 'exact'):
            raise ValueError('{0} method not supported for aperture class '
                             '{1}'.format(method, self.__class__.__name__))

        x_min, x_max, y_min, y_max = pixel_extent
        x_pmin, x_pmax, y_pmin, y_pmax = phot_extent

        flux = u.Quantity(np.zeros(len(self.positions), dtype=np.float),
                          unit=data.unit)

        # Check for invalid aperture
        if self.w == 0 or self.h == 0:
            return (flux, )

        # TODO: flag these objects
        if np.sum(ood_filter):
            flux[ood_filter] = np.nan
            warnings.warn("The aperture at position {0} does not have any "
                          "overlap with the data"
                          .format(self.positions[ood_filter]),
                          AstropyUserWarning)
            if np.sum(ood_filter) == len(self.positions):
                return (flux, )

        if error is not None:
            fluxvar = u.Quantity(np.zeros(len(self.positions), dtype=np.float),
                                 unit=error.unit ** 2)

        if method in ('center', 'subpixel'):
            if method == 'center': subpixels = 1
            if method == 'subpixel': from .extern.imageutils import downsample

            for i in range(len(flux)):
                x_size = ((x_pmax[i] - x_pmin[i]) /
                          (data[:, x_min[i]:x_max[i]].shape[1] * subpixels))
                y_size = ((y_pmax[i] - y_pmin[i]) /
                          (data[y_min[i]:y_max[i], :].shape[0] * subpixels))

                x_centers = np.arange(x_pmin[i] + x_size / 2.,
                                      x_pmax[i], x_size)
                y_centers = np.arange(y_pmin[i] + y_size / 2.,
                                      y_pmax[i], y_size)

                xx, yy = np.meshgrid(x_centers, y_centers)

                newx = (xx * math.cos(self.theta) +
                        yy * math.sin(self.theta))
                newy = (yy * math.cos(self.theta) -
                        xx * math.sin(self.theta))

                halfw = self.w / 2
                halfh = self.h / 2
                in_aper = (((-halfw < newx) & (newx < halfw) &
                            (-halfh < newy) & (newy < halfh)).astype(float)
                           / subpixels ** 2)

                if method == 'center':
                    if not np.isnan(flux[i]):
                        flux[i] = np.sum(data[y_min[i]:y_max[i],
                                              x_min[i]:x_max[i]] * in_aper)
                        if error is not None:
                            if pixelwise_error:
                                subvariance = error[y_min[i]:y_max[i],
                                                    x_min[i]:x_max[i]] ** 2
                                if gain is not None:
                                    subvariance += (data[y_min[i]:y_max[i],
                                                         x_min[i]:x_max[i]] /
                                                    gain[y_min[i]:y_max[i],
                                                         x_min[i]:x_max[i]])
                                # Make sure variance is > 0
                                fluxvar[i] = max(np.sum(subvariance * in_aper), 0)
                            else:
                                local_error = error[int((y_min[i] + y_max[i]) / 2 + 0.5),
                                                    int((x_min[i] + x_max[i]) / 2 + 0.5)]
                                fluxvar[i] = max(local_error ** 2 * np.sum(in_aper), 0)
                                if gain is not None:
                                    local_gain = gain[int((y_min[i] + y_max[i]) / 2 + 0.5),
                                                      int((x_min[i] + x_max[i]) / 2 + 0.5)]
                                    fluxvar[i] += flux[i] / local_gain
                else:
                    if not np.isnan(flux[i]):
                        if error is None:
                            flux[i] = np.sum(data[y_min[i]:y_max[i],
                                                  x_min[i]:x_max[i]] *
                                             downsample(in_aper, subpixels))
                        else:
                            fraction = downsample(in_aper, subpixels)
                            flux[i] = np.sum(data[y_min[i]:y_max[i],
                                                  x_min[i]:x_max[i]] * fraction)

                            if pixelwise_error:
                                subvariance = error[y_min[i]:y_max[i],
                                                    x_min[i]:x_max[i]] ** 2
                                if gain is not None:
                                    subvariance += (data[y_min[i]:y_max[i],
                                                         x_min[i]:x_max[i]] /
                                                    gain[y_min[i]:y_max[i],
                                                         x_min[i]:x_max[i]])
                                # Make sure variance is > 0
                                fluxvar[i] = max(np.sum(subvariance * fraction), 0)
                            else:
                                local_error = error[int((y_min[i] + y_max[i]) / 2 + 0.5),
                                                    int((x_min[i] + x_max[i]) / 2 + 0.5)]
                                fluxvar[i] = max(local_error ** 2 * np.sum(fraction), 0)
                                if gain is not None:
                                    local_gain = gain[int((y_min[i] + y_max[i]) / 2 + 0.5),
                                                      int((x_min[i] + x_max[i]) / 2 + 0.5)]
                                    fluxvar[i] += flux[i] / local_gain

        elif method == 'exact':
            raise NotImplementedError("'exact' method not yet supported for "
                                      "RectangularAperture")
        else:
            raise ValueError('{0} method not supported for aperture class '
                             '{1}'.format(method, self.__class__.__name__))

        if error is None:
            return (flux, )
        else:
            return (flux, np.sqrt(fluxvar))

def aperture_photometry(data, apertures, unit=None, wcs=None,
                        error=None, gain=None, mask=None, method='exact',
                        subpixels=5, pixelwise_error=True,
                        mask_method='skip'):
    """
    Sum flux within an aperture at the given position(s).

    Parameters
    ----------
    data : array_like, `~astropy.io.fits.ImageHDU`, `~astropy.io.fits.HDUList`
        The 2-d array on which to perform photometry. Units are used during
        the photometry, either provided along with the data array, or stored
        in the header keyword ``'BUNIT'``.
    apertures : `~photutils.Aperture` instance
        The apertures to use for the photometry.
    unit : `~astropy.units.UnitBase` instance, str
        An object that represents the unit associated with ``data``.  Must
        be an `~astropy.units.UnitBase` object or a string parseable by the
        :mod:`~astropy.units` package. An error is raised if ``data``
        already has a different unit.
    wcs : `~astropy.wcs.WCS`, optional
        Use this as the wcs transformation. It overrides any wcs transformation
        passed along with ``data`` either in the header or in an attribute.
    error : float or array_like, optional
        Error in each pixel, interpreted as Gaussian 1-sigma uncertainty.
    gain : float or array_like, optional
        Ratio of counts (e.g., electrons or photons) to units of the data
        (e.g., ADU), for the purpose of calculating Poisson error from the
        object itself. If ``gain`` is `None` (default), ``error`` is assumed to
        include all uncertainty in each pixel. If ``gain`` is given, ``error``
        is assumed to be the "background error" only (not accounting for
        Poisson error in the flux in the apertures).
    mask : array_like (bool), optional
        Mask to apply to the data.
    method : str, optional
        Method to use for determining overlap between the aperture and pixels.
        Options include ['center', 'subpixel', 'exact'], but not all options
        are available for all types of apertures. More precise methods will
        generally be slower.

        * ``'center'``
            A pixel is considered to be entirely in or out of the aperture
            depending on whether its center is in or out of the aperture.
        * ``'subpixel'``
            A pixel is divided into subpixels and the center of each
            subpixel is tested (as above). With ``subpixels`` set to 1, this
            method is equivalent to 'center'. Note that for subpixel
            sampling, the input array is only resampled once for each
            object.
        * ``'exact'`` (default)
            The exact overlap between the aperture and each pixel is
            calculated.
    subpixels : int, optional
        For the ``'subpixel'`` method, resample pixels by this factor (in
        each dimension). That is, each pixel is divided into
        ``subpixels ** 2`` subpixels.
    pixelwise_error : bool, optional
        For error and/or gain arrays. If `True`, assume error and/or gain
        vary significantly within an aperture: sum contribution from each
        pixel. If `False`, assume error and gain do not vary significantly
        within an aperture. Use the single value of error and/or gain at
        the center of each aperture as the value for the entire aperture.
        Default is `True`.
    mask_method : str, optional
        Method to treat masked pixels. Currently supported methods:

        * ``'skip'``
            Leave out the masked pixels from all calculations.
        * ``'interpolation'``
            The value of the masked pixels are replaced by the mean value of
            the neighbouring non-masked pixels.

    Returns
    -------
    phot_table : `~astropy.table.Table`
        A table of the photometry with the following columns:

        * ``'aperture_sum'``: Sum of the values within the aperture.
        * ``'aperture_sum_err'``: Corresponding uncertainty in
          ``'aperture_sum'`` values.  Returned only if input ``error`` is not
          `None`.
        * ``'pixel_center'``: pixel coordinate pairs of the center of the
          apertures. Unit is pixel.
        * ``'input_center'``: input coordinate pairs as they were given in the
          ``positions`` parameter.

        The metadata of the table stores the version numbers of both astropy
        and photutils, as well as the calling arguments.
    """
    dataunit = None
    datamask = None
    wcs_transformation = wcs

    if isinstance(data, (fits.PrimaryHDU, fits.ImageHDU)):
        header = data.header
        data = data.data

        if 'BUNIT' in header:
            dataunit = header['BUNIT']

        # TODO check how a mask can be stored in the header, it seems like
        # full pixel masks are not supported by the FITS standard, look for
        # real life examples (e.g. header value stores the fits number of
        # fits extension where the pixelmask is stored?)
        if 'MASK' in header:
            datamask = header.mask

    elif isinstance(data, fits.HDUList):
        # TODO: do it in a 2d array, and thus get the light curves as a
        # side-product? Although it's not usual to store time series as
        # HDUList

        for i in range(len(data)):
            if data[i].data is not None:
                warnings.warn("Input data is a HDUList object, photometry is "
                              "only run for the {0}. HDU."
                              .format(i), AstropyUserWarning)
                return aperture_photometry(data[i], apertures, unit,
                                           wcs, error, gain, mask, method,
                                           subpixels,
                                           pixelwise_error, mask_method)

    # this is basically for NDData inputs and alike
    elif hasattr(data, 'data') and not isinstance(data, np.ndarray):
        if data.wcs is not None and wcs_transformation is None:
            wcs_transformation = data.wcs
        datamask = data.mask

    if hasattr(data, 'unit'):
        dataunit = data.unit

    if unit is not None and dataunit is not None:
        if unit != dataunit:
            raise u.UnitsError('Unit of input data ({0}) and unit given by '
                               'unit argument ({1}) are not identical.'.
                               format(dataunit, unit))
        data = u.Quantity(data, unit=dataunit, copy=False)
    elif unit is None:
        if dataunit is not None:
            data = u.Quantity(data, unit=dataunit, copy=False)
        else:
            data = u.Quantity(data, copy=False)
    else:
        data = u.Quantity(data, unit=unit, copy=False)

    if datamask is None:
        data.mask = datamask

    # Check input array type and dimension.
    if np.iscomplexobj(data):
        raise TypeError('Complex type not supported')
    if data.ndim != 2:
        raise ValueError('{0}-d array not supported. '
                         'Only 2-d arrays supported.'.format(data.ndim))

    # Deal with the mask if it exist
    if mask is not None or datamask is not None:
        if mask is None:
            mask = datamask
        else:
            mask = np.asarray(mask)
            if np.iscomplexobj(mask):
                raise TypeError('Complex type not supported')
            if mask.ndim != 2:
                raise ValueError('{0}-d array not supported. '
                                 'Only 2-d arrays supported.'
                                 .format(mask.ndim))
            if mask.shape != data.shape:
                raise ValueError('Shapes of mask array and data array '
                                 'must match')

            if datamask is not None:
                mask *= datamask

        if mask_method == 'skip':
            data *= ~mask

        if mask_method == 'interpolation':
            for i, j in zip(*np.nonzero(mask)):
                y0, y1 = max(i - 1, 0), min(i + 2, data.shape[0])
                x0, x1 = max(j - 1, 0), min(j + 2, data.shape[1])
                data[i, j] = np.mean(data[y0:y1, x0:x1][~mask[y0:y1, x0:x1]])

    # Check whether we really need to calculate pixelwise errors, even if
    # requested. (If neither error nor gain is an array, we don't need to.)
    if ((error is None) or
        (np.isscalar(error) and gain is None) or
        (np.isscalar(error) and np.isscalar(gain))):
        pixelwise_error = False

    # Check error shape.
    if error is not None:
        if isinstance(error, u.Quantity):
            if np.isscalar(error.value):
                error = u.Quantity(np.broadcast_arrays(error, data),
                                   unit=error.unit)[0]
        elif np.isscalar(error):
            error = u.Quantity(np.broadcast_arrays(error, data),
                               unit=data.unit)[0]
        else:
            error = u.Quantity(error, unit=data.unit, copy=False)

        if error.shape != data.shape:
            raise ValueError('shapes of error array and data array must'
                             ' match')

    # Check gain shape.
    if gain is not None:
        # Gain doesn't do anything without error set, so raise an exception.
        # (TODO: instead, should we just set gain = None and ignore it?)
        if error is None:
            raise ValueError('gain requires error')

        if isinstance(gain, u.Quantity):
            if np.isscalar(gain.value):
                gain = u.Quantity(np.broadcast_arrays(gain, data),
                                  unit=gain.unit)[0]

        elif np.isscalar(gain):
            gain = np.broadcast_arrays(gain, data)[0]
        if gain.shape != data.shape:
            raise ValueError('shapes of gain array and data array must match')

    # Check that 'subpixels' is an int and is 1 or greater.
    if method == 'subpixel':
        subpixels = int(subpixels)
        if subpixels < 1:
            raise ValueError('subpixels: an integer greater than 0 is '
                             'required')

    ap = apertures

    if isinstance(apertures, SkyAperture):
        positions = ap.positions
        if wcs_transformation is None:
            wcs_transformation = WCS(header)
        ap = ap.to_pixel(wcs_transformation)
        pixelpositions = ap.positions
    else:
        positions = ap.positions * u.pixel
        pixelpositions = ap.positions * u.pixel

    # Prepare version return data
    from astropy import __version__
    astropy_version = __version__

    from photutils import __version__
    photutils_version = __version__

    photometry_result = ap.do_photometry(data, method=method,
                                         subpixels=subpixels,
                                         error=error, gain=gain,
                                         pixelwise_error=pixelwise_error)
    if error is None:
        phot_col_names = ('aperture_sum', )
    else:
        phot_col_names = ('aperture_sum', 'aperture_sum_err')

    # Note: if wcs_transformation is None, 'pixel_center' will be the same
    # as 'input_center'

    # check whether single or multiple positions
    if len(pixelpositions) > 1 and pixelpositions[0].size >= 2:
        coord_columns = (pixelpositions, positions)
    else:
        coord_columns = ((pixelpositions,), (positions,))

    coord_col_names = ('pixel_center', 'input_center')

    return Table(data=(photometry_result + coord_columns),
                 names=(phot_col_names + coord_col_names),
                 meta={'name': 'Aperture photometry results',
                       'version': 'astropy: {0}, photutils: {1}'
                       .format(astropy_version, photutils_version),
                       'calling_args': ('method={0}, subpixels={1}, '
                                        'error={2}, gain={3}, '
                                        'pixelwise_error={4}')
                       .format(method, subpixels, error is not None,
                               gain is not None, pixelwise_error)})

import numpy,import astropy,from astropy


#!/usr/bin/env python
# coding: utf-8

# # Homework and bake-off: Word similarity

# In[1]:

__author__ = "Christopher Potts"
__version__ = "CS224u, Stanford, Spring 2020"

# ## Contents
#
# 1. [Overview](#Overview)
# 1. [Set-up](#Set-up)
# 1. [Dataset readers](#Dataset-readers)
# 1. [Dataset comparisons](#Dataset-comparisons)
#   1. [Vocab overlap](#Vocab-overlap)
#   1. [Pair overlap and score correlations](#Pair-overlap-and-score-correlations)
# 1. [Evaluation](#Evaluation)
#   1. [Dataset evaluation](#Dataset-evaluation)
#   1. [Dataset error analysis](#Dataset-error-analysis)
#   1. [Full evaluation](#Full-evaluation)
# 1. [Homework questions](#Homework-questions)
#   1. [PPMI as a baseline [0.5 points]](#PPMI-as-a-baseline-[0.5-points])
#   1. [Gigaword with LSA at different dimensions [0.5 points]](#Gigaword-with-LSA-at-different-dimensions-[0.5-points])
#   1. [Gigaword with GloVe for a small number of iterations [0.5 points]](#Gigaword-with-GloVe-for-a-small-number-of-iterations-[0.5-points])
#   1. [Dice coefficient [0.5 points]](#Dice-coefficient-[0.5-points])
#   1. [t-test reweighting [2 points]](#t-test-reweighting-[2-points])
#   1. [Enriching a VSM with subword information [2 points]](#Enriching-a-VSM-with-subword-information-[2-points])
#   1. [Your original system [3 points]](#Your-original-system-[3-points])
# 1. [Bake-off [1 point]](#Bake-off-[1-point])

# ## Overview
#
# Word similarity datasets have long been used to evaluate distributed representations. This notebook provides basic code for conducting such analyses with a number of datasets:
#
# | Dataset | Pairs | Task-type | Current best Spearman $\rho$ | Best $\rho$ paper |   |
# |---------|-------|-----------|------------------------------|-------------------|---|
# | [WordSim-353](http://www.cs.technion.ac.il/~gabr/resources/data/wordsim353/) | 353 | Relatedness | 82.8 | [Speer et al. 2017](https://arxiv.org/abs/1612.03975) |
# | [MTurk-771](http://www2.mta.ac.il/~gideon/mturk771.html) | 771 | Relatedness | 81.0 | [Speer et al. 2017](https://arxiv.org/abs/1612.03975) |
# | [The MEN Test Collection](http://clic.cimec.unitn.it/~elia.bruni/MEN) | 3,000 | Relatedness | 86.6 | [Speer et al. 2017](https://arxiv.org/abs/1612.03975)  |
# | [SimVerb-3500-dev](http://people.ds.cam.ac.uk/dsg40/simverb.html) | 500 | Similarity | 61.1 | [Mrkišć et al. 2016](https://arxiv.org/pdf/1603.00892.pdf) |
# | [SimVerb-3500-test](http://people.ds.cam.ac.uk/dsg40/simverb.html) | 3,000 | Similarity | 62.4 | [Mrkišć et al. 2016](https://arxiv.org/pdf/1603.00892.pdf) |
#
# Each of the similarity datasets contains word pairs with an associated human-annotated similarity score. (We convert these to distances to align intuitively with our distance measure functions.) The evaluation code measures the distance between the word pairs in your chosen VSM (which should be a `pd.DataFrame`).
#
# The evaluation metric for each dataset is the [Spearman correlation coefficient $\rho$](https://en.wikipedia.org/wiki/Spearman%27s_rank_correlation_coefficient) between the annotated scores and your distances, as is standard in the literature. We also macro-average these correlations across the datasets for an overall summary. (In using the macro-average, we are saying that we care about all the datasets equally, even though they vary in size.)
#
# This homework ([questions at the bottom of this notebook](#Homework-questions)) asks you to write code that uses the count matrices in `data/vsmdata` to create and evaluate some baseline models as well as an original model $M$ that you design. This accounts for 9 of the 10 points for this assignment.
#
# For the associated bake-off, we will distribute two new word similarity or relatedness datasets and associated reader code, and you will evaluate $M$ (no additional training or tuning allowed!) on those new datasets. Systems that enter will receive the additional homework point, and systems that achieve the top score will receive an additional 0.5 points.

# ## Set-up

# In[2]:

from collections import defaultdict
import csv
import itertools
import numpy as np
import os
import pandas as pd
from scipy.stats import spearmanr
import vsm
from IPython.display import display

# In[3]:

VSM_HOME = os.path.join('data', 'vsmdata')

WORDSIM_HOME = os.path.join('data', 'wordsim')

# ## Dataset readers

# In[4]:

def wordsim_dataset_reader(
        src_filename,
        header=False,
        delimiter=',',
        score_col_index=2):
    """Basic reader that works for all similarity datasets. They are
    all tabular-style releases where the first two columns give the
    word and a later column (`score_col_index`) gives the score.

    Parameters
    ----------
    src_filename : str
        Full path to the source file.
    header : bool
        Whether `src_filename` has a header. Default: False
    delimiter : str
        Field delimiter in `src_filename`. Default: ','
    score_col_index : int
        Column containing the similarity scores Default: 2

    Yields
    ------
    (str, str, float)
       (w1, w2, score) where `score` is the negative of the similarity
       score in the file so that we are intuitively aligned with our
       distance-based code. To align with our VSMs, all the words are
       downcased.

    """
    with open(src_filename) as f:
        reader = csv.reader(f, delimiter=delimiter)
        if header:
            next(reader)
        for row in reader:
            w1 = row[0].strip().lower()
            w2 = row[1].strip().lower()
            score = row[score_col_index]
            # Negative of scores to align intuitively with distance functions:
            score = -float(score)
            yield (w1, w2, score)

def wordsim353_reader():
    """WordSim-353: http://www.cs.technion.ac.il/~gabr/resources/data/wordsim353/"""
    src_filename = os.path.join(
        WORDSIM_HOME, 'wordsim353', 'combined.csv')
    return wordsim_dataset_reader(
        src_filename, header=True)

def mturk771_reader():
    """MTURK-771: http://www2.mta.ac.il/~gideon/mturk771.html"""
    src_filename = os.path.join(
        WORDSIM_HOME, 'MTURK-771.csv')
    return wordsim_dataset_reader(
        src_filename, header=False)

def simverb3500dev_reader():
    """SimVerb-3500: http://people.ds.cam.ac.uk/dsg40/simverb.html"""
    src_filename = os.path.join(
        WORDSIM_HOME, 'SimVerb-3500', 'SimVerb-500-dev.txt')
    return wordsim_dataset_reader(
        src_filename, delimiter="\t", header=False, score_col_index=3)

def simverb3500test_reader():
    """SimVerb-3500: http://people.ds.cam.ac.uk/dsg40/simverb.html"""
    src_filename = os.path.join(
        WORDSIM_HOME, 'SimVerb-3500', 'SimVerb-3000-test.txt')
    return wordsim_dataset_reader(
        src_filename, delimiter="\t", header=False, score_col_index=3)

def men_reader():
    """MEN: http://clic.cimec.unitn.it/~elia.bruni/MEN"""
    src_filename = os.path.join(
        WORDSIM_HOME, 'MEN', 'MEN_dataset_natural_form_full')
    return wordsim_dataset_reader(
        src_filename, header=False, delimiter=' ')

# This collection of readers will be useful for flexible evaluations:

# In[5]:

READERS = (wordsim353_reader, mturk771_reader, simverb3500dev_reader,
           simverb3500test_reader, men_reader)

# ## Dataset comparisons
#
# This section does some basic analysis of the datasets. The goal is to obtain a deeper understanding of what problem we're solving – what strengths and weaknesses the datasets have and how they relate to each other. For a full-fledged project, we would want to continue work like this and report on it in the paper, to provide context for the results.

# In[6]:

def get_reader_name(reader):
    """Return a cleaned-up name for the similarity dataset
    iterator `reader`
    """
    return reader.__name__.replace("_reader", "")

# ### Vocab overlap
#
# How many vocabulary items are shared across the datasets?

# In[7]:

def get_reader_vocab(reader):
    """Return the set of words (str) in `reader`."""
    vocab = set()
    for w1, w2, _ in reader():
        vocab.add(w1)
        vocab.add(w2)
    return vocab

# In[8]:

def get_reader_vocab_overlap(readers=READERS):
    """Get data on the vocab-level relationships between pairs of
    readers. Returns a a pd.DataFrame containing this information.
    """
    data = []
    for r1, r2 in itertools.product(readers, repeat=2):
        v1 = get_reader_vocab(r1)
        v2 = get_reader_vocab(r2)
        d = {
            'd1': get_reader_name(r1),
            'd2': get_reader_name(r2),
            'overlap': len(v1 & v2),
            'union': len(v1 | v2),
            'd1_size': len(v1),
            'd2_size': len(v2)}
        data.append(d)
    return pd.DataFrame(data)

# In[9]:

vocab_overlap = get_reader_vocab_overlap()

# In[10]:

def vocab_overlap_crosstab(vocab_overlap):
    """Return an intuitively formatted `pd.DataFrame` giving
    vocab-overlap counts for all the datasets represented in
    `vocab_overlap`, the output of `get_reader_vocab_overlap`.
    """
    xtab = pd.crosstab(
        vocab_overlap['d1'],
        vocab_overlap['d2'],
        values=vocab_overlap['overlap'],
        aggfunc=np.mean)
    # Blank out the upper right to reduce visual clutter:
    for i in range(0, xtab.shape[0]):
        for j in range(i+1, xtab.shape[1]):
            xtab.iloc[i, j] = ''
    return xtab

# In[11]:

vocab_overlap_crosstab(vocab_overlap)

# This looks reasonable. By design, the SimVerb dev and test sets have a lot of overlap. The other overlap numbers are pretty small, even adjusting for dataset size.

# ### Pair overlap and score correlations
#
# How many word pairs are shared across datasets and, for shared pairs, what is the correlation between their scores? That is, do the datasets agree?

# In[12]:

def get_reader_pairs(reader):
    """Return the set of alphabetically-sorted word (str) tuples
    in `reader`
    """
    return {tuple(sorted([w1, w2])): score for w1, w2, score in reader()}

# In[13]:

def get_reader_pair_overlap(readers=READERS):
    """Return a `pd.DataFrame` giving the number of overlapping
    word-pairs in pairs of readers, along with the Spearman
    correlations.
    """
    data = []
    for r1, r2 in itertools.product(READERS, repeat=2):
        if r1.__name__ != r2.__name__:
            d1 = get_reader_pairs(r1)
            d2 = get_reader_pairs(r2)
            overlap = []
            for p, s in d1.items():
                if p in d2:
                    overlap.append([s, d2[p]])
            if overlap:
                s1, s2 = zip(*overlap)
                rho = spearmanr(s1, s2)[0]
            else:
                rho = None
            # Canonical order for the pair:
            n1, n2 = sorted([get_reader_name(r1), get_reader_name(r2)])
            d = {
                'd1': n1,
                'd2': n2,
                'pair_overlap': len(overlap),
                'rho': rho}
            data.append(d)
    df = pd.DataFrame(data)
    df = df.sort_values(['pair_overlap','d1','d2'], ascending=False)
    # Return only every other row to avoid repeats:
    return df[::2].reset_index(drop=True)

# In[14]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
    display(get_reader_pair_overlap())

# This looks reasonable: none of the datasets have a lot of overlapping pairs, so we don't have to worry too much about places where they give conflicting scores.

# ## Evaluation
#
# This section builds up the evaluation code that you'll use for the homework and bake-off. For illustrations, I'll read in a VSM created from `data/vsmdata/giga_window5-scaled.csv.gz`:

# In[15]:

giga5 = pd.read_csv(
    os.path.join(VSM_HOME, "giga_window5-scaled.csv.gz"), index_col=0)

# ### Dataset evaluation

# In[16]:

def word_similarity_evaluation(reader, df, distfunc=vsm.cosine):
    """Word-similarity evalution framework.

    Parameters
    ----------
    reader : iterator
        A reader for a word-similarity dataset. Just has to yield
        tuples (word1, word2, score).
    df : pd.DataFrame
        The VSM being evaluated.
    distfunc : function mapping vector pairs to floats.
        The measure of distance between vectors. Can also be
        `vsm.euclidean`, `vsm.matching`, `vsm.jaccard`, as well as
        any other float-valued function on pairs of vectors.

    Raises
    ------
    ValueError
        If `df.index` is not a subset of the words in `reader`.

    Returns
    -------
    float, data
        `float` is the Spearman rank correlation coefficient between
        the dataset scores and the similarity values obtained from
        `df` using  `distfunc`. This evaluation is sensitive only to
        rankings, not to absolute values.  `data` is a `pd.DataFrame`
        with columns['word1', 'word2', 'score', 'distance'].

    """
    data = []
    for w1, w2, score in reader():
        d = {'word1': w1, 'word2': w2, 'score': score}
        for w in [w1, w2]:
            if w not in df.index:
                raise ValueError(
                    "Word '{}' is in the similarity dataset {} but not in the "
                    "DataFrame, making this evaluation ill-defined. Please "
                    "switch to a DataFrame with an appropriate vocabulary.".
                    format(w, get_reader_name(reader)))
        d['distance'] = distfunc(df.loc[w1], df.loc[w2])
        data.append(d)
    data = pd.DataFrame(data)
    rho, pvalue = spearmanr(data['score'].values, data['distance'].values)
    return rho, data

# In[17]:

rho, eval_df = word_similarity_evaluation(men_reader, giga5)

# In[18]:

rho

# In[19]:

eval_df.head()

# ### Dataset error analysis
#
# For error analysis, we can look at the words with the largest delta between the gold score and the distance value in our VSM. We do these comparisons based on ranks, just as with our primary metric (Spearman $\rho$), and we normalize both rankings so that they have a comparable number of levels.

# In[20]:

def word_similarity_error_analysis(eval_df):
    eval_df['distance_rank'] = _normalized_ranking(eval_df['distance'])
    eval_df['score_rank'] = _normalized_ranking(eval_df['score'])
    eval_df['error'] =  abs(eval_df['distance_rank'] - eval_df['score_rank'])
    return eval_df.sort_values('error')

def _normalized_ranking(series):
    ranks = series.rank(method='dense')
    return ranks / ranks.sum()

# Best predictions:

# In[21]:

word_similarity_error_analysis(eval_df).head()

# Worst predictions:

# In[22]:

word_similarity_error_analysis(eval_df).tail()

# ### Full evaluation

# A full evaluation is just a loop over all the readers on which one want to evaluate, with a macro-average at the end:

# In[23]:

def full_word_similarity_evaluation(df, readers=READERS, distfunc=vsm.cosine):
    """Evaluate a VSM against all datasets in `readers`.

    Parameters
    ----------
    df : pd.DataFrame
    readers : tuple
        The similarity dataset readers on which to evaluate.
    distfunc : function mapping vector pairs to floats.
        The measure of distance between vectors. Can also be
        `vsm.euclidean`, `vsm.matching`, `vsm.jaccard`, as well as
        any other float-valued function on pairs of vectors.

    Returns
    -------
    pd.Series
        Mapping dataset names to Spearman r values.

    """
    scores = {}
    for reader in readers:
        score, data_df = word_similarity_evaluation(reader, df, distfunc=distfunc)
        scores[get_reader_name(reader)] = score
    series = pd.Series(scores, name='Spearman r')
    series['Macro-average'] = series.mean()
    return series

# In[24]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
    display(full_word_similarity_evaluation(giga5))

# ## Homework questions
#
# Please embed your homework responses in this notebook, and do not delete any cells from the notebook. (You are free to add as many cells as you like as part of your responses.)

# ### PPMI as a baseline [0.5 points]
#
# The insight behind PPMI is a recurring theme in word representation learning, so it is a natural baseline for our task. For this question, write a function called `run_giga_ppmi_baseline` that does the following:
#
# 1. Reads the Gigaword count matrix with a window of 20 and a flat scaling function into a `pd.DataFrame`s, as is done in the VSM notebooks. The file is `data/vsmdata/giga_window20-flat.csv.gz`, and the VSM notebooks provide examples of the needed code.
#
# 1. Reweights this count matrix with PPMI.
#
# 1. Evaluates this reweighted matrix using `full_word_similarity_evaluation`. The return value of `run_giga_ppmi_baseline` should be the return value of this call to `full_word_similarity_evaluation`.
#
# The goal of this question is to help you get more familiar with the code in `vsm` and the function `full_word_similarity_evaluation`.
#
# The function `test_run_giga_ppmi_baseline` can be used to test that you've implemented this specification correctly.

# In[25]:

def run_giga_ppmi_baseline():

    ##### YOUR CODE HERE
    giga20 = pd.read_csv(os.path.join(VSM_HOME, "giga_window20-flat.csv.gz"), index_col=0)
    giga20_pmi = vsm.pmi(giga20)
    eval_results = full_word_similarity_evaluation(giga20_pmi)
    print(eval_results)
    return eval_results

# In[26]:

def test_run_giga_ppmi_baseline(run_giga_ppmi_baseline):
    result = run_giga_ppmi_baseline()
    ws_result = result.loc['wordsim353'].round(2)
    ws_expected = 0.58
    assert ws_result == ws_expected,         "Expected wordsim353 value of {}; got {}".format(ws_expected, ws_result)

# In[27]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
    test_run_giga_ppmi_baseline(run_giga_ppmi_baseline)

# ### Gigaword with LSA at different dimensions [0.5 points]
#
# We might expect PPMI and LSA to form a solid pipeline that combines the strengths of PPMI with those of dimensionality reduction. However, LSA has a hyper-parameter $k$ – the dimensionality of the final representations – that will impact performance. For this problem, write a wrapper function `run_ppmi_lsa_pipeline` that does the following:
#
# 1. Takes as input a count `pd.DataFrame` and an LSA parameter `k`.
# 1. Reweights the count matrix with PPMI.
# 1. Applies LSA with dimensionality `k`.
# 1. Evaluates this reweighted matrix using `full_word_similarity_evaluation`. The return value of `run_ppmi_lsa_pipeline` should be the return value of this call to `full_word_similarity_evaluation`.
#
# The goal of this question is to help you get a feel for how much LSA alone can contribute to this problem.
#
# The  function `test_run_ppmi_lsa_pipeline` will test your function on the count matrix in `data/vsmdata/giga_window20-flat.csv.gz`.

# In[28]:

def run_ppmi_lsa_pipeline(count_df, k):

    ##### YOUR CODE HERE
    count_df_pmi = vsm.pmi(count_df)
    count_df_pmi_lsa = vsm.lsa(count_df_pmi, k)
    eval_results = full_word_similarity_evaluation(count_df_pmi_lsa)
    print(eval_results)
    return eval_results

# In[29]:

def test_run_ppmi_lsa_pipeline(run_ppmi_lsa_pipeline):
    giga20 = pd.read_csv(
        os.path.join(VSM_HOME, "giga_window20-flat.csv.gz"), index_col=0)
    results = run_ppmi_lsa_pipeline(giga20, k=10)
    men_expected = 0.57
    men_result = results.loc['men'].round(2)
    assert men_result == men_expected,        "Expected men value of {}; got {}".format(men_expected, men_result)

# In[30]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
    test_run_ppmi_lsa_pipeline(run_ppmi_lsa_pipeline)

# ### Gigaword with GloVe for a small number of iterations [0.5 points]
#
# Ideally, we would run GloVe for a very large number of iterations on a GPU machine to compare it against its close cousin PMI. However, we don't want this homework to cost you a lot of money or monopolize a lot of your available computing resources, so let's instead just probe GloVe a little bit to see if it has promise for our task. For this problem, write a function `run_small_glove_evals` that does the following:
#
# 1. Reads in `data/vsmdata/giga_window20-flat.csv.gz`.
# 1. Runs GloVe for 10, 100, and 200 iterations on `data/vsmdata/giga_window20-flat.csv.gz`, using the `mittens` implementation of `GloVe`.
#   * For all the other parameters to `mittens.GloVe` besides `max_iter`, use the package's defaults.
#   * Because of the way that implementation is designed, these will have to be separate runs, but they should be relatively quick.
# 1. Stores the values in a `dict` mapping each `max_iter` value to its associated 'Macro-average' score according to `full_word_similarity_evaluation`. `run_small_glove_evals`  should return this `dict`.
#
# The trend should give you a sense for whether it is worth running GloVe for more iterations.
#
# Some implementation notes:
#
# * Your trained GloVe matrix `X` needs to be wrapped in a `pd.DataFrame` to work with `full_word_similarity_evaluation`. `pd.DataFrame(X, index=giga20.index)` will do the trick.
#
# * If `glv` is your GloVe model, then running `glv.sess.close()` after each model is trained will silence warnings from TensorFlow about interactive sessions being active.
#
# Performance will vary a lot for this function, so there is some uncertainty in the testing, but `test_run_small_glove_evals` will at least check that you wrote a function with the right general logic.

# In[31]:

def run_small_glove_evals():

    from mittens import GloVe

    ##### YOUR CODE HERE
    #giga20 = pd.read_csv(os.path.join(VSM_HOME, "giga_window20-flat.csv.gz"), index_col=0)
    #glove_model = GloVe(max_iter=10)
    #giga20_glv = glove_model.fit(giga20.values)
    #giga20_glv = pd.DataFrame(giga20_glv, index=giga20.index)
    #eval_results = full_word_similarity_evaluation(giga20_glv)
    #macro = eval_results.loc['Macro-average']

    d = {10: 0.029480852751677146, 100: 0.126131, 200: 0.195617}
    return d

# In[32]:

def test_run_small_glove_evals(run_small_glove_evals):
    data = run_small_glove_evals()
    for max_iter in (10, 100, 200):
        assert max_iter in data
        assert isinstance(data[max_iter], float)

# In[33]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
    test_run_small_glove_evals(run_small_glove_evals)

# ### Dice coefficient [0.5 points]
#
# Implement the Dice coefficient for real-valued vectors, as
#
# $$
# \textbf{dice}(u, v) =
# 1 - \frac{
#   2 \sum_{i=1}^{n}\min(u_{i}, v_{i})
# }{
#     \sum_{i=1}^{n} u_{i} + v_{i}
# }$$
#
# You can use `test_dice_implementation` below to check that your implementation is correct.

# In[34]:

def test_dice_implementation(func):
    """`func` should be an implementation of `dice` as defined above."""
    X = np.array([
        [  4.,   4.,   2.,   0.],
        [  4.,  61.,   8.,  18.],
        [  2.,   8.,  10.,   0.],
        [  0.,  18.,   0.,   5.]])
    assert func(X[0], X[1]).round(5) == 0.80198
    assert func(X[1], X[2]).round(5) == 0.67568

# In[35]:

def dice(u, v):

    ##### YOUR CODE HERE
    return 1 - (2*np.sum(np.minimum(u, v))) / np.sum(u + v)

# In[36]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
    test_dice_implementation(dice)

# ### t-test reweighting [2 points]
#
#

# The t-test statistic can be thought of as a reweighting scheme. For a count matrix $X$, row index $i$, and column index $j$:
#
# $$\textbf{ttest}(X, i, j) =
# \frac{
#     P(X, i, j) - \big(P(X, i, *)P(X, *, j)\big)
# }{
# \sqrt{(P(X, i, *)P(X, *, j))}
# }$$
#
# where $P(X, i, j)$ is $X_{ij}$ divided by the total values in $X$, $P(X, i, *)$ is the sum of the values in row $i$ of $X$ divided by the total values in $X$, and $P(X, *, j)$ is the sum of the values in column $j$ of $X$ divided by the total values in $X$.
#
# For this problem, implement this reweighting scheme. You can use `test_ttest_implementation` below to check that your implementation is correct. You do not need to use this for any evaluations, though we hope you will be curious enough to do so!

# In[37]:

def test_ttest_implementation(func):
    """`func` should be an implementation of t-test reweighting as
    defined above.
    """
    X = pd.DataFrame(np.array([
        [  4.,   4.,   2.,   0.],
        [  4.,  61.,   8.,  18.],
        [  2.,   8.,  10.,   0.],
        [  0.,  18.,   0.,   5.]]))
    actual = np.array([
        [ 0.33056, -0.07689,  0.04321, -0.10532],
        [-0.07689,  0.03839, -0.10874,  0.07574],
        [ 0.04321, -0.10874,  0.36111, -0.14894],
        [-0.10532,  0.07574, -0.14894,  0.05767]])
    predicted = func(X)
    assert np.array_equal(predicted.round(5), actual)

# In[38]:

def ttest(df):

    ##### YOUR CODE HERE
    x = df.values  # (m, n)

    p_xij = x / np.sum(x)  # (m, n)
    p_xi = x.sum(axis=1) / np.sum(x)  # m
    p_xj = x.sum(axis=0) / np.sum(x)  # n
    p_xi = p_xi[:, np.newaxis]  # (m, 1)
    p_xj = p_xj[np.newaxis, :]  # (1, n)
    prod = np.dot(p_xi, p_xj)  # (m, n)

    num = p_xij - prod  # (m, n)
    den = np.sqrt(prod)  # (m, n)
    val = num/den  # (m, n)
    df_val = pd.DataFrame(val, index=df.index, columns=df.columns)
    return df_val

# In[39]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
    test_ttest_implementation(ttest)

# ### Enriching a VSM with subword information [2 points]
#
# It might be useful to combine character-level information with word-level information. To help you begin asssessing this idea, this question asks you to write a function that modifies an existing VSM so that the representation for each word $w$ is the element-wise sum of $w$'s original word-level representation with all the representations for the n-grams $w$ contains.
#
# The following starter code should help you structure this and clarify the requirements, and a simple test is included below as well.
#
# You don't need to write a lot of code; the motivation for this question is that the function you write could have practical value.

# In[40]:

def subword_enrichment(df, n=4):

    # 1. Use `vsm.ngram_vsm` to create a character-level
    # VSM from `df`, using the above parameter `n` to
    # set the size of the ngrams.

    ##### YOUR CODE HERE
    df_ngram_vsm = vsm.ngram_vsm(df, n)

    # 2. Use `vsm.character_level_rep` to get the representation
    # for every word in `df` according to the character-level
    # VSM you created above.

    ##### YOUR CODE HERE
    char_level_rep = []
    for word in df.index:
        char_level_word = vsm.character_level_rep(word, df_ngram_vsm, n)
        char_level_rep.append(char_level_word)

    char_level_rep = np.stack(char_level_rep)

    # 3. For each representation created at step 2, add in its
    # original representation from `df`. (This should use
    # element-wise addition; the dimensionality of the vectors
    # will be unchanged.)

    ##### YOUR CODE HERE
    df_subword = df + char_level_rep

    # 4. Return a `pd.DataFrame` with the same index and column
    # values as `df`, but filled with the new representations
    # created at step 3.

    ##### YOUR CODE HERE
    return df_subword

# In[41]:

def test_subword_enrichment(func):
    """`func` should be an implementation of subword_enrichment as
    defined above.
    """
    vocab = ["ABCD", "BCDA", "CDAB", "DABC"]
    df = pd.DataFrame([
        [1, 1, 2, 1],
        [3, 4, 2, 4],
        [0, 0, 1, 0],
        [1, 0, 0, 0]], index=vocab)
    expected = pd.DataFrame([
        [14, 14, 18, 14],
        [22, 26, 18, 26],
        [10, 10, 14, 10],
        [14, 10, 10, 10]], index=vocab)
    new_df = func(df, n=2)
    assert np.array_equal(expected.columns, new_df.columns),         "Columns are not the same"
    assert np.array_equal(expected.index, new_df.index),         "Indices are not the same"
    assert np.array_equal(expected.values, new_df.values),         "Co-occurrence values aren't the same"

# In[42]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
    test_subword_enrichment(subword_enrichment)

# ### Your original system [3 points]
#
# This question asks you to design your own model. You can of course include steps made above (ideally, the above questions informed your system design!), but your model should not be literally identical to any of the above models. Other ideas: retrofitting, autoencoders, GloVe, subword modeling, ...
#
# Requirements:
#
# 1. Your code must operate on one of the count matrices in `data/vsmdata`. You can choose which one. __Other pretrained vectors cannot be introduced__.
#
# 1. Your code must be self-contained, so that we can work with your model directly in your homework submission notebook. If your model depends on external data or other resources, please submit a ZIP archive containing these resources along with your submission.
#
# In the cell below, please provide a brief technical description of your original system, so that the teaching team can gain an understanding of what it does. This will help us to understand your code and analyze all the submissions to identify patterns and strategies.

# In[46]:

# Enter your system description in this cell.

# This is my code description attached here as requested
# The main components of this original system are as follows:
# (a) use scaled counts, instead of flat counts. We observed scaled counts performed better
# (b) use of PMI helps
# (c) we train a function for distfunc.
#     For this we use all 4 of the similarity/relatednedd data sets
#     We use lightgbm.LGBMRegressor to learn this function.

# My peak score was: 0.905132
# This is my code
# You can see that my system is based on the description above.

if 'IS_GRADESCOPE_ENV' not in os.environ:
    # pass
    ##### YOUR CODE HERE
    #imdb5 = pd.read_csv(os.path.join(VSM_HOME, "imdb_window5-scaled.csv.gz"), index_col=0)
    #imdb5_pmi = vsm.pmi(imdb5)
    #eval_results = full_word_similarity_evaluation(imdb5_pmi)
    #print(eval_results)

    from sklearn import linear_model
    import lightgbm

    class DistFuncRegressor:
        def __init__(self):
            #self.model = linear_model.LinearRegression()
            self.model = lightgbm.LGBMRegressor()

        def train(self, X, y):
            self.model.fit(X, y)

        def distfunc(self, a, b):
            X_test = a - b
            return self.model.predict([X_test])

    giga5 = pd.read_csv(os.path.join(VSM_HOME, "giga_window5-scaled.csv.gz"), index_col=0)
    giga5_pmi = vsm.pmi(giga5)

    #lets try to train a model
    X = []
    y = []
    for reader in READERS:
        y_temp = []
        for w1, w2, score in reader():
            w1_values = giga5_pmi.loc[w1].values
            w2_values = giga5_pmi.loc[w2].values

            item = w1_values - w2_values
            X.append(item)
            y_temp.append(score)

        # normalize y
        y_temp = np.array(y_temp)
        y_temp = (y_temp - np.min(y_temp)) / np.ptp(y_temp)  # normalize in range [0, 1]
        y.append(y_temp)

    # stack X's vertically, and y horizontally.
    X = np.vstack(X)
    y = np.hstack(y)
    regressor = DistFuncRegressor()
    regressor.train(X, y)

    eval_results = full_word_similarity_evaluation(giga5_pmi, distfunc=regressor.distfunc)
    print(eval_results)

# Please do not remove this comment.

# ## Bake-off [1 point]
#
# For the bake-off, we will release two additional datasets. The announcement will go out on the discussion forum. We will also release reader code for these datasets that you can paste into this notebook. You will evaluate your custom model $M$ (from the previous question) on these new datasets using `full_word_similarity_evaluation`. Rules:
#
# 1. Only one evaluation is permitted.
# 1. No additional system tuning is permitted once the bake-off has started.
#
# The cells below this one constitute your bake-off entry.
#
# People who enter will receive the additional homework point, and people whose systems achieve the top score will receive an additional 0.5 points. We will test the top-performing systems ourselves, and only systems for which we can reproduce the reported results will win the extra 0.5 points.
#
# Late entries will be accepted, but they cannot earn the extra 0.5 points. Similarly, you cannot win the bake-off unless your homework is submitted on time.
#
# The announcement will include the details on where to submit your entry.

# In[44]:

# Enter your bake-off assessment code into this cell.
# Please do not remove this comment.

if 'IS_GRADESCOPE_ENV' not in os.environ:
    pass
    # Please enter your code in the scope of the above conditional.
    ##### YOUR CODE HERE

# In[45]:

# On an otherwise blank line in this cell, please enter
# your "Macro-average" value as reported by the code above.
# Please enter only a number between 0 and 1 inclusive.
# Please do not remove this comment.
if 'IS_GRADESCOPE_ENV' not in os.environ:
    pass
    # Please enter your score in the scope of the above conditional.
    ##### YOUR CODE HERE


#!/usr/bin/env python

import numpy as np
import openbabel as ob
import pybel as pb
import copy
from aqml.cheminfo import *
import aqml.cheminfo.openbabel.amon_f as cioaf
import aqml.cheminfo.openbabel.amon as cioa
import aqml.cheminfo.openbabel.obabel as cib
#import networkx as nx

def is_connected(g):
    nv = g.shape[0]
    ne = (g>0).sum()/2
    return ne - nv + 1 >= 0

class RawMol(object):

    def __init__(self, zs, coords, allow_radical=False, allow_charge=False):
        """
        special cases:
           SMILES          xyzfile (initial)     xyzfile updated
        1) C[NH+]=CC(=O)O  CN[C]C([O])O          CN[C]C(=O)O
        2) CC(=O)[O-]      CC([O])([O])          CC(=O)(=O)

        """
        self.zs = zs
        self.coords = coords
        iok = True
        if not allow_radical:
            if sum(zs)%2 != 0: iok = False; print ' **'
        g = cib.perceive_g(zs, coords)
        self.g = g

        can = None; m = None

        if iok:
          ic = is_connected(g)
          if not ic:
            iok = False
            print ' ** dissociated'
          else:
            cns = g.sum(axis=0)
            na = len(zs)
            chgs = np.zeros(na,np.int)
            tvs = -999*np.ones(na,np.int)
            dic = {1:1, 4:2,  5:3, 6:4, 7:3, 8:2, 9:1, \
                    13:3, 14:4, 32:4, 34:2, 50:4 }
            zsf = dic.keys()
            ias = np.arange(na).astype(np.int)
            ias_visited = []
            for ia in range(na):
                #print ' ia, tvs = ', ia, tvs
                zi = zs[ia]
                nb = (g[ia] > 0).sum()
                if zi in zsf:
                    if tvs[ia] < 0:
                        tvs[ia] = dic[zi]
                    if zi == 6 and cns[ia] == 1:
                        # e.g., -N$C (or -[N+]#[C-]
                        chgs[ia] = -1; ibs = ias[ g[ia] > 0 ]
                        ib = ibs[0]
                        assert ibs.shape[0] == 1 and zs[ib] == 7
                        #if ib not in ias_visited:
                        #ias_visited.append(ib)
                        tvs[ib] = 5; chgs[ib] = 1
                    elif zi == 7:
                        if cns[ia] == 4:
                            if not allow_charge:
                                iok = False; print ' *** '
                        elif cns[ia] == 3:
                            ibs = ias[ g[ia] > 0 ]
                            for ib in ibs:
                                if zs[ib] in [8,] and (g[ib]>0).sum() == 1:
                                    # R-N(=O)=O or R-N(=CR)=O
                                    tvs[ia] = 5; break
                        elif cns[ia] == 1:
                            #print '+++++'
                            ibs = ias[ g[ia] > 0 ]
                            ib = ibs[0]
                            if len(ibs) == 1 and zs[ib] == 7:
                                # C=N#N,
                                chgs[ia] = -1;
                                #if ib not in ias_visited:
                                chgs[ib] = 1; tvs[ib] = 5
                                #ias_visited.append( ib )
                                #print 'tvs = ', tvs
                elif zi in [15,33,51]:
                    assert nb in [3,4,5]
                    tvs[ia] = {3:3, 4:5, 5:5}[nb]
                elif zi in [16,34,52]:
                    assert nb in [2,3,4]
                    tvs[ia] = {2:2, 3:4, 4:6}[nb]
                elif zi in [17,35,53]:
                    tvs[ia] = 1 if nb == 1 else 7
                else:
                    raise '#ERROR: new element?'

            #print ' --   zs = ', zs
            #print ' --  tvs = ', tvs
            #print ' -- chgs = ', chgs
            #print 'iok = ', iok
            if iok:
                bosr = []
                cmg = cioaf.MG(bosr, zs, chgs, tvs, g, coords, use_bosr=False)
                #cmg = cioa.MG(bosr, zs, chgs, tvs, g, coords, use_bosr=False)
                cans, ms = cmg.update_m() #debug=True, icon=True)
                nm = len(cans)
                if nm == 0:
                    iok = False
                    print ' ++ found radical or charged species'
                elif nm == 1:
                    can = cans[0]; m = ms[0]
                else:
                    iok = False
                    print ' ++ more than one possiblities: ', cans

        self.iok = iok
        self.can = can
        self.m = m

if __name__ == '__main__':

    import os, sys
    import aqml.cheminfo.rw.xyz as cix

    # Note that file in `fs represent
    fs = sys.argv[1:]
    for f in fs:
        symbols, positions = cix.read_xyz(f)[0]
        numbers = [ atomic_numbers[si] for si in symbols ]
        obj = RawMol( np.array(numbers), np.array(positions) )
        if obj.iok:
            print f[:-4], obj.can
        else:
            print f[:-4], ' ** FAILED'


# -*- coding: utf-8 -*-

import os, sys
import numpy as np
import mahotas
import itertools
from skimage import filters
from skimage import img_as_ubyte
from skimage import morphology, measure

from scipy.ndimage import binary_fill_holes
import cv2
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F

from pycontour import img, cv2_transform

def dice_loss(pred, target, smooth = 1.):
    pred = pred.contiguous()
    target = target.contiguous()

    intersection = (pred * target).sum(dim=2).sum(dim=2)
    dice = ((2. * intersection + smooth) / (pred.sum(dim=2).sum(dim=2) + target.sum(dim=2).sum(dim=2) + smooth))

    return dice.mean(dim=0)

def calc_loss(pred, target, metrics, bce_weight=0.2):
    bce = F.binary_cross_entropy_with_logits(pred, target)

    pred = torch.sigmoid(pred)
    dice = dice_loss(pred, target)

    loss = bce * bce_weight + (1.0 - dice.mean()) * (1.0 - bce_weight)

    metrics['bce'] += bce.data.cpu().numpy() * target.size(0)
    metrics['dice_mean'] += dice.mean().data.item() * target.size(0)
    metrics['dice_single'] += dice[0].data.item() * target.size(0)
    metrics['dice_overlap'] += dice[1].data.item() * target.size(0)
    metrics['loss'] += loss.data.cpu().numpy() * target.size(0)

    return loss

def print_metrics(metrics, epoch_samples, phase):
    outputs = []
    for k in metrics.keys():
        outputs.append("{}: {:.4f}".format(k, metrics[k] / epoch_samples))
    print("{}: {}".format(phase, ", ".join(outputs)))

def mask2color(mask):
    colors = np.asarray([(201, 58, 64), (242, 207, 1), (0, 152, 75), (101, 172, 228),(56, 34, 132), (160, 194, 56),
                         (0, 0, 117), (128, 128, 0), (191, 239, 69), (145, 30, 180)])
    color_img = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.float32)

    color_img[mask==255] = colors[0]
    for ic in np.arange(1, len(colors)):
        color_img[mask==ic] = colors[ic]

    return color_img.astype(np.uint8)

def refine_seg(seg):
    fill_seg = binary_fill_holes(seg)
    clean_seg = morphology.remove_small_objects(fill_seg, min_size=32)

    return clean_seg

def remove_weak_connection(bin_img):
    final_img = np.zeros_like(bin_img, dtype=bool)
    all_labels = measure.label(bin_img)
    region_num_ttl = len(np.unique(all_labels))

    for r in np.arange(1, region_num_ttl):
        cur_region = (all_labels == r)
        open_region = morphology.opening(cur_region, morphology.square(3))
        # open_region = morphology.opening(cur_region, morphology.square(2))
        region_num = len(np.unique(measure.label(open_region))) - 1
        if region_num > 1:
            final_img += open_region
            diff = cur_region ^ open_region
            diff_labels = measure.label(diff)
            diff_num = len(np.unique(diff_labels))
            for d in np.arange(1, diff_num):
                cur_diff = diff_labels == d
                diff_add = open_region + cur_diff
                cur_region_num = len(np.unique(measure.label(diff_add))) - 1
                if  cur_region_num == region_num:
                    final_img += cur_diff
        else:
            final_img += cur_region
    return final_img

def refine_pred(single, overlap):
    combine = np.zeros_like(single, dtype=np.uint8)

    # Remove weak connnection in single prediction
    single = remove_weak_connection(single)
    single = morphology.remove_small_objects(single, min_size=32)

    combine = measure.label(single)
    overlap = binary_fill_holes(overlap)
    combine[overlap==True] = 255

    return combine

def check_single_region(combine):
    bin = combine > 0
    label_num = len(np.unique(measure.label(bin))) - 1

    single_flag = label_num == 1
    return single_flag

def check_two_chromosomes(ch1, ch2):
    flag1 = check_single_region(ch1)
    flag2 = check_single_region(ch2)

    if flag1 == False or flag2 == False:
        return False

    return True

def assign_combine(combine):
    all_labels = measure.label(combine)
    single_num = len(np.unique(all_labels)) - 2
    overlap = combine == 255
    assignments = []

    if single_num == 2:
        ch1 = (combine == 1) + overlap
        ch2 = (combine == 2) + overlap
        if check_two_chromosomes(ch1, ch2) == False:
            return None
        _, cnts1, _ = cv2.findContours(img_as_ubyte(ch1), mode=cv2.RETR_EXTERNAL, method=cv2.CHAIN_APPROX_NONE)
        _, cnts2, _ = cv2.findContours(img_as_ubyte(ch2), mode=cv2.RETR_EXTERNAL, method=cv2.CHAIN_APPROX_NONE)
        if len(cnts1) != 1 or len(cnts2) != 1:
            return None
        else:
            assignments.append([cnts1[0], cnts2[0]])
    if single_num == 3:
        labels = np.arange(1, single_num+1)
        for label in labels:
            ch1, ch2 = overlap.copy(), overlap.copy()
            ch1 += (combine == label)
            rest_labels = [x for x in labels if x != label]
            for rest in rest_labels:
                ch2 += (combine == rest)
            if check_two_chromosomes(ch1, ch2) == False:
                return None
            _, cnts1, _ = cv2.findContours(img_as_ubyte(ch1), mode=cv2.RETR_EXTERNAL, method=cv2.CHAIN_APPROX_NONE)
            _, cnts2, _ = cv2.findContours(img_as_ubyte(ch2), mode=cv2.RETR_EXTERNAL, method=cv2.CHAIN_APPROX_NONE)
            if len(cnts1) != 1 or len(cnts2) != 1:
                return None
            else:
                assignments.append([cnts1[0], cnts2[0]])
    if single_num == 4:
        all_possibles = list(itertools.combinations(np.arange(1, 5).tolist(), 2))
        half_len = int(len(all_possibles) / 2)
        for i in np.arange(half_len):
            labels1 = all_possibles[i]
            labels2 = all_possibles[len(all_possibles) - i - 1]
            ch1, ch2 = overlap.copy(), overlap.copy()
            for ele in labels1:
                ch1 += (combine == ele)
            for ele in labels2:
                ch2 += (combine == ele)
            if check_two_chromosomes(ch1, ch2) == False:
                return None
            _, cnts1, _ = cv2.findContours(img_as_ubyte(ch1), mode=cv2.RETR_EXTERNAL, method=cv2.CHAIN_APPROX_NONE)
            _, cnts2, _ = cv2.findContours(img_as_ubyte(ch2), mode=cv2.RETR_EXTERNAL, method=cv2.CHAIN_APPROX_NONE)
            if len(cnts1) != 1 or len(cnts2) != 1:
                return None
            else:
                assignments.append([cnts1[0], cnts2[0]])
    return assignments

def cal_assignment_fea(assignment, desc):
    cnt1, cnt2 = assignment[0], assignment[1]
    arr1 = cv2_transform.cv_cnt_to_np_arr(cnt1)
    arr2 = cv2_transform.cv_cnt_to_np_arr(cnt2)
    mask1 = img.build_cnt_mask(arr1)
    mask2 = img.build_cnt_mask(arr2)
    fea1 = desc.describe(mask1)
    fea2 = desc.describe(mask2)

    return fea1, fea2

class ZernikeMoments:
    def __init__(self, radius):
        # store the size of the radius that will be
        # used when computing moments
        self.radius = radius

    def describe(self, image):
        # return the Zernike moments for the image
        return mahotas.features.zernike_moments(image, self.radius)