import numpy as nm

from sfepy.base.base import assert_, Struct
import sfepy.linalg as la

class ContactPlane(Struct):

    def __init__(self, anchor, normal, bounds):
        Struct.__init__(self, anchor=nm.array(anchor, dtype=nm.float64),
                        bounds=nm.asarray(bounds, dtype=nm.float64))
        self.normal = nm.asarray(normal, dtype=nm.float64)

        norm = nm.linalg.norm
        self.normal /= norm(self.normal)

        e3 = [0.0, 0.0, 1.0]
        dd = nm.dot(e3, self.normal)
        rot_angle = nm.arccos(dd)

        if nm.abs(rot_angle) < 1e-14:
            mtx = nm.eye(3, dtype=nm.float64)
            bounds2d = self.bounds[:, :2]

        else:
            rot_axis = nm.cross([0.0, 0.0, 1.0], self.normal)
            mtx = la.make_axis_rotation_matrix(rot_axis, rot_angle)

            mm = la.insert_strided_axis(mtx, 0, self.bounds.shape[0])
            rbounds = la.dot_sequences(mm, self.bounds)
            bounds2d = rbounds[:, :2]

        assert_(nm.allclose(nm.dot(mtx, self.normal), e3,
                            rtol=0.0, atol=1e-12))

        self.adotn = nm.dot(self.anchor, self.normal)

        self.rot_angle = rot_angle
        self.mtx = mtx
        self.bounds2d = bounds2d

    def mask_points(self, points):
        mm = la.insert_strided_axis(self.mtx, 0, points.shape[0])
        points2d = la.dot_sequences(mm, points)[:, :2]

        return la.flag_points_in_polygon2d(self.bounds2d, points2d)

    def get_distance(self, points):
        dist = la.dot_sequences(points, self.normal) - self.adotn

        return dist

class ContactSphere(Struct):

    def __init__(self, centre, radius):
        self.centre = nm.asarray(centre)
        self.radius = radius

    def mask_points(self, points, eps):
        dist2 = la.norm_l2_along_axis(points - self.centre, squared=True)
        radius2 = self.radius**2
        mask = dist2 <= ((1 + eps)**2) * radius2
        return mask

    def get_distance(self, points):
        """
        Get the penetration distance and normals of points w.r.t. the sphere
        surface.

        Returns
        -------
        d : array
            The penetration distance.
        normals : array
            The normals from the points to the sphere centre.
        """
        vecs = self.centre - points
        dist = la.norm_l2_along_axis(vecs)
        # Prevent zero division.
        ii = dist > 1e-8
        normals = nm.where(ii[:, None], vecs[ii] / dist[ii][:, None],
                           vecs[ii])
        return self.radius - dist, normals

    def _get_derivatives(self, points):
        vecs = self.centre - points
        dist = la.norm_l2_along_axis(vecs)

        # Distance derivative w.r.t. point coordinates.
        dd = vecs / dist[:, None]

        normals = dd
        # Unit normal derivative w.r.t. point coordinates.
        dim = points.shape[1]
        ee = nm.eye(dim)[None, ...]
        nnt = normals[..., None] * normals[..., None, :]
        dn = - (ee - nnt) / dist[:, None, None]

        return dd, dn

def plot_polygon(ax, polygon):
    from sfepy.postprocess.plot_dofs import _get_axes

    dim = polygon.shape[1]
    ax = 
_get_axes(ax, dim)