#!/usr/bin/env python

# stdlib modules
import copy

# third party imports
import numpy as np
from openquake.hazardlib.geo.mesh import Mesh
from openquake.hazardlib.geo.point import Point
from openquake.hazardlib.geo.utils import get_orthographic_projection

from impactutils.vectorutils.ecef import latlon2ecef
from impactutils.vectorutils.ecef import ecef2latlon
from impactutils.vectorutils.vector import Vector
from impactutils.time.ancient_time import HistoricTime
from shakelib.utils.exception import ShakeLibException
from shakelib.rupture.base import Rupture
from shakelib.rupture import utils
from shakelib.rupture import gc2

class QuadRupture(Rupture):
    """
    Rupture class that represents the rupture surface as a combination of
    quadrilaterals. Each quadrilateral must have horizontal top and bottom
    edges and must be coplanar. These restrictions make the computation of
    rupture distances more efficient. The number of points in the top edges
    must match the number of points in the bottom edge.
    """

    def __init__(self, d, origin):
        """
        Create a QuadRupture instance from a GeoJSON dictionary and an Origin.

        Args:
           d (dict): Rupture GeoJSON dictionary.
           origin (Origin): Reference to a ShakeMap Origin object.

        Returns:
            QuadRupture instance.

        """

        polys = d['features'][0]['geometry']['coordinates'][0]
        n_polygons = len(polys)
        lon = []
        lat = []
        dep = []
        for i in range(n_polygons):
            p = polys[i]
            p_lons = [pt[0] for pt in p][0:-1]
            p_lats = [pt[1] for pt in p][0:-1]
            p_depths = [pt[2] for pt in p][0:-1]
            lon = lon + p_lons + [np.nan]
            lat = lat + p_lats + [np.nan]
            dep = dep + p_depths + [np.nan]

        # Add origin information to metadata
        odict = origin.__dict__
        for k, v in odict.items():
            if isinstance(v, HistoricTime):
                d['metadata'][k] = v.strftime('%Y-%m-%dT%H:%M:%SZ')
            else:
                d['metadata'][k] = v

        self._geojson = d
        self._lon = lon
        self._lat = lat
        self._depth = dep
        self._origin = origin
        self._reference = d['metadata']['reference']
        self._setQuadrilaterals()

    def getDepthAtPoint(self, lat, lon):
        SMALL_DISTANCE = 2e-03  # 2 meters
        depth = np.nan

        tmp = self.computeRjb(np.array([lon]), np.array([lat]), np.array([0]))
        if tmp > SMALL_DISTANCE:
            return depth

        i = 0
        imin = -1
        dmin = 9999999999999999
        for quad in self.getQuadrilaterals():
            pX = Vector.fromPoint(Point(lon, lat, 0))
            points = np.reshape(np.array([pX.x, pX.y, pX.z]), (1, 3))
            rjb = utils._quad_distance(quad, points, horizontal=True)
            if rjb[0][0] < dmin:
                dmin = rjb[0][0]
                imin = i
            i += 1

        quad = self._quadrilaterals[imin]
        P0, P1, P2, P3 = quad
        # project the quad and the point in question to orthographic defined by
        # quad
        xmin = np.min([P0.x, P1.x, P2.x, P3.x])
        xmax = np.max([P0.x, P1.x, P2.x, P3.x])
        ymin = np.min([P0.y, P1.y, P2.y, P3.y])
        ymax = np.max([P0.y, P1.y, P2.y, P3.y])
        proj = get_orthographic_projection(xmin, xmax, ymax, ymin)

        # project each vertex of quad (at 0 depth)
        s0x, s0y = proj(P0.x, P0.y)
        s1x, s1y = proj(P1.x, P1.y)
        s2x, s2y = proj(P2.x, P2.y)
        s3x, s3y = proj(P3.x, P3.y)
        sxx, sxy = proj(lon, lat)

        # turn these to vectors
        s0 = Vector(s0x, s0y, 0)
        s1 = Vector(s1x, s1y, 0)
        s3 = Vector(s3x, s3y, 0)
        sx = Vector(sxx, sxy, 0)

        # Compute vector from s0 to s1
        s0s1 = s1 - s0
        # Compute the vector from s0 to s3
        s0s3 = s3 - s0
        # Compute the vector from s0 to sx
        s0sx = sx - s0

        # cross products
        s0normal = s0s3.cross(s0s1)
        dd = s0s1.cross(s0normal)
        # normalize dd (down dip direction)
        ddn = dd.norm()
        # dot product
        sxdd = ddn.dot(s0sx)

        # get width of quad
        w = utils.get_quad_width(quad)

        # Get weights for top and bottom edge depths
        N = utils.get_quad_normal(quad)
        V = utils.get_vertical_vector(quad)
        dip = np.degrees(np.arccos(Vector.dot(N, V)))
        ws = (w * np.cos(np.radians(dip)))
        wtt = (ws - sxdd) / ws
        wtb = sxdd / ws

        # Compute the depth of of the plane at Px:
        depth = wtt * P0.z + wtb * P3.z * 1000

        return depth

    def getLength(self):
        """
        Compute length of rupture based on top edge in km.

        Returns:
            float: Length of rupture (km).

        """
        flength = 0
        for quad in self._quadrilaterals:
            flength = flength + utils.get_quad_length(quad)
        return flength

    def getWidth(self):
        """
        Compute average rupture width (km) for all quadrilaterals defined for
        the rupture.

        Returns:
            float: Average width in km of all rupture quadrilaterals.
        """
        wsum = 0.0
        for quad in self._quadrilaterals:
            wsum = wsum + utils.get_quad_width(quad)
        mwidth = (wsum / len(self._quadrilaterals)) / 1000.0
        return mwidth

    def getArea(self):
        """
        Compute area of rupture.

        Returns:
            float: Rupture area in square km.

        """
        asum = 0.0
        for quad in self._quadrilaterals:
            width = utils.get_quad_width(quad)
            length = utils.get_quad_length(quad)
            asum = asum + width * length
        return asum

    @classmethod
    def fromTrace(cls, xp0, yp0, xp1, yp1, zp, widths, dips, origin,
                  strike=None, group_index=None, reference=""):
        """
        Create a QuadRupture instance from a set of vertices that define the
        top of the rupture, and an array of widths/dips.

        Each rupture quadrilaterial is defined by specifying the latitude,
        longitude, and depth of the two vertices on the top edges, which must
        have the dame depths. The other verticies are then constructed from
        the top edges and the width and dip of the quadrilateral.

        Args:
            xp0 (array): Array or list of longitudes (floats) of p0.
            yp0 (array): Array or list of latitudes (floats) of p0.
            xp1 (array): Array or list of longitudes (floats) of p1.
            yp1 (array): Array or list of latitudes (floats) of p1.
            zp (array): Array or list of depths for each of the top of rupture
                rectangles (km).
            widths (array): Array of widths for each of rectangle (km).
            dips (array): Array of dips for each of rectangle (degrees).
            origin (Origin): Reference to a ShakeMap origin object.
            strike (array): If None then strike is computed from verticies of
                top edge of each quadrilateral. If a scalar, then all
                quadrilaterals are constructed assuming this strike direction.
                If an array with the same length as the trace coordinates then
                it specifies the strike for each quadrilateral.
            group_index (list): List of integers to indicate group index. If
                None then each quadrilateral is assumed to be in a different
                group since there is no guarantee that any of them are
                continuous.
            reference (str): String explaining where the rupture definition
                came from (publication style reference, etc.).

        Returns:
            QuadRupture instance.

        """
        if len(xp0) == len(yp0) == len(xp1) == len(
                yp1) == len(zp) == len(dips) == len(widths):
            pass
        else:
            raise ShakeLibException(
                'Number of xp0,yp0,xp1,yp1,zp,widths,dips points must be '
                'equal.')
        if strike is None:
            pass
        else:
            if (len(xp0) == len(strike)) | (len(strike) == 1):
                pass
            else:
                raise ShakeLibException(
                    'Strike must be None, scalar, or same length as '
                    'trace coordinates.')

        if group_index is None:
            group_index = np.array(range(len(xp0)))

        # Convert dips to radians
        dips = np.radians(dips)

        # Ensure that all input sequences are numpy arrays
        xp0 = np.array(xp0, dtype='d')
        xp1 = np.array(xp1, dtype='d')
        yp0 = np.array(yp0, dtype='d')
        yp1 = np.array(yp1, dtype='d')
        zp = np.array(zp, dtype='d')
        widths = np.array(widths, dtype='d')
        dips = np.array(dips, dtype='d')

        # Get a projection object
        west = np.min((xp0.min(), xp1.min()))
        east = np.max((xp0.max(), xp1.max()))
        south = np.min((yp0.min(), yp1.min()))
        north = np.max((yp0.max(), yp1.max()))

        # Projected coordinates are in km
        proj = get_orthographic_projection(west, east, north, south)
        xp2 = np.zeros_like(xp0)
        xp3 = np.zeros_like(xp0)
        yp2 = np.zeros_like(xp0)
        yp3 = np.zeros_like(xp0)
        zpdown = np.zeros_like(zp)
        for i in range(0, len(xp0)):
            # Project the top edge coordinates
            p0x, p0y = proj(xp0[i], yp0[i])
            p1x, p1y = proj(xp1[i], yp1[i])

            # Get the rotation angle defined by these two points
            if strike is None:
                dx = p1x - p0x
                dy = p1y - p0y
                theta = np.arctan2(dx, dy)  # theta is angle from north
            elif len(strike) == 1:
                theta = np.radians(strike[0])
            else:
                theta = np.radians(strike[i])

            R = np.array([[np.cos(theta), -np.sin(theta)],
                          [np.sin(theta), np.cos(theta)]])

            # Rotate the top edge points into a new coordinate system (vertical
            # line)
            p0 = np.array([p0x, p0y])
            p1 = np.array([p1x, p1y])
            p0p = np.dot(R, p0)
            p1p = np.dot(R, p1)

            # Get right side coordinates in project, rotated system
            dz = np.sin(dips[i]) * widths[i]
            dx = np.cos(dips[i]) * widths[i]
            p3xp = p0p[0] + dx
            p3yp = p0p[1]
            p2xp = p1p[0] + dx
            p2yp = p1p[1]

            # Get right side coordinates in un-rotated projected system
            p3p = np.array([p3xp, p3yp])
            p2p = np.array([p2xp, p2yp])
            Rback = np.array([[np.cos(-theta), -np.sin(-theta)],
                              [np.sin(-theta), np.cos(-theta)]])
            p3 = np.dot(Rback, p3p)
            p2 = np.dot(Rback, p2p)
            p3x = np.array([p3[0]])
            p3y = np.array([p3[1]])
            p2x = np.array([p2[0]])
            p2y = np.array([p2[1]])

            # project lower edge points back to lat/lon coordinates
            lon3, lat3 = proj(p3x, p3y, reverse=True)
            lon2, lat2 = proj(p2x, p2y, reverse=True)

            xp2[i] = lon2
            xp3[i] = lon3
            yp2[i] = lat2
            yp3[i] = lat3
            zpdown[i] = zp[i] + dz

        # ---------------------------------------------------------------------
        # Create GeoJSON object
        # ---------------------------------------------------------------------

        coords = []
        u_groups = np.unique(group_index)
        n_groups = len(u_groups)
        for i in range(n_groups):
            ind = np.where(u_groups[i] == group_index)[0]
            lons = np.concatenate(
                    [xp0[ind[0]].reshape((1,)),
                     xp1[ind], xp2[ind][::-1],
                     xp3[ind][::-1][-1].reshape((1,)),
                     xp0[ind[0]].reshape((1,))
                     ])
            lats = np.concatenate(
                    [yp0[ind[0]].reshape((1,)),
                     yp1[ind],
                     yp2[ind][::-1],
                     yp3[ind][::-1][-1].reshape((1,)),
                     yp0[ind[0]].reshape((1,))
                     ])
            deps = np.concatenate(
                    [zp[ind[0]].reshape((1,)),
                     zp[ind],
                     zpdown[ind][::-1],
                     zpdown[ind][::-1][-1].reshape((1,)),
                     zp[ind[0]].reshape((1,))])

            poly = []
            for lon, lat, dep in zip(lons, lats, deps):
                poly.append([lon, lat, dep])
            coords.append(poly)

        d = {"type": "FeatureCollection",
             "metadata": {
                 "reference": reference
             },
             "features": [{
                 "type": "Feature",
                 "properties": {
                     "rupture type": "rupture extent"
                 },
                 "geometry": {
                     "type": "MultiPolygon",
                     "coordinates": [coords]
                 }
             }]}

        # Add origin information to metadata
        odict = origin.__dict__
        for k, v in odict.items():
            if isinstance(v, HistoricTime):
                d['metadata'][k] = v.strftime('%Y-%m-%dT%H:%M:%SZ')
            else:
                d['metadata'][k] = v

        return cls(d, origin)

    def writeTextFile(self, rupturefile):
        """
        Write rupture data to rupture file format as defined in ShakeMap
        Software Guide.

        Note that this currently treats each quadrilateral as a separate
        polygon. This needs to be udpated.

        Args:
            rupturefile (str): Filename of output data file OR file-like
                object.

        """
        if not hasattr(rupturefile, 'read'):
            f = open(rupturefile, 'wt')
        else:
            f = rupturefile  # just a reference to the input file-like object
        f.write('#%s\n' % self._reference)
        for quad in self.getQuadrilaterals():
            P0, P1, P2, P3 = quad
            f.write('%.4f %.4f %.4f\n' % (P0.latitude, P0.longitude, P0.depth))
            f.write('%.4f %.4f %.4f\n' % (P1.latitude, P1.longitude, P1.depth))
            f.write('%.4f %.4f %.4f\n' % (P2.latitude, P2.longitude, P2.depth))
            f.write('%.4f %.4f %.4f\n' % (P3.latitude, P3.longitude, P3.depth))
            f.write('%.4f %.4f %.4f\n' % (P0.latitude, P0.longitude, P0.depth))
            f.write(u'>\n')
        if not hasattr(rupturefile, 'read'):
            f.close()

    @classmethod
    def fromVertices(cls,
                     xp0, yp0, zp0, xp1, yp1, zp1,
                     xp2, yp2, zp2, xp3, yp3, zp3,
                     origin,
                     group_index=None,
                     reference=None):
        """
        Create a QuadDrupture instance from the vector of vertices that fully
        define the quadrilaterals. The points p0, ..., p3 are labeled below for
        a trapezoid:

        ::

              p0--------p1
             /          |
            /           |
           p3-----------p2

        All of the following vector arguments must have the same length.

        Args:
            xp0 (array): Array or list of longitudes (floats) of p0.
            yp0 (array): Array or list of latitudes (floats) of p0.
            zp0 (array): Array or list of depths (floats) of p0.
            xp1 (array): Array or list of longitudes (floats) of p1.
            yp1 (array): Array or list of latitudes (floats) of p1.
            zp1 (array): Array or list of depths (floats) of p1.
            xp2 (array): Array or list of longitudes (floats) of p2.
            yp2 (array): Array or list of latitudes (floats) of p2.
            zp2 (array): Array or list of depths (floats) of p2.
            xp3 (array): Array or list of longitudes (floats) of p3.
            yp3 (array): Array or list of latitudes (floats) of p3.
            zp3 (array): Array or list of depths (floats) of p3.
            origin (Origin): Reference to a ShakeMap Origin object.
            group_index (list): List of integers to indicate group index. If
                None then each quadrilateral is assumed to be in a different
                group since there is no guarantee that any of them are
                continuous.
            reference (str): String explaining where the rupture definition
                came from (publication style reference, etc.)

        Returns:
            QuadRupture object, where the rupture is modeled as a series of
                trapezoids.

        """
        if len(xp0) == len(yp0) == len(zp0) == len(xp1) == len(yp1) == \
           len(zp1) == len(xp2) == len(yp2) == len(zp2) == len(xp3) == \
           len(yp3) == len(zp3):
            pass
        else:
            raise ShakeLibException('All vectors specifying quadrilateral '
                                    'vertices must have the same length.')

        nq = len(xp0)
        if group_index is not None:
            if len(group_index) != nq:
                raise Exception(
                    "group_index must have same length as vertices.")
        else:
            group_index = np.array(range(nq))

        xp0 = np.array(xp0, dtype='d')
        yp0 = np.array(yp0, dtype='d')
        zp0 = np.array(zp0, dtype='d')
        xp1 = np.array(xp1, dtype='d')
        yp1 = np.array(yp1, dtype='d')
        zp1 = np.array(zp1, dtype='d')
        xp2 = np.array(xp2, dtype='d')
        yp2 = np.array(yp2, dtype='d')
        zp2 = np.array(zp2, dtype='d')
        xp3 = np.array(xp3, dtype='d')
        yp3 = np.array(yp3, dtype='d')
        zp3 = np.array(zp3, dtype='d')

        # ---------------------------------------------------------------------
        # Create GeoJSON object
        # ---------------------------------------------------------------------

        coords = []
        u_groups = np.unique(group_index)
        n_groups = len(u_groups)
        for i in range(n_groups):
            ind = np.where(u_groups[i] == group_index)[0]
            lons = np.concatenate(
                    [xp0[ind[0]].reshape((1,)),
                     xp1[ind],
                     xp2[ind][::-1],
                     xp3[ind][::-1][-1].reshape((1,)),
                     xp0[ind[0]].reshape((1,))
                     ])
            lats = np.concatenate(
                    [yp0[ind[0]].reshape((1,)),
                     yp1[ind],
                     yp2[ind][::-1],
                     yp3[ind][::-1][-1].reshape((1,)),
                     yp0[ind[0]].reshape((1,))
                     ])
            deps = np.concatenate(
                    [zp0[ind[0]].reshape((1,)),
                     zp1[ind],
                     zp2[ind][::-1],
                     zp3[ind][::-1][-1].reshape((1,)),
                     zp0[ind[0]].reshape((1,))
                     ])

            poly = []
            for lon, lat, dep in zip(lons, lats, deps):
                poly.append([lon, lat, dep])
            coords.append(poly)

        d = {"type": "FeatureCollection",
             "metadata": {
                 "reference": reference
             },
             "features": [{
                 "type": "Feature",
                 "properties": {
                     "rupture type": "rupture extent"
                 },
                 "geometry": {
                     "type": "MultiPolygon",
                     "coordinates": [coords]
                 }
             }]}

        # Add origin information to metadata
        odict = origin.__dict__
        for k, v in odict.items():
            if isinstance(v, HistoricTime):
                d['metadata'][k] = v.strftime('%Y-%m-%dT%H:%M:%SZ')
            else:
                d['metadata'][k] = v
        if hasattr(origin, 'id'):
            d['metadata']['eventid'] = origin.id

        return cls(d, origin)

    def getQuadrilaterals(self):
        """
        Return a list of quadrilaterals.

        Returns:
            list: List of quadrilaterals where each quad is a tuple of four
                `Point <https://github.com/gem/oq-hazardlib/blob/master/openquake/hazardlib/geo/point.py>`__
                objects.
        """  # noqa
        return copy.deepcopy(self._quadrilaterals)

    def getStrike(self):
        """
        Return strike angle. If rupture consists of multiple quadrilaterals,
        the average strike angle, weighted by quad length, is returned.
        Note: for ruptures with quads where the strike angle changes by 180 deg
        due to reverses in dip direction are problematic and not handeled well
        by this algorithm.

        Returns:
            float: Strike angle in degrees.

        """
        nq = len(self._quadrilaterals)
        strikes = np.zeros(nq)
        lengths = np.zeros(nq)
        for i in range(nq):
            P0 = self._quadrilaterals[i][0]
            P1 = self._quadrilaterals[i][1]
            strikes[i] = P0.azimuth(P1)
            lengths[i] = utils.get_quad_length(self._quadrilaterals[i])
        x = np.sin(np.radians(strikes))
        y = np.cos(np.radians(strikes))
        xbar = np.sum(x * lengths) / np.sum(lengths)
        ybar = np.sum(y * lengths) / np.sum(lengths)
        return np.degrees(np.arctan2(xbar, ybar))

    def getDepthToTop(self):
        """
        Determine shallowest vertex of entire rupture.

        :returns:
            Shallowest depth of all vertices (float).
        """
        mindep = 9999999
        for quad in self._quadrilaterals:
            P0, P1, P2, P3 = quad
            depths = np.array([P0.depth, P1.depth, P2.depth, P3.depth])
            if np.min(depths) < mindep:
                mindep = np.min(depths)
        return mindep

    def getDip(self):
        """
        Return average dip of all quadrilaterals in the rupture.

        Returns:
           float: Average dip in degrees.

        """
        dipsum = 0.0
        for quad in self._quadrilaterals:
            N = utils.get_quad_normal(quad)
            V = utils.get_vertical_vector(quad)
            dipsum = dipsum + np.degrees(np.arccos(Vector.dot(N, V)))
        dip = dipsum / len(self._quadrilaterals)
        return dip

    def getIndividualWidths(self):
        """
        Return an array of rupture widths (km), one for each quadrilateral
        defined for the rupture.

        Returns:
            Array of quad widths in km of all rupture quadrilaterals.
        """
        nquad = self.getNumQuads()
        widths = np.zeros(nquad)
        for i in range(nquad):
            q = self._quadrilaterals[i]
            widths[i] = utils.get_quad_width(q) / 1000.0
        return widths

    def getIndividualTopLengths(self):
        """
        Return an array of rupture lengths along top edge (km),
        one for each quadrilateral defined for the rupture.

        :returns:
            Array of lengths in km of top edge of quadrilaterals.
        """
        nquad = self.getNumQuads()
        lengths = np.zeros(nquad)
        for i in range(nquad):
            P0, P1, P2, P3 = self._quadrilaterals[i]
            p0 = Vector.fromPoint(P0)
            p1 = Vector.fromPoint(P1)
            lengths[i] = (p1 - p0).mag() / 1000.0
        return lengths

    @staticmethod
    def _fixStrikeDirection(quad):
        P0, P1, P2, P3 = quad
        eps = 1e-6
        p0 = Vector.fromPoint(P0)  # fromPoint converts to ECEF
        p1 = Vector.fromPoint(P1)
        p2 = Vector.fromPoint(P2)
        p1p0 = p1 - p0
        p2p0 = p2 - p0
        qnv = Vector.cross(p2p0, p1p0).norm()
        tmp = p0 + qnv
        tmplat, tmplon, tmpz = ecef2latlon(tmp.x, tmp.y, tmp.z)
        if (tmpz - P0.depth) < eps:  # If True then do nothing
            fixed = quad
        else:
            newP0 = copy.deepcopy(P1)
            newP1 = copy.deepcopy(P0)
            newP2 = copy.deepcopy(P3)
            newP3 = copy.deepcopy(P2)
            fixed = [newP0, newP1, newP2, newP3]
        return fixed

    def _setQuadrilaterals(self):
        """
        Create internal list of N quadrilaterals. Reverses quad if dip
        direction is incorrect.
        """

        # Make sure arrays are numpy arrays.
        self._lon = np.array(self._lon)
        self._lat = np.array(self._lat)
        self._depth = np.array(self._depth)

        # Find the nans, which tells is where the separate polygons/groups are
        group_ends = np.where(np.isnan(self._lon))[0]
        n_groups = len(group_ends)

        # Check that arrays are the same length
        if len(self._lon) != len(self._lat) != len(self._depth):
            raise IndexError(
                'Length of input lon, lat, depth arrays must be equal')

        # Construct quads
        group_start = 0

        self._quadrilaterals = []
        self._group_index = []
        groupind = 0
        for i in range(n_groups):
            lonseg = self._lon[group_start:group_ends[i]]
            latseg = self._lat[group_start:group_ends[i]]
            depthseg = self._depth[group_start:group_ends[i]]

            # Each group can have many contiguous quadrilaterals defined in it
            # separations (nans) between segments mean that segments are not
            # contiguous.

            npoints = len(lonseg)
            nquads = int((npoints - 4) / 2) + 1
            quad_start = 0
            quad_end = -1
            for j in range(nquads):
                P0 = Point(lonseg[quad_start],
                           latseg[quad_start],
                           depthseg[quad_start])
                P1 = Point(lonseg[quad_start + 1],
                           latseg[quad_start + 1],
                           depthseg[quad_start + 1])
                P2 = Point(lonseg[quad_end - 1],
                           latseg[quad_end - 1],
                           depthseg[quad_end - 1])
                P3 = Point(lonseg[quad_end],
                           latseg[quad_end],
                           depthseg[quad_end])
                quad = [P0, P1, P2, P3]

                # Enforce plane by moving P2 -- already close because of check
                # in read_rupture_file/is_quadrupture_class/is_quad

                dummy, fixed_quad = utils.is_quad(quad)

                # Reverse quad if necessary
                fixed_quad = self._fixStrikeDirection(fixed_quad)

                self._quadrilaterals.append(fixed_quad)

                quad_start = quad_start + 1
                quad_end = quad_end - 1

            group_start = group_ends[i] + 1
            self._group_index.extend([groupind] * nquads)
            groupind = groupind + 1

    def _getGroupIndex(self):
        """
        Return a list of segment group indexes.

        Returns:
            list: Segment group indexes; length equals the number of
                quadrilaterals.
        """
        return copy.deepcopy(self._group_index)

    @property
    def lats(self):
        """
        Return an array of latitudes for the rupture verticies arranged for
        plotting purposes; will give an outline of each group connected
        segments.

        Returns:
            array: Numpy array of closed-loop latitude values; disconnected
                segments are separated by nans.
        """
        lats = []
        quads = self.getQuadrilaterals()
        groups = self._getGroupIndex()
        u_groups = np.unique(groups)
        ng = len(u_groups)
        for i in range(ng):
            q_ind = np.where(groups == u_groups[i])[0]
            nq = len(q_ind)
            top_lats = []
            bot_lats = []
            for j in range(nq):
                if j == 0:
                    top0 = [quads[q_ind[j]][0].latitude]
                    bot0 = [quads[q_ind[j]][3].latitude]
                    top_lats = top_lats + top0
                    bot_lats = bot_lats + bot0
                top_lats = top_lats + [quads[q_ind[j]][1].latitude]
                bot_lats = bot_lats + [quads[q_ind[j]][2].latitude]
            lats = lats + top_lats + bot_lats[::-1] + top0 + [np.nan]

        return np.array(lats)

    @property
    def lons(self):
        """
        Return an array of longitudes for the rupture verticies arranged for
        plotting purposes; will give an outline of each group connected
        segments.

        Returns:
            array: Numpy array of closed-loop longitude values; disconnected
                segments are separated by nans.
        """
        lons = []
        quads = self.getQuadrilaterals()
        groups = self._getGroupIndex()
        u_groups = np.unique(groups)
        ng = len(u_groups)
        for i in range(ng):
            q_ind = np.where(groups == u_groups[i])[0]
            nq = len(q_ind)
            top_lons = []
            bot_lons = []
            for j in range(nq):
                if j == 0:
                    top0 = [quads[q_ind[j]][0].longitude]
                    bot0 = [quads[q_ind[j]][3].longitude]
                    top_lons = top_lons + top0
                    bot_lons = bot_lons + bot0
                top_lons = top_lons + [quads[q_ind[j]][1].longitude]
                bot_lons = bot_lons + [quads[q_ind[j]][2].longitude]
            lons = lons + top_lons + bot_lons[::-1] + top0 + [np.nan]
        return np.array(lons)

    @property
    def depths(self):
        """
        Return an array of depths for the rupture verticies arranged for
        plotting purposes; will give an outline of each group connected
        segments.

        Returns:
            array: Numpy array of closed-loop depths; disconnected
                segments are separated by nans.
        """
        deps = []
        quads = self.getQuadrilaterals()
        groups = self._getGroupIndex()
        u_groups = np.unique(groups)
        ng = len(u_groups)
        for i in range(ng):
            q_ind = np.where(groups == u_groups[i])[0]
            nq = len(q_ind)
            top_deps = []
            bot_deps = []
            for j in range(nq):
                if j == 0:
                    top0 = [quads[q_ind[j]][0].depth]
                    bot0 = [quads[q_ind[j]][3].depth]
                    top_deps = top_deps + top0
                    bot_deps = bot_deps + bot0
                top_deps = top_deps + [quads[q_ind[j]][1].depth]
                bot_deps = bot_deps + [quads[q_ind[j]][2].depth]
            deps = deps + top_deps + bot_deps[::-1] + top0 + [np.nan]

        return np.array(deps)

    def getDeps(self):
        """
        Return a copy of the array of depths for the rupture verticies.

        Returns:
            array: Numpy array of latitude values.
        """
        return self._depth.copy()

    def getNumGroups(self):
        """
        Return a count of the number of rupture groups.

        Returns:
            int:Rnumber of rupture groups.

        """
        return len(np.unique(self._group_index))

    def getNumQuads(self):
        """
        Return a count of the number of rupture quadrilaterals.

        Returns:
            int: Number of rupture quadrilaterals.
        """
        return len(self._quadrilaterals)

    def getRuptureAsArrays(self):
        """
        Return a 3-tuple of numpy arrays indicating X, Y, Z (lon,lat,depth)
        coordinates. Rupture groups are separated by numpy.NaN values.

        Returns:
            tuple: 3-tuple of numpy arrays indicating X,Y,Z (lon,lat,depth)
                coordinates.
        """
        return (np.array(self._lon),
                np.array(self._lat),
                np.array(self._depth))

    def getRuptureAsMesh(self):
        """
        Return rupture segments as a OQ-Hazardlib Mesh object.

        Returns:
            Mesh (https://github.com/gem/oq-hazardlib/blob/master/openquake/hazardlib/geo/mesh.py)
        """  # noqa
        rupture = Mesh(self._lon, self._lat, self._depth)
        return rupture

    def computeRjb(self, lon, lat, depth):
        """
        Method for computing Joyner-Boore distance.

        Args:
            lon (array): Numpy array of longitudes.
            lat (array): Numpy array of latitudes.
            depth (array): Numpy array of depths (km; positive down).

        Returns:
           array: Joyner-Boore distance (km).

        """

        # ---------------------------------------------------------------------
        # Sort out sites
        # ---------------------------------------------------------------------
        oldshape = lon.shape

        if len(oldshape) == 2:
            newshape = (oldshape[0] * oldshape[1], 1)
        else:
            newshape = (oldshape[0], 1)

        x, y, z = latlon2ecef(lat, lon, depth)
        x.shape = newshape
        y.shape = newshape
        z.shape = newshape
        sites_ecef = np.hstack((x, y, z))

        minrjb = np.ones(newshape, dtype=lon.dtype) * 1e16
        quads = self.getQuadrilaterals()

        for i in range(len(quads)):
            P0, P1, P2, P3 = quads[i]
            S0 = copy.deepcopy(P0)
            S1 = copy.deepcopy(P1)
            S2 = copy.deepcopy(P2)
            S3 = copy.deepcopy(P3)
            S0.depth = 0.0
            S1.depth = 0.0
            S2.depth = 0.0
            S3.depth = 0.0
            squad = [S0, S1, S2, S3]
            rjbdist = utils._quad_distance(squad, sites_ecef, horizontal=True)
            minrjb = np.minimum(minrjb, rjbdist)

        minrjb = minrjb.reshape(oldshape)
        return minrjb

    def computeRrup(self, lon, lat, depth):
        """
        Method for computing rupture distance.

        Args:
            lon (array): Numpy array of longitudes.
            lat (array): Numpy array of latitudes.
            depth (array): Numpy array of depths (km; positive down).

        Returns:
           array: Rupture distance (km).

        """

        # ---------------------------------------------------------------------
        # Sort out sites
        # ---------------------------------------------------------------------
        oldshape = lon.shape

        if len(oldshape) == 2:
            newshape = (oldshape[0] * oldshape[1], 1)
        else:
            newshape = (oldshape[0], 1)

        x, y, z = latlon2ecef(lat, lon, depth)
        x.shape = newshape
        y.shape = newshape
        z.shape = newshape
        sites_ecef = np.hstack((x, y, z))

        minrrup = np.ones(newshape, dtype=lon.dtype) * 1e16
        quads = self.getQuadrilaterals()

        for i in range(len(quads)):
            rrupdist = utils._quad_distance(quads[i], sites_ecef)
            minrrup = np.minimum(minrrup, rrupdist)

        minrrup = minrrup.reshape(oldshape)
        return minrrup

    def computeGC2(self, lon, lat, depth):
        """
        Method for computing version 2 of the Generalized Coordinate system
        (GC2) by Spudich and Chiou OFR 2015-1028.

        Args:
            lon (array): Numpy array of longitudes.
            lat (array): Numpy array of latitudes.
            depth (array): Numpy array of depths (km; positive down).

        Returns:
            dict: Dictionary with keys for each of the GC2-related distances,
                which include 'rx', 'ry', 'ry0', 'U', 'T'.
        """
        # This just hands off to the module-level method
        dict = gc2._computeGC2(self, lon, lat, depth)
        return dict


# -*- coding: utf-8 -*-
"""
Spyder Editor

This is a temporary script file.
"""

# standard imports
import numpy as np
import matplotlib.pyplot as plt

# custom imports
import GaN_fun
import GaN_type_peak_assignments

import peak_param_determination as ppd
from histogram_functions import bin_dat

plt.close('all')

# Read in data
epos = GaN_fun.load_epos(run_number='R44_03146',
                         epos_trim=[5000, 5000],
                         fig_idx=999)

pk_data = GaN_type_peak_assignments.In_doped_GaN()
bg_rois=[[0.4,0.9]]
#bg_rois=[[10,11]]

pk_params, glob_bg_param, Ga1p_idxs, Ga2p_idxs = GaN_fun.fit_spectrum(
        epos=epos,
        pk_data=pk_data,
        peak_height_fraction=0.1,
        bg_rois=bg_rois)

cts, compositions, is_peak = GaN_fun.count_and_get_compositions(
        epos=epos,
        pk_data=pk_data,
        pk_params=pk_params,
        glob_bg_param=glob_bg_param,
        bg_frac=1,
        noise_threshhold=2)

# Print out the composition of the full dataset
ppd.pretty_print_compositions(compositions,pk_data)

# Plot the full spectrum
xs, ys_sm = GaN_fun.bin_and_smooth_spectrum(epos=epos,
                                            user_roi=[0,150],
                                            bin_wid_mDa=30,
                                            smooth_wid_mDa=-1)

fig = plt.figure(num=1)
fig.set_size_inches(w=6.69, h=3)
fig.clear()
ax = fig.gca()

ax.plot(xs, ys_sm, lw=1, label='full spec')

glob_bg = ppd.physics_bg(xs,glob_bg_param)
ax.plot(xs, glob_bg, lw=1, label='bg', alpha=1)

ax.set_xlim(0,120)
ax.set_ylim(1e0,1e5)
ax.grid(b=True)
ax.set(xlabel='m/z', ylabel='counts')
ax.set_yscale('log')
ax.legend()
fig.tight_layout()

fig.savefig('InGaN_full_spectrum.pdf')
fig.savefig('InGaN_full_spectrum.jpg', dpi=300)

# Find the pole center and show it
m2q_roi = [3, 100]
sel_idxs = np.where((epos['m2q']>m2q_roi[0]) & (epos['m2q']<m2q_roi[1]))
xc,yc = GaN_fun.mean_shift(epos['x_det'][sel_idxs],epos['y_det'][sel_idxs])

# Find all the Indium events
is_In = ~np.isfinite(epos['m2q'])
for pk,param in zip(pk_data,pk_params):
    if pk['In']>0:
        is_In = is_In | ((epos['m2q']>=param['pre_rng']) & (epos['m2q']<=param['post_rng']))

# Fit the QW to a plane and rotate the point cloud to 'flatten' wrt z-axis
p = GaN_fun.qw_plane_fit(epos['x'][is_In],epos['y'][is_In],epos['z'][is_In],np.array([0,0,1.8]))
epos['x'],epos['y'],epos['z'] = GaN_fun.rotate_data_flat(p,epos['x'],epos['y'],epos['z'])

# USER DEFINED!!!
# Z-ROIs TO TAKE COMPOSITIONS IN
z_roi_qw = [38, 38.6]
z_roi_buf = [43, 48]
z_roi_gan = [40, 42]

# Find all the Ga events
is_Ga = ~np.isfinite(epos['m2q'])
for pk,param in zip(pk_data,pk_params):
    if pk['Ga']>0:
        is_Ga = is_Ga | ((epos['m2q']>=param['pre_rng']) & (epos['m2q']<=param['post_rng']))

import colorcet as cc
cm=cc.cm.glasbey

# Plot the 'flat' interface to verify vector algebra didn't go awry
fig = plt.figure(num=11)
fig.clear()
ax = fig.gca()
#ax.plot(epos['x'][is_In],epos['z'][is_In],'.')
#ax.plot(epos['y'][is_In],epos['z'][is_In],'.')
ax.plot(epos['x'][is_Ga],epos['z'][is_Ga],'.', alpha=0.05, color=cm(2), ms=2)
ax.plot(epos['x'][is_In],epos['z'][is_In],'.', alpha=0.5, color=cm(10), ms=2)

# Plot the ROIs for visual inspection
ax.fill([-8,-8, 8,8], [z_roi_qw[0], z_roi_qw[1], z_roi_qw[1], z_roi_qw[0]], color=[1,0,0,0.5])
ax.fill([-8,-8, 8,8], [z_roi_buf[0], z_roi_buf[1], z_roi_buf[1], z_roi_buf[0]], color=[1,0,0,0.5])
ax.fill([-8,-8, 8,8], [z_roi_gan[0], z_roi_gan[1], z_roi_gan[1], z_roi_gan[0]], color=[1,0,0,0.5])

# Calculate the QW composition
is_roi = (epos['z']>=z_roi_qw[0]) & (epos['z']<=z_roi_qw[1])
sub_epos = epos[is_roi]

bg_frac_roi = [120,150]
bg_frac = np.sum((sub_epos['m2q']>bg_frac_roi[0]) & (sub_epos['m2q']<bg_frac_roi[1])) \
                    / np.sum((epos['m2q']>bg_frac_roi[0]) & (epos['m2q']<bg_frac_roi[1]))

# Count the peaks, local bg, and global bg.  Ignore the local bg based info
cts, compositions, is_peak = GaN_fun.count_and_get_compositions(
        epos=sub_epos,
        pk_data=pk_data,
        pk_params=pk_params,
        glob_bg_param=glob_bg_param,
        bg_frac=bg_frac,
        noise_threshhold=2)

ppd.pretty_print_compositions(compositions,pk_data)

# Plot the QW spectrum
xs, ys_sm = GaN_fun.bin_and_smooth_spectrum(epos=sub_epos,
                                            user_roi=[0,150],
                                            bin_wid_mDa=30,
                                            smooth_wid_mDa=-1)

fig = plt.figure(num=2)
fig.set_size_inches(w=6.69, h=3)
fig.clear()
ax = fig.gca()

ax.plot(xs, ys_sm, lw=1, label='QW spec',color='k')

glob_bg = ppd.physics_bg(xs,bg_frac*glob_bg_param)
ax.plot(xs, glob_bg, lw=1, label='bg', alpha=1,color='r')

ax.set_xlim(0,120)
ax.set_ylim(1e0,1e3)
ax.grid(b=True)
ax.set(xlabel='m/z', ylabel='counts')
ax.set_yscale('log')
ax.legend()
fig.tight_layout()

fig.savefig('InGaN_QW_spectrum.pdf')
fig.savefig('InGaN_QW_spectrum.jpg', dpi=300)

# Calculate the buffer composition
is_roi = (epos['z']>=z_roi_buf[0]) & (epos['z']<=z_roi_buf[1])
sub_epos = epos[is_roi]

bg_frac_roi = [120,150]
bg_frac = np.sum((sub_epos['m2q']>bg_frac_roi[0]) & (sub_epos['m2q']<bg_frac_roi[1])) \
                    / np.sum((epos['m2q']>bg_frac_roi[0]) & (epos['m2q']<bg_frac_roi[1]))

# Count the peaks, local bg, and global bg.  Ignore the local bg based info
cts, compositions, is_peak = GaN_fun.count_and_get_compositions(
        epos=sub_epos,
        pk_data=pk_data,
        pk_params=pk_params,
        glob_bg_param=glob_bg_param,
        bg_frac=bg_frac,
        noise_threshhold=2)

ppd.pretty_print_compositions(compositions,pk_data)

# Plot the buffer spectrum
xs, ys_sm = GaN_fun.bin_and_smooth_spectrum(epos=sub_epos,
                                            user_roi=[0,150],
                                            bin_wid_mDa=30,
                                            smooth_wid_mDa=-1)

fig = plt.figure(num=3)
fig.set_size_inches(w=6.69, h=3)
fig.clear()
ax = fig.gca()

ax.plot(xs, ys_sm, lw=1, label='buffer spec',color='k')

glob_bg = ppd.physics_bg(xs,bg_frac*glob_bg_param)
ax.plot(xs, glob_bg, lw=1, label='bg', alpha=1,color='r')

ax.set_xlim(0,120)
ax.set_ylim(1e0,1e4)
ax.grid(b=True)
ax.set(xlabel='m/z', ylabel='counts')
ax.set_yscale('log')
ax.legend()
fig.tight_layout()

fig.savefig('InGaN_buffer_spectrum.pdf')
fig.savefig('InGaN_buffer_spectrum.jpg', dpi=300)

# Calculate the barrier composition
is_roi = (epos['z']>=z_roi_gan[0]) & (epos['z']<=z_roi_gan[1])
sub_epos = epos[is_roi]

bg_frac_roi = [120,150]
bg_frac = np.sum((sub_epos['m2q']>bg_frac_roi[0]) & (sub_epos['m2q']<bg_frac_roi[1])) \
                    / np.sum((epos['m2q']>bg_frac_roi[0]) & (epos['m2q']<bg_frac_roi[1]))

# Count the peaks, local bg, and global bg.  Ignore the local bg based info
cts, compositions, is_peak = GaN_fun.count_and_get_compositions(
        epos=sub_epos,
        pk_data=pk_data,
        pk_params=pk_params,
        glob_bg_param=glob_bg_param,
        bg_frac=bg_frac,
        noise_threshhold=2)

ppd.pretty_print_compositions(compositions,pk_data)

# Plot the barrier spectrum
xs, ys_sm = GaN_fun.bin_and_smooth_spectrum(epos=sub_epos,
                                            user_roi=[0,150],
                                            bin_wid_mDa=30,
                                            smooth_wid_mDa=-1)

fig = plt.figure(num=4)
fig.set_size_inches(w=6.69, h=3)
fig.clear()
ax = fig.gca()

ax.plot(xs, ys_sm, lw=1, label='barrier spec',color='k')

glob_bg = ppd.physics_bg(xs,bg_frac*glob_bg_param)
ax.plot(xs, glob_bg, lw=1, label='bg$\pm', alpha=1,color='r')

ax.set_xlim(0,120)
ax.set_ylim(1e0,1e4)
ax.grid(b=True)
ax.set(xlabel='m/z', ylabel='counts')
ax.set_yscale('log')
ax.legend()
fig.tight_layout()

fig.savefig('InGaN_barrier_spectrum.pdf')
fig.savefig('InGaN_barrier_spectrum.jpg', dpi=300)


# Copyright 2018 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np

QUANTIZATION_BYTES_TO_DTYPES = {1: np.uint8, 2: np.uint16}

def quantize_weights(data, quantization_dtype):
  """Quantizes the weights by linearly re-scaling across available bits.

  The weights are quantized by linearly re-scaling the values between the
  minimum and maximum value, and representing them with the number of bits
  provided by the `quantization_dtype`.

  In order to guarantee that 0 is perfectly represented by one of the quantized
  values, the range is "nudged" in the same manner as in TF-Lite.

  Weights can be de-quantized by multiplying by the returned `scale` and adding
  `min`.

  Args:
    data: A numpy array of dtype 'float32' or 'int32'.
    quantization_dtype: A numpy dtype to quantize weights to. Only np.uint8 and
      np.uint16 are supported.

  Returns:
    quantized_data: The quantized weights as a numpy array with dtype
      `quantization_dtype`.
    scale: The linearly scaling constant used for quantization.
    min_val: The minimum value of the linear range.
  Raises:
    ValueError: if `quantization_dtype` is not a valid type.
  """
  if quantization_dtype not in QUANTIZATION_BYTES_TO_DTYPES.values():
    raise ValueError('Invalid `quantization_dtype`: %r' % quantization_dtype)

  # Compute the min and max for the group.
  min_val = data.min().astype(np.float64)
  max_val = data.max().astype(np.float64)
  if min_val == max_val:
    # If there is only a single value, we can represent everything as zeros.
    quantized_data = np.zeros_like(data, dtype=quantization_dtype)
    scale = 1.0
  else:
    # Quantize data.
    scale, min_val, max_val = _get_quantization_range(
        min_val, max_val, quantization_dtype)
    quantized_data = np.round(
        (data.clip(min_val, max_val) - min_val) / scale).astype(
            quantization_dtype)

  return quantized_data, scale, min_val

def dequantize_weights(
    quantized_data, scale, min_val, original_dtype=np.float32):
  return np.round(quantized_data * scale + min_val).astype(original_dtype)

def _get_quantization_range(min_val, max_val, quantization_dtype):
  """Computes quantization range to ensure that zero is represented if covered.

  Gymnastics with nudged zero point is to ensure that real zero maps to an
  integer, which is required for e.g. zero-padding in convolutional layers.

  Based on `NudgeQuantizationRange` in
  tensorflow/contrib/lite/kernels/internal/quantization_util.h, except we do not
  nudge if 0 is not in the range.

  Args:
    min_val: The actual minimum value of the data.
    max_val: The actual maximum value of the data.
    quantization_dtype: A numpy dtype to quantize weights to. Only np.uint8 and
      np.uint16 are supported.

  Returns:
    scale: The linear scaling constant used for quantization.
    nudged_min: The adjusted minimum value to ensure zero is represented, if
      covered.
    nudged_max: The adjusted maximum value to ensure zero is represented, if
      covered.
  Raises:
    ValueError: if `quantization_dtype` is not a valid type.
  """
  if quantization_dtype not in QUANTIZATION_BYTES_TO_DTYPES.values():
    raise ValueError('Invalid `quantization_dtype`: %r' % quantization_dtype)

  quant_max = np.iinfo(quantization_dtype).max
  scale = (max_val - min_val) / quant_max

  if min_val <= 0 <= max_val:
    quantized_zero_point = (0 - min_val) / scale
    nudged_zero_point = np.round(quantized_zero_point)

    # Solve `0 = nudged_zero_point * scale + nudged_min` for `nudged_min`.
    nudged_min = -nudged_zero_point * scale
    nudged_max = quant_max * scale + nudged_min
  else:
    nudged_min, nudged_max = min_val, max_val

  return scale, nudged_min, nudged_max


# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Base classes for probability distributions."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import abc
import contextlib
import types

import numpy as np
import six

from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops.distributions import kullback_leibler
from tensorflow.python.ops.distributions import util
from tensorflow.python.util import tf_inspect
from tensorflow.python.util.tf_export import tf_export

__all__ = [
    "ReparameterizationType",
    "FULLY_REPARAMETERIZED",
    "NOT_REPARAMETERIZED",
    "Distribution",
]

_DISTRIBUTION_PUBLIC_METHOD_WRAPPERS = [
    "batch_shape",
    "batch_shape_tensor",
    "cdf",
    "covariance",
    "cross_entropy",
    "entropy",
    "event_shape",
    "event_shape_tensor",
    "kl_divergence",
    "log_cdf",
    "log_prob",
    "log_survival_function",
    "mean",
    "mode",
    "prob",
    "sample",
    "stddev",
    "survival_function",
    "variance",
]

@six.add_metaclass(abc.ABCMeta)
class _BaseDistribution(object):
  """Abstract base class needed for resolving subclass hierarchy."""
  pass

def _copy_fn(fn):
  """Create a deep copy of fn.

  Args:
    fn: a callable

  Returns:
    A `FunctionType`: a deep copy of fn.

  Raises:
    TypeError: if `fn` is not a callable.
  """
  if not callable(fn):
    raise TypeError("fn is not callable: %s" % fn)
  # The blessed way to copy a function. copy.deepcopy fails to create a
  # non-reference copy. Since:
  #   types.FunctionType == type(lambda: None),
  # and the docstring for the function type states:
  #
  #   function(code, globals[, name[, argdefs[, closure]]])
  #
  #   Create a function object from a code object and a dictionary.
  #   ...
  #
  # Here we can use this to create a new function with the old function's
  # code, globals, closure, etc.
  return types.FunctionType(
      code=fn.__code__, globals=fn.__globals__,
      name=fn.__name__, argdefs=fn.__defaults__,
      closure=fn.__closure__)

def _update_docstring(old_str, append_str):
  """Update old_str by inserting append_str just before the "Args:" section."""
  old_str = old_str or ""
  old_str_lines = old_str.split("\n")

  # Step 0: Prepend spaces to all lines of append_str. This is
  # necessary for correct markdown generation.
  append_str = "\n".join("    %s" % line for line in append_str.split("\n"))

  # Step 1: Find mention of "Args":
  has_args_ix = [
      ix for ix, line in enumerate(old_str_lines)
      if line.strip().lower() == "args:"]
  if has_args_ix:
    final_args_ix = has_args_ix[-1]
    return ("\n".join(old_str_lines[:final_args_ix])
            + "\n\n" + append_str + "\n\n"
            + "\n".join(old_str_lines[final_args_ix:]))
  else:
    return old_str + "\n\n" + append_str

class _DistributionMeta(abc.ABCMeta):

  def __new__(mcs, classname, baseclasses, attrs):
    """Control the creation of subclasses of the Distribution class.

    The main purpose of this method is to properly propagate docstrings
    from private Distribution methods, like `_log_prob`, into their
    public wrappers as inherited by the Distribution base class
    (e.g. `log_prob`).

    Args:
      classname: The name of the subclass being created.
      baseclasses: A tuple of parent classes.
      attrs: A dict mapping new attributes to their values.

    Returns:
      The class object.

    Raises:
      TypeError: If `Distribution` is not a subclass of `BaseDistribution`, or
        the new class is derived via multiple inheritance and the first
        parent class is not a subclass of `BaseDistribution`.
      AttributeError:  If `Distribution` does not implement e.g. `log_prob`.
      ValueError:  If a `Distribution` public method lacks a docstring.
    """
    if not baseclasses:  # Nothing to be done for Distribution
      raise TypeError("Expected non-empty baseclass. Does Distribution "
                      "not subclass _BaseDistribution?")
    which_base = [
        base for base in baseclasses
        if base == _BaseDistribution or issubclass(base, Distribution)]
    base = which_base[0]
    if base == _BaseDistribution:  # Nothing to be done for Distribution
      return abc.ABCMeta.__new__(mcs, classname, baseclasses, attrs)
    if not issubclass(base, Distribution):
      raise TypeError("First parent class declared for %s must be "
                      "Distribution, but saw '%s'" % (classname, base.__name__))
    for attr in _DISTRIBUTION_PUBLIC_METHOD_WRAPPERS:
      special_attr = "_%s" % attr
      class_attr_value = attrs.get(attr, None)
      if attr in attrs:
        # The method is being overridden, do not update its docstring
        continue
      base_attr_value = getattr(base, attr, None)
      if not base_attr_value:
        raise AttributeError(
            "Internal error: expected base class '%s' to implement method '%s'"
            % (base.__name__, attr))
      class_special_attr_value = attrs.get(special_attr, None)
      if class_special_attr_value is None:
        # No _special method available, no need to update the docstring.
        continue
      class_special_attr_docstring = tf_inspect.getdoc(class_special_attr_value)
      if not class_special_attr_docstring:
        # No docstring to append.
        continue
      class_attr_value = _copy_fn(base_attr_value)
      class_attr_docstring = tf_inspect.getdoc(base_attr_value)
      if class_attr_docstring is None:
        raise ValueError(
            "Expected base class fn to contain a docstring: %s.%s"
            % (base.__name__, attr))
      class_attr_value.__doc__ = _update_docstring(
          class_attr_value.__doc__,
          ("Additional documentation from `%s`:\n\n%s"
           % (classname, class_special_attr_docstring)))
      attrs[attr] = class_attr_value

    return abc.ABCMeta.__new__(mcs, classname, baseclasses, attrs)

@tf_export("distributions.ReparameterizationType")
class ReparameterizationType(object):
  """Instances of this class represent how sampling is reparameterized.

  Two static instances exist in the distributions library, signifying
  one of two possible properties for samples from a distribution:

  `FULLY_REPARAMETERIZED`: Samples from the distribution are fully
    reparameterized, and straight-through gradients are supported.

  `NOT_REPARAMETERIZED`: Samples from the distribution are not fully
    reparameterized, and straight-through gradients are either partially
    unsupported or are not supported at all. In this case, for purposes of
    e.g. RL or variational inference, it is generally safest to wrap the
    sample results in a `stop_gradients` call and use policy
    gradients / surrogate loss instead.
  """

  def __init__(self, rep_type):
    self._rep_type = rep_type

  def __repr__(self):
    return "<Reparameteriation Type: %s>" % self._rep_type

  def __eq__(self, other):
    """Determine if this `ReparameterizationType` is equal to another.

    Since RepaparameterizationType instances are constant static global
    instances, equality checks if two instances' id() values are equal.

    Args:
      other: Object to compare against.

    Returns:
      `self is other`.
    """
    return self is other

# Fully reparameterized distribution: samples from a fully
# reparameterized distribution support straight-through gradients with
# respect to all parameters.
FULLY_REPARAMETERIZED = ReparameterizationType("FULLY_REPARAMETERIZED")
tf_export("distributions.FULLY_REPARAMETERIZED").export_constant(
    __name__, "FULLY_REPARAMETERIZED")

# Not reparameterized distribution: samples from a non-
# reparameterized distribution do not support straight-through gradients for
# at least some of the parameters.
NOT_REPARAMETERIZED = ReparameterizationType("NOT_REPARAMETERIZED")
tf_export("distributions.NOT_REPARAMETERIZED").export_constant(
    __name__, "NOT_REPARAMETERIZED")

@six.add_metaclass(_DistributionMeta)
@tf_export("distributions.Distribution")
class Distribution(_BaseDistribution):
  """A generic probability distribution base class.

  `Distribution` is a base class for constructing and organizing properties
  (e.g., mean, variance) of random variables (e.g, Bernoulli, Gaussian).

  #### Subclassing

  Subclasses are expected to implement a leading-underscore version of the
  same-named function. The argument signature should be identical except for
  the omission of `name="..."`. For example, to enable `log_prob(value,
  name="log_prob")` a subclass should implement `_log_prob(value)`.

  Subclasses can append to public-level docstrings by providing
  docstrings for their method specializations. For example:

  ```python
  @util.AppendDocstring("Some other details.")
  def _log_prob(self, value):
    ...
  ```

  would add the string "Some other details." to the `log_prob` function
  docstring. This is implemented as a simple decorator to avoid python
  linter complaining about missing Args/Returns/Raises sections in the
  partial docstrings.

  #### Broadcasting, batching, and shapes

  All distributions support batches of independent distributions of that type.
  The batch shape is determined by broadcasting together the parameters.

  The shape of arguments to `__init__`, `cdf`, `log_cdf`, `prob`, and
  `log_prob` reflect this broadcasting, as does the return value of `sample` and
  `sample_n`.

  `sample_n_shape = [n] + batch_shape + event_shape`, where `sample_n_shape` is
  the shape of the `Tensor` returned from `sample_n`, `n` is the number of
  samples, `batch_shape` defines how many independent distributions there are,
  and `event_shape` defines the shape of samples from each of those independent
  distributions. Samples are independent along the `batch_shape` dimensions, but
  not necessarily so along the `event_shape` dimensions (depending on the
  particulars of the underlying distribution).

  Using the `Uniform` distribution as an example:

  ```python
  minval = 3.0
  maxval = [[4.0, 6.0],
            [10.0, 12.0]]

  # Broadcasting:
  # This instance represents 4 Uniform distributions. Each has a lower bound at
  # 3.0 as the `minval` parameter was broadcasted to match `maxval`'s shape.
  u = Uniform(minval, maxval)

  # `event_shape` is `TensorShape([])`.
  event_shape = u.event_shape
  # `event_shape_t` is a `Tensor` which will evaluate to [].
  event_shape_t = u.event_shape_tensor()

  # Sampling returns a sample per distribution. `samples` has shape
  # [5, 2, 2], which is [n] + batch_shape + event_shape, where n=5,
  # batch_shape=[2, 2], and event_shape=[].
  samples = u.sample_n(5)

  # The broadcasting holds across methods. Here we use `cdf` as an example. The
  # same holds for `log_cdf` and the likelihood functions.

  # `cum_prob` has shape [2, 2] as the `value` argument was broadcasted to the
  # shape of the `Uniform` instance.
  cum_prob_broadcast = u.cdf(4.0)

  # `cum_prob`'s shape is [2, 2], one per distribution. No broadcasting
  # occurred.
  cum_prob_per_dist = u.cdf([[4.0, 5.0],
                             [6.0, 7.0]])

  # INVALID as the `value` argument is not broadcastable to the distribution's
  # shape.
  cum_prob_invalid = u.cdf([4.0, 5.0, 6.0])
  ```

  #### Shapes

  There are three important concepts associated with TensorFlow Distributions
  shapes:
  - Event shape describes the shape of a single draw from the distribution;
    it may be dependent across dimensions. For scalar distributions, the event
    shape is `[]`. For a 5-dimensional MultivariateNormal, the event shape is
    `[5]`.
  - Batch shape describes independent, not identically distributed draws, aka a
    "collection" or "bunch" of distributions.
  - Sample shape describes independent, identically distributed draws of batches
    from the distribution family.

  The event shape and the batch shape are properties of a Distribution object,
  whereas the sample shape is associated with a specific call to `sample` or
  `log_prob`.

  For detailed usage examples of TensorFlow Distributions shapes, see
  [this tutorial](
  https://github.com/tensorflow/probability/blob/master/tensorflow_probability/examples/jupyter_notebooks/Understanding_TensorFlow_Distributions_Shapes.ipynb)

  #### Parameter values leading to undefined statistics or distributions.

  Some distributions do not have well-defined statistics for all initialization
  parameter values. For example, the beta distribution is parameterized by
  positive real numbers `concentration1` and `concentration0`, and does not have
  well-defined mode if `concentration1 < 1` or `concentration0 < 1`.

  The user is given the option of raising an exception or returning `NaN`.

  ```python
  a = tf.exp(tf.matmul(logits, weights_a))
  b = tf.exp(tf.matmul(logits, weights_b))

  # Will raise exception if ANY batch member has a < 1 or b < 1.
  dist = distributions.beta(a, b, allow_nan_stats=False)
  mode = dist.mode().eval()

  # Will return NaN for batch members with either a < 1 or b < 1.
  dist = distributions.beta(a, b, allow_nan_stats=True)  # Default behavior
  mode = dist.mode().eval()
  ```

  In all cases, an exception is raised if *invalid* parameters are passed, e.g.

  ```python
  # Will raise an exception if any Op is run.
  negative_a = -1.0 * a  # beta distribution by definition has a > 0.
  dist = distributions.beta(negative_a, b, allow_nan_stats=True)
  dist.mean().eval()
  ```

  """

  def __init__(self,
               dtype,
               reparameterization_type,
               validate_args,
               allow_nan_stats,
               parameters=None,
               graph_parents=None,
               name=None):
    """Constructs the `Distribution`.

    **This is a private method for subclass use.**

    Args:
      dtype: The type of the event samples. `None` implies no type-enforcement.
      reparameterization_type: Instance of `ReparameterizationType`.
        If `distributions.FULLY_REPARAMETERIZED`, this
        `Distribution` can be reparameterized in terms of some standard
        distribution with a function whose Jacobian is constant for the support
        of the standard distribution. If `distributions.NOT_REPARAMETERIZED`,
        then no such reparameterization is available.
      validate_args: Python `bool`, default `False`. When `True` distribution
        parameters are checked for validity despite possibly degrading runtime
        performance. When `False` invalid inputs may silently render incorrect
        outputs.
      allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
        (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
        result is undefined. When `False`, an exception is raised if one or
        more of the statistic's batch members are undefined.
      parameters: Python `dict` of parameters used to instantiate this
        `Distribution`.
      graph_parents: Python `list` of graph prerequisites of this
        `Distribution`.
      name: Python `str` name prefixed to Ops created by this class. Default:
        subclass name.

    Raises:
      ValueError: if any member of graph_parents is `None` or not a `Tensor`.
    """
    graph_parents = [] if graph_parents is None else graph_parents
    for i, t in enumerate(graph_parents):
      if t is None or not tensor_util.is_tensor(t):
        raise ValueError("Graph parent item %d is not a Tensor; %s." % (i, t))
    if not name or name[-1] != "/":  # `name` is not a name scope
      non_unique_name = name or type(self).__name__
      with ops.name_scope(non_unique_name) as name:
        pass
    self._dtype = dtype
    self._reparameterization_type = reparameterization_type
    self._allow_nan_stats = allow_nan_stats
    self._validate_args = validate_args
    self._parameters = parameters or {}
    self._graph_parents = graph_parents
    self._name = name

  @property
  def _parameters(self):
    return self._parameter_dict

  @_parameters.setter
  def _parameters(self, value):
    """Intercept assignments to self._parameters to avoid reference cycles.

    Parameters are often created using locals(), so we need to clean out any
    references to `self` before assigning it to an attribute.

    Args:
      value: A dictionary of parameters to assign to the `_parameters` property.
    """
    if "self" in value:
      del value["self"]
    self._parameter_dict = value

  @classmethod
  def param_shapes(cls, sample_shape, name="DistributionParamShapes"):
    """Shapes of parameters given the desired shape of a call to `sample()`.

    This is a class method that describes what key/value arguments are required
    to instantiate the given `Distribution` so that a particular shape is
    returned for that instance's call to `sample()`.

    Subclasses should override class method `_param_shapes`.

    Args:
      sample_shape: `Tensor` or python list/tuple. Desired shape of a call to
        `sample()`.
      name: name to prepend ops with.

    Returns:
      `dict` of parameter name to `Tensor` shapes.
    """
    with ops.name_scope(name, values=[sample_shape]):
      return cls._param_shapes(sample_shape)

  @classmethod
  def param_static_shapes(cls, sample_shape):
    """param_shapes with static (i.e. `TensorShape`) shapes.

    This is a class method that describes what key/value arguments are required
    to instantiate the given `Distribution` so that a particular shape is
    returned for that instance's call to `sample()`. Assumes that the sample's
    shape is known statically.

    Subclasses should override class method `_param_shapes` to return
    constant-valued tensors when constant values are fed.

    Args:
      sample_shape: `TensorShape` or python list/tuple. Desired shape of a call
        to `sample()`.

    Returns:
      `dict` of parameter name to `TensorShape`.

    Raises:
      ValueError: if `sample_shape` is a `TensorShape` and is not fully defined.
    """
    if isinstance(sample_shape, tensor_shape.TensorShape):
      if not sample_shape.is_fully_defined():
        raise ValueError("TensorShape sample_shape must be fully defined")
      sample_shape = sample_shape.as_list()

    params = cls.param_shapes(sample_shape)

    static_params = {}
    for name, shape in params.items():
      static_shape = tensor_util.constant_value(shape)
      if static_shape is None:
        raise ValueError(
            "sample_shape must be a fully-defined TensorShape or list/tuple")
      static_params[name] = tensor_shape.TensorShape(static_shape)

    return static_params

  @staticmethod
  def _param_shapes(sample_shape):
    raise NotImplementedError("_param_shapes not implemented")

  @property
  def name(self):
    """Name prepended to all ops created by this `Distribution`."""
    return self._name

  @property
  def dtype(self):
    """The `DType` of `Tensor`s handled by this `Distribution`."""
    return self._dtype

  @property
  def parameters(self):
    """Dictionary of parameters used to instantiate this `Distribution`."""
    # Remove "self", "__class__", or other special variables. These can appear
    # if the subclass used:
    # `parameters = dict(locals())`.
    return {k: v for k, v in self._parameters.items()
            if not k.startswith("__") and k != "self"}

  @property
  def reparameterization_type(self):
    """Describes how samples from the distribution are reparameterized.

    Currently this is one of the static instances
    `distributions.FULLY_REPARAMETERIZED`
    or `distributions.NOT_REPARAMETERIZED`.

    Returns:
      An instance of `ReparameterizationType`.
    """
    return self._reparameterization_type

  @property
  def allow_nan_stats(self):
    """Python `bool` describing behavior when a stat is undefined.

    Stats return +/- infinity when it makes sense. E.g., the variance of a
    Cauchy distribution is infinity. However, sometimes the statistic is
    undefined, e.g., if a distribution's pdf does not achieve a maximum within
    the support of the distribution, the mode is undefined. If the mean is
    undefined, then by definition the variance is undefined. E.g. the mean for
    Student's T for df = 1 is undefined (no clear way to say it is either + or -
    infinity), so the variance = E[(X - mean)**2] is also undefined.

    Returns:
      allow_nan_stats: Python `bool`.
    """
    return self._allow_nan_stats

  @property
  def validate_args(self):
    """Python `bool` indicating possibly expensive checks are enabled."""
    return self._validate_args

  def copy(self, **override_parameters_kwargs):
    """Creates a deep copy of the distribution.

    Note: the copy distribution may continue to depend on the original
    initialization arguments.

    Args:
      **override_parameters_kwargs: String/value dictionary of initialization
        arguments to override with new values.

    Returns:
      distribution: A new instance of `type(self)` initialized from the union
        of self.parameters and override_parameters_kwargs, i.e.,
        `dict(self.parameters, **override_parameters_kwargs)`.
    """
    parameters = dict(self.parameters, **override_parameters_kwargs)
    return type(self)(**parameters)

  def _batch_shape_tensor(self):
    raise NotImplementedError(
        "batch_shape_tensor is not implemented: {}".format(type(self).__name__))

  def batch_shape_tensor(self, name="batch_shape_tensor"):
    """Shape of a single sample from a single event index as a 1-D `Tensor`.

    The batch dimensions are indexes into independent, non-identical
    parameterizations of this distribution.

    Args:
      name: name to give to the op

    Returns:
      batch_shape: `Tensor`.
    """
    with self._name_scope(name):
      if self.batch_shape.is_fully_defined():
        return ops.convert_to_tensor(self.batch_shape.as_list(),
                                     dtype=dtypes.int32,
                                     name="batch_shape")
      return self._batch_shape_tensor()

  def _batch_shape(self):
    return tensor_shape.TensorShape(None)

  @property
  def batch_shape(self):
    """Shape of a single sample from a single event index as a `TensorShape`.

    May be partially defined or unknown.

    The batch dimensions are indexes into independent, non-identical
    parameterizations of this distribution.

    Returns:
      batch_shape: `TensorShape`, possibly unknown.
    """
    return tensor_shape.as_shape(self._batch_shape())

  def _event_shape_tensor(self):
    raise NotImplementedError(
        "event_shape_tensor is not implemented: {}".format(type(self).__name__))

  def event_shape_tensor(self, name="event_shape_tensor"):
    """Shape of a single sample from a single batch as a 1-D int32 `Tensor`.

    Args:
      name: name to give to the op

    Returns:
      event_shape: `Tensor`.
    """
    with self._name_scope(name):
      if self.event_shape.is_fully_defined():
        return ops.convert_to_tensor(self.event_shape.as_list(),
                                     dtype=dtypes.int32,
                                     name="event_shape")
      return self._event_shape_tensor()

  def _event_shape(self):
    return tensor_shape.TensorShape(None)

  @property
  def event_shape(self):
    """Shape of a single sample from a single batch as a `TensorShape`.

    May be partially defined or unknown.

    Returns:
      event_shape: `TensorShape`, possibly unknown.
    """
    return tensor_shape.as_shape(self._event_shape())

  def is_scalar_event(self, name="is_scalar_event"):
    """Indicates that `event_shape == []`.

    Args:
      name: Python `str` prepended to names of ops created by this function.

    Returns:
      is_scalar_event: `bool` scalar `Tensor`.
    """
    with self._name_scope(name):
      return ops.convert_to_tensor(
          self._is_scalar_helper(self.event_shape, self.event_shape_tensor),
          name="is_scalar_event")

  def is_scalar_batch(self, name="is_scalar_batch"):
    """Indicates that `batch_shape == []`.

    Args:
      name: Python `str` prepended to names of ops created by this function.

    Returns:
      is_scalar_batch: `bool` scalar `Tensor`.
    """
    with self._name_scope(name):
      return ops.convert_to_tensor(
          self._is_scalar_helper(self.batch_shape, self.batch_shape_tensor),
          name="is_scalar_batch")

  def _sample_n(self, n, seed=None):
    raise NotImplementedError("sample_n is not implemented: {}".format(
        type(self).__name__))

  def _call_sample_n(self, sample_shape, seed, name, **kwargs):
    with self._name_scope(name, values=[sample_shape]):
      sample_shape = ops.convert_to_tensor(
          sample_shape, dtype=dtypes.int32, name="sample_shape")
      sample_shape, n = self._expand_sample_shape_to_vector(
          sample_shape, "sample_shape")
      samples = self._sample_n(n, seed, **kwargs)
      batch_event_shape = array_ops.shape(samples)[1:]
      final_shape = array_ops.concat([sample_shape, batch_event_shape], 0)
      samples = array_ops.reshape(samples, final_shape)
      samples = self._set_sample_static_shape(samples, sample_shape)
      return samples

  def sample(self, sample_shape=(), seed=None, name="sample"):
    """Generate samples of the specified shape.

    Note that a call to `sample()` without arguments will generate a single
    sample.

    Args:
      sample_shape: 0D or 1D `int32` `Tensor`. Shape of the generated samples.
      seed: Python integer seed for RNG
      name: name to give to the op.

    Returns:
      samples: a `Tensor` with prepended dimensions `sample_shape`.
    """
    return self._call_sample_n(sample_shape, seed, name)

  def _log_prob(self, value):
    raise NotImplementedError("log_prob is not implemented: {}".format(
        type(self).__name__))

  def _call_log_prob(self, value, name, **kwargs):
    with self._name_scope(name, values=[value]):
      value = ops.convert_to_tensor(value, name="value")
      try:
        return self._log_prob(value, **kwargs)
      except NotImplementedError as original_exception:
        try:
          return math_ops.log(self._prob(value, **kwargs))
        except NotImplementedError:
          raise original_exception

  def log_prob(self, value, name="log_prob"):
    """Log probability density/mass function.

    Args:
      value: `float` or `double` `Tensor`.
      name: Python `str` prepended to names of ops created by this function.

    Returns:
      log_prob: a `Tensor` of shape `sample_shape(x) + self.batch_shape` with
        values of type `self.dtype`.
    """
    return self._call_log_prob(value, name)

  def _prob(self, value):
    raise NotImplementedError("prob is not implemented: {}".format(
        type(self).__name__))

  def _call_prob(self, value, name, **kwargs):
    with self._name_scope(name, values=[value]):
      value = ops.convert_to_tensor(value, name="value")
      try:
        return self._prob(value, **kwargs)
      except NotImplementedError as original_exception:
        try:
          return math_ops.exp(self._log_prob(value, **kwargs))
        except NotImplementedError:
          raise original_exception

  def prob(self, value, name="prob"):
    """Probability density/mass function.

    Args:
      value: `float` or `double` `Tensor`.
      name: Python `str` prepended to names of ops created by this function.

    Returns:
      prob: a `Tensor` of shape `sample_shape(x) + self.batch_shape` with
        values of type `self.dtype`.
    """
    return self._call_prob(value, name)

  def _log_cdf(self, value):
    raise NotImplementedError("log_cdf is not implemented: {}".format(
        type(self).__name__))

  def _call_log_cdf(self, value, name, **kwargs):
    with self._name_scope(name, values=[value]):
      value = ops.convert_to_tensor(value, name="value")
      try:
        return self._log_cdf(value, **kwargs)
      except NotImplementedError as original_exception:
        try:
          return math_ops.log(self._cdf(value, **kwargs))
        except NotImplementedError:
          raise original_exception

  def log_cdf(self, value, name="log_cdf"):
    """Log cumulative distribution function.

    Given random variable `X`, the cumulative distribution function `cdf` is:

    ```none
    log_cdf(x) := Log[ P[X <= x] ]
    ```

    Often, a numerical approximation can be used for `log_cdf(x)` that yields
    a more accurate answer than simply taking the logarithm of the `cdf` when
    `x << -1`.

    Args:
      value: `float` or `double` `Tensor`.
      name: Python `str` prepended to names of ops created by this function.

    Returns:
      logcdf: a `Tensor` of shape `sample_shape(x) + self.batch_shape` with
        values of type `self.dtype`.
    """
    return self._call_log_cdf(value, name)

  def _cdf(self, value):
    raise NotImplementedError("cdf is not implemented: {}".format(
        type(self).__name__))

  def _call_cdf(self, value, name, **kwargs):
    with self._name_scope(name, values=[value]):
      value = ops.convert_to_tensor(value, name="value")
      try:
        return self._cdf(value, **kwargs)
      except NotImplementedError as original_exception:
        try:
          return math_ops.exp(self._log_cdf(value, **kwargs))
        except NotImplementedError:
          raise original_exception

  def cdf(self, value, name="cdf"):
    """Cumulative distribution function.

    Given random variable `X`, the cumulative distribution function `cdf` is:

    ```none
    cdf(x) := P[X <= x]
    ```

    Args:
      value: `float` or `double` `Tensor`.
      name: Python `str` prepended to names of ops created by this function.

    Returns:
      cdf: a `Tensor` of shape `sample_shape(x) + self.batch_shape` with
        values of type `self.dtype`.
    """
    return self._call_cdf(value, name)

  def _log_survival_function(self, value):
    raise NotImplementedError(
        "log_survival_function is not implemented: {}".format(
            type(self).__name__))

  def _call_log_survival_function(self, value, name, **kwargs):
    with self._name_scope(name, values=[value]):
      value = ops.convert_to_tensor(value, name="value")
      try:
        return self._log_survival_function(value, **kwargs)
      except NotImplementedError as original_exception:
        try:
          return math_ops.log1p(-self.cdf(value, **kwargs))
        except NotImplementedError:
          raise original_exception

  def log_survival_function(self, value, name="log_survival_function"):
    """Log survival function.

    Given random variable `X`, the survival function is defined:

    ```none
    log_survival_function(x) = Log[ P[X > x] ]
                             = Log[ 1 - P[X <= x] ]
                             = Log[ 1 - cdf(x) ]
    ```

    Typically, different numerical approximations can be used for the log
    survival function, which are more accurate than `1 - cdf(x)` when `x >> 1`.

    Args:
      value: `float` or `double` `Tensor`.
      name: Python `str` prepended to names of ops created by this function.

    Returns:
      `Tensor` of shape `sample_shape(x) + self.batch_shape` with values of type
        `self.dtype`.
    """
    return self._call_log_survival_function(value, name)

  def _survival_function(self, value):
    raise NotImplementedError("survival_function is not implemented: {}".format(
        type(self).__name__))

  def _call_survival_function(self, value, name, **kwargs):
    with self._name_scope(name, values=[value]):
      value = ops.convert_to_tensor(value, name="value")
      try:
        return self._survival_function(value, **kwargs)
      except NotImplementedError as original_exception:
        try:
          return 1. - self.cdf(value, **kwargs)
        except NotImplementedError:
          raise original_exception

  def survival_function(self, value, name="survival_function"):
    """Survival function.

    Given random variable `X`, the survival function is defined:

    ```none
    survival_function(x) = P[X > x]
                         = 1 - P[X <= x]
                         = 1 - cdf(x).
    ```

    Args:
      value: `float` or `double` `Tensor`.
      name: Python `str` prepended to names of ops created by this function.

    Returns:
      `Tensor` of shape `sample_shape(x) + self.batch_shape` with values of type
        `self.dtype`.
    """
    return self._call_survival_function(value, name)

  def _entropy(self):
    raise NotImplementedError("entropy is not implemented: {}".format(
        type(self).__name__))

  def entropy(self, name="entropy"):
    """Shannon entropy in nats."""
    with self._name_scope(name):
      return self._entropy()

  def _mean(self):
    raise NotImplementedError("mean is not implemented: {}".format(
        type(self).__name__))

  def mean(self, name="mean"):
    """Mean."""
    with self._name_scope(name):
      return self._mean()

  def _quantile(self, value):
    raise NotImplementedError("quantile is not implemented: {}".format(
        type(self).__name__))

  def _call_quantile(self, value, name, **kwargs):
    with self._name_scope(name, values=[value]):
      value = ops.convert_to_tensor(value, name="value")
      return self._quantile(value, **kwargs)

  def quantile(self, value, name="quantile"):
    """Quantile function. Aka "inverse cdf" or "percent point function".

    Given random variable `X` and `p in [0, 1]`, the `quantile` is:

    ```none
    quantile(p) := x such that P[X <= x] == p
    ```

    Args:
      value: `float` or `double` `Tensor`.
      name: Python `str` prepended to names of ops created by this function.

    Returns:
      quantile: a `Tensor` of shape `sample_shape(x) + self.batch_shape` with
        values of type `self.dtype`.
    """
    return self._call_quantile(value, name)

  def _variance(self):
    raise NotImplementedError("variance is not implemented: {}".format(
        type(self).__name__))

  def variance(self, name="variance"):
    """Variance.

    Variance is defined as,

    ```none
    Var = E[(X - E[X])**2]
    ```

    where `X` is the random variable associated with this distribution, `E`
    denotes expectation, and `Var.shape = batch_shape + event_shape`.

    Args:
      name: Python `str` prepended to names of ops created by this function.

    Returns:
      variance: Floating-point `Tensor` with shape identical to
        `batch_shape + event_shape`, i.e., the same shape as `self.mean()`.
    """
    with self._name_scope(name):
      try:
        return self._variance()
      except NotImplementedError as original_exception:
        try:
          return math_ops.square(self._stddev())
        except NotImplementedError:
          raise original_exception

  def _stddev(self):
    raise NotImplementedError("stddev is not implemented: {}".format(
        type(self).__name__))

  def stddev(self, name="stddev"):
    """Standard deviation.

    Standard deviation is defined as,

    ```none
    stddev = E[(X - E[X])**2]**0.5
    ```

    where `X` is the random variable associated with this distribution, `E`
    denotes expectation, and `stddev.shape = batch_shape + event_shape`.

    Args:
      name: Python `str` prepended to names of ops created by this function.

    Returns:
      stddev: Floating-point `Tensor` with shape identical to
        `batch_shape + event_shape`, i.e., the same shape as `self.mean()`.
    """

    with self._name_scope(name):
      try:
        return self._stddev()
      except NotImplementedError as original_exception:
        try:
          return math_ops.sqrt(self._variance())
        except NotImplementedError:
          raise original_exception

  def _covariance(self):
    raise NotImplementedError("covariance is not implemented: {}".format(
        type(self).__name__))

  def covariance(self, name="covariance"):
    """Covariance.

    Covariance is (possibly) defined only for non-scalar-event distributions.

    For example, for a length-`k`, vector-valued distribution, it is calculated
    as,

    ```none
    Cov[i, j] = Covariance(X_i, X_j) = E[(X_i - E[X_i]) (X_j - E[X_j])]
    ```

    where `Cov` is a (batch of) `k x k` matrix, `0 <= (i, j) < k`, and `E`
    denotes expectation.

    Alternatively, for non-vector, multivariate distributions (e.g.,
    matrix-valued, Wishart), `Covariance` shall return a (batch of) matrices
    under some vectorization of the events, i.e.,

    ```none
    Cov[i, j] = Covariance(Vec(X)_i, Vec(X)_j) = [as above]
    ```

    where `Cov` is a (batch of) `k' x k'` matrices,
    `0 <= (i, j) < k' = reduce_prod(event_shape)`, and `Vec` is some function
    mapping indices of this distribution's event dimensions to indices of a
    length-`k'` vector.

    Args:
      name: Python `str` prepended to names of ops created by this function.

    Returns:
      covariance: Floating-point `Tensor` with shape `[B1, ..., Bn, k', k']`
        where the first `n` dimensions are batch coordinates and
        `k' = reduce_prod(self.event_shape)`.
    """
    with self._name_scope(name):
      return self._covariance()

  def _mode(self):
    raise NotImplementedError("mode is not implemented: {}".format(
        type(self).__name__))

  def mode(self, name="mode"):
    """Mode."""
    with self._name_scope(name):
      return self._mode()

  def _cross_entropy(self, other):
    return kullback_leibler.cross_entropy(
        self, other, allow_nan_stats=self.allow_nan_stats)

  def cross_entropy(self, other, name="cross_entropy"):
    """Computes the (Shannon) cross entropy.

    Denote this distribution (`self`) by `P` and the `other` distribution by
    `Q`. Assuming `P, Q` are absolutely continuous with respect to
    one another and permit densities `p(x) dr(x)` and `q(x) dr(x)`, (Shanon)
    cross entropy is defined as:

    ```none
    H[P, Q] = E_p[-log q(X)] = -int_F p(x) log q(x) dr(x)
    ```

    where `F` denotes the support of the random variable `X ~ P`.

    Args:
      other: `tfp.distributions.Distribution` instance.
      name: Python `str` prepended to names of ops created by this function.

    Returns:
      cross_entropy: `self.dtype` `Tensor` with shape `[B1, ..., Bn]`
        representing `n` different calculations of (Shanon) cross entropy.
    """
    with self._name_scope(name):
      return self._cross_entropy(other)

  def _kl_divergence(self, other):
    return kullback_leibler.kl_divergence(
        self, other, allow_nan_stats=self.allow_nan_stats)

  def kl_divergence(self, other, name="kl_divergence"):
    """Computes the Kullback--Leibler divergence.

    Denote this distribution (`self`) by `p` and the `other` distribution by
    `q`. Assuming `p, q` are absolutely continuous with respect to reference
    measure `r`, the KL divergence is defined as:

    ```none
    KL[p, q] = E_p[log(p(X)/q(X))]
             = -int_F p(x) log q(x) dr(x) + int_F p(x) log p(x) dr(x)
             = H[p, q] - H[p]
    ```

    where `F` denotes the support of the random variable `X ~ p`, `H[., .]`
    denotes (Shanon) cross entropy, and `H[.]` denotes (Shanon) entropy.

    Args:
      other: `tfp.distributions.Distribution` instance.
      name: Python `str` prepended to names of ops created by this function.

    Returns:
      kl_divergence: `self.dtype` `Tensor` with shape `[B1, ..., Bn]`
        representing `n` different calculations of the Kullback-Leibler
        divergence.
    """
    with self._name_scope(name):
      return self._kl_divergence(other)

  def __str__(self):
    return ("tfp.distributions.{type_name}("
            "\"{self_name}\""
            "{maybe_batch_shape}"
            "{maybe_event_shape}"
            ", dtype={dtype})".format(
                type_name=type(self).__name__,
                self_name=self.name,
                maybe_batch_shape=(", batch_shape={}".format(self.batch_shape)
                                   if self.batch_shape.ndims is not None
                                   else ""),
                maybe_event_shape=(", event_shape={}".format(self.event_shape)
                                   if self.event_shape.ndims is not None
                                   else ""),
                dtype=self.dtype.name))

  def __repr__(self):
    return ("<tfp.distributions.{type_name} "
            "'{self_name}'"
            " batch_shape={batch_shape}"
            " event_shape={event_shape}"
            " dtype={dtype}>".format(
                type_name=type(self).__name__,
                self_name=self.name,
                batch_shape=self.batch_shape,
                event_shape=self.event_shape,
                dtype=self.dtype.name))

  @contextlib.contextmanager
  def _name_scope(self, name=None, values=None):
    """Helper function to standardize op scope."""
    with ops.name_scope(self.name):
      with ops.name_scope(name, values=(
          ([] if values is None else values) + self._graph_parents)) as scope:
        yield scope

  def _expand_sample_shape_to_vector(self, x, name):
    """Helper to `sample` which ensures input is 1D."""
    x_static_val = tensor_util.constant_value(x)
    if x_static_val is None:
      prod = math_ops.reduce_prod(x)
    else:
      prod = np.prod(x_static_val, dtype=x.dtype.as_numpy_dtype())

    ndims = x.get_shape().ndims  # != sample_ndims
    if ndims is None:
      # Maybe expand_dims.
      ndims = array_ops.rank(x)
      expanded_shape = util.pick_vector(
          math_ops.equal(ndims, 0),
          np.array([1], dtype=np.int32), array_ops.shape(x))
      x = array_ops.reshape(x, expanded_shape)
    elif ndims == 0:
      # Definitely expand_dims.
      if x_static_val is not None:
        x = ops.convert_to_tensor(
            np.array([x_static_val], dtype=x.dtype.as_numpy_dtype()),
            name=name)
      else:
        x = array_ops.reshape(x, [1])
    elif ndims != 1:
      raise ValueError("Input is neither scalar nor vector.")

    return x, prod

  def _set_sample_static_shape(self, x, sample_shape):
    """Helper to `sample`; sets static shape info."""
    # Set shape hints.
    sample_shape = tensor_shape.TensorShape(
        tensor_util.constant_value(sample_shape))

    ndims = x.get_shape().ndims
    sample_ndims = sample_shape.ndims
    batch_ndims = self.batch_shape.ndims
    event_ndims = self.event_shape.ndims

    # Infer rank(x).
    if (ndims is None and
        sample_ndims is not None and
        batch_ndims is not None and
        event_ndims is not None):
      ndims = sample_ndims + batch_ndims + event_ndims
      x.set_shape([None] * ndims)

    # Infer sample shape.
    if ndims is not None and sample_ndims is not None:
      shape = sample_shape.concatenate([None]*(ndims - sample_ndims))
      x.set_shape(x.get_shape().merge_with(shape))

    # Infer event shape.
    if ndims is not None and event_ndims is not None:
      shape = tensor_shape.TensorShape(
          [None]*(ndims - event_ndims)).concatenate(self.event_shape)
      x.set_shape(x.get_shape().merge_with(shape))

    # Infer batch shape.
    if batch_ndims is not None:
      if ndims is not None:
        if sample_ndims is None and event_ndims is not None:
          sample_ndims = ndims - batch_ndims - event_ndims
        elif event_ndims is None and sample_ndims is not None:
          event_ndims = ndims - batch_ndims - sample_ndims
      if sample_ndims is not None and event_ndims is not None:
        shape = tensor_shape.TensorShape([None]*sample_ndims).concatenate(
            self.batch_shape).concatenate([None]*event_ndims)
        x.set_shape(x.get_shape().merge_with(shape))

    return x

  def _is_scalar_helper(self, static_shape, dynamic_shape_fn):
    """Implementation for `is_scalar_batch` and `is_scalar_event`."""
    if static_shape.ndims is not None:
      return static_shape.ndims == 0
    shape = dynamic_shape_fn()
    if (shape.get_shape().ndims is not None and
        shape.get_shape()[0].value is not None):
      # If the static_shape is correctly written then we should never execute
      # this branch. We keep it just in case there's some unimagined corner
      # case.
      return shape.get_shape().as_list() == [0]
    return math_ops.equal(array_ops.shape(shape)[0], 0)