from matplotlib.offsetbox import AnchoredOffsetbox
#from matplotlib.transforms import IdentityTransform

import matplotlib.transforms as mtrans
from matplotlib.axes import Axes

from matplotlib.transforms import Bbox, TransformedBbox, IdentityTransform

from matplotlib.patches import Patch
from matplotlib.path import Path

from matplotlib.patches import Rectangle


class InsetPosition(object):
   def __init__(self, parent, lbwh):
       self.parent = parent
       self.lbwh = lbwh # position of the inset axes in the normalized coordinate of the parent axes

   def __call__(self, ax, renderer):
       bbox_parent = self.parent.get_position(original=False)
       trans = mtrans.BboxTransformTo(bbox_parent)
       bbox_inset = mtrans.Bbox.from_bounds(*self.lbwh)
       bb = mtrans.TransformedBbox(bbox_inset, trans)
       return bb


class AnchoredLocatorBase(AnchoredOffsetbox):
   def __init__(self, bbox_to_anchor, offsetbox, loc,
                borderpad=0.5, bbox_transform=None):

       super(AnchoredLocatorBase, self).__init__(loc,
                                                 pad=0., child=None,
                                                 borderpad=borderpad,
                                                 bbox_to_anchor=bbox_to_anchor,
                                                 bbox_transform=bbox_transform)


   def draw(self, renderer):
       raise RuntimeError("No draw method should be called")


   def __call__(self, ax, renderer):

       fontsize = renderer.points_to_pixels(self.prop.get_size_in_points())
       self._update_offset_func(renderer, fontsize)

       width, height, xdescent, ydescent = self.get_extent(renderer)

       px, py = self.get_offset(width, height, 0, 0)
       bbox_canvas = mtrans.Bbox.from_bounds(px, py, width, height)
       tr = ax.figure.transFigure.inverted()
       bb = mtrans.TransformedBbox(bbox_canvas, tr)

       return bb




from mpl_toolkits.axes_grid.axes_divider import Size

class AnchoredSizeLocator(AnchoredLocatorBase):
   def __init__(self, bbox_to_anchor, x_size, y_size, loc,
                borderpad=0.5, bbox_transform=None):

      self.axes = None
      self.x_size = Size.from_any(x_size)
      self.y_size = Size.from_any(y_size)

      super(AnchoredSizeLocator, self).__init__(bbox_to_anchor, None, loc,
                                                borderpad=borderpad,
                                                bbox_transform=bbox_transform)

   def get_extent(self, renderer):

      x, y, w, h = self.get_bbox_to_anchor().bounds

      dpi = renderer.points_to_pixels(72.)

      r, a = self.x_size.get_size(renderer)
      width = w*r + a*dpi

      r, a = self.y_size.get_size(renderer)
      height = h*r + a*dpi
      xd, yd = 0, 0

      fontsize = renderer.points_to_pixels(self.prop.get_size_in_points())
      pad = self.pad * fontsize

      return width+2*pad, height+2*pad, xd+pad, yd+pad


   def __call__(self, ax, renderer):

      self.axes = ax
      return super(AnchoredSizeLocator, self).__call__(ax, renderer)


class AnchoredZoomLocator(AnchoredLocatorBase):
   def __init__(self, parent_axes, zoom, loc,
                borderpad=0.5,
                bbox_to_anchor=None,
                bbox_transform=None):

      self.parent_axes = parent_axes
      self.zoom = zoom

      if bbox_to_anchor is None:
         bbox_to_anchor = parent_axes.bbox

      super(AnchoredZoomLocator, self).__init__(bbox_to_anchor, None, loc,
                                                borderpad=borderpad,
                                                bbox_transform=bbox_transform)

      self.axes = None


   def get_extent(self, renderer):

      bb = mtrans.TransformedBbox(self.axes.viewLim, self.parent_axes.transData)

      x, y, w, h = bb.bounds

      xd, yd = 0, 0

      fontsize = renderer.points_to_pixels(self.prop.get_size_in_points())
      pad = self.pad * fontsize

      return w*self.zoom+2*pad, h*self.zoom+2*pad, xd+pad, yd+pad


   def __call__(self, ax, renderer):

      self.axes = ax
      return super(AnchoredZoomLocator, self).__call__(ax, renderer)






class BboxPatch(Patch):
   def __init__(self, bbox, **kwargs):
        if "transform" in kwargs:
           raise ValueError("trnasform should nt be set")

        kwargs["transform"] = IdentityTransform()
        Patch.__init__(self, **kwargs)
        self.bbox = bbox

   def get_path(self):
       x0, y0, x1, y1 = self.bbox.extents

       verts = [(x0, y0),
                (x1, y0),
                (x1, y1),
                (x0, y1),
                (x0, y0),
                (0,0)]

       codes = [Path.MOVETO,
                Path.LINETO,
                Path.LINETO,
                Path.LINETO,
                Path.LINETO,
                Path.CLOSEPOLY]

       return Path(verts, codes)




class BboxConnector(Patch):

    @staticmethod
    def get_bbox_edge_pos(bbox, loc):
      x0, y0, x1, y1 = bbox.extents
      if loc==1:
         return x1, y1
      elif loc==2:
         return x0, y1
      elif loc==3:
         return x0, y0
      elif loc==4:
         return x1, y0

    @staticmethod
    def connect_bbox(bbox1, bbox2, loc1, loc2=None):
       if isinstance(bbox1, Rectangle):
          transform = bbox1.get_transfrom()
          bbox1 = Bbox.from_bounds(0, 0, 1, 1)
          bbox1 = TransformedBbox(bbox1, transform)

       if isinstance(bbox2, Rectangle):
          transform = bbox2.get_transform()
          bbox2 = Bbox.from_bounds(0, 0, 1, 1)
          bbox2 = TransformedBbox(bbox2, transform)

       if loc2 is None:
          loc2 = loc1

       x1, y1 = BboxConnector.get_bbox_edge_pos(bbox1, loc1)
       x2, y2 = BboxConnector.get_bbox_edge_pos(bbox2, loc2)

       verts = [[x1, y1], [x2,y2]]
       #Path()

       codes = [Path.MOVETO, Path.LINETO]

       return Path(verts, codes)


    def __init__(self, bbox1, bbox2, loc1, loc2=None, **kwargs):
        """
        *path* is a :class:`matplotlib.path.Path` object.

        Valid kwargs are:
        %(Patch)s

        .. seealso::

            :class:`Patch`
                For additional kwargs

        """
        if "transform" in kwargs:
           raise ValueError("trnasform should nt be set")

        kwargs["transform"] = IdentityTransform()
        Patch.__init__(self, **kwargs)
        self.bbox1 = bbox1
        self.bbox2 = bbox2
        self.loc1 = loc1
        self.loc2 = loc2


    def get_path(self):
       return self.connect_bbox(self.bbox1, self.bbox2,
                                self.loc1, self.loc2)


class BboxConnectorPatch(BboxConnector):

    def __init__(self, bbox1, bbox2, loc1a, loc2a, loc1b, loc2b, **kwargs):
        if "transform" in kwargs:
            raise ValueError("transform should not be set")
        BboxConnector.__init__(self, bbox1, bbox2, loc1a, loc2a, **kwargs)
        self.loc1b = loc1b
        self.loc2b = loc2b

    def get_path(self):
        path1 = self.connect_bbox(self.bbox1, self.bbox2, self.loc1, self.loc2)
        path2 = self.connect_bbox(self.bbox2, self.bbox1, self.loc2b, self.loc1b)
        path_merged = list(path1.vertices) + list (path2.vertices) + [path1.vertices[0]]
        return Path(path_merged)



def _add_inset_axes(parent_axes, inset_axes):
   parent_axes.figure.add_axes(inset_axes)
   inset_axes.set_navigate(False)


def inset_axes(parent_axes, width, height, loc=1,
               bbox_to_anchor=None, bbox_transform=None,
               axes_class=None,
               axes_kwargs=None,
               **kwargs):

   if axes_class is None:
      axes_class = Axes

   if axes_kwargs is None:
      inset_axes = axes_class(parent_axes.figure, parent_axes.get_position())
   else:
      inset_axes = axes_class(parent_axes.figure, parent_axes.get_position(),
                              **axes_kwargs)

   axes_locator = AnchoredSizeLocator(parent_axes.bbox,
                                      width, height,
                                      loc=loc)

   inset_axes.set_axes_locator(axes_locator)

   _add_inset_axes(parent_axes, inset_axes)

   return inset_axes


def zoomed_inset_axes(parent_axes, zoom, loc=1,
                      bbox_to_anchor=None, bbox_transform=None,
                      axes_class=None,
                      axes_kwargs=None,
                      **kwargs):

   if axes_class is None:
      axes_class = Axes

   if axes_kwargs is None:
      inset_axes = axes_class(parent_axes.figure, parent_axes.get_position())
   else:
      inset_axes = axes_class(parent_axes.figure, parent_axes.get_position(),
                              **axes_kwargs)

   axes_locator = AnchoredZoomLocator(parent_axes, zoom=zoom, loc=loc,
                                      bbox_to_anchor=None, bbox_transform=None,
                                      **kwargs)
   inset_axes.set_axes_locator(axes_locator)

   _add_inset_axes(parent_axes, inset_axes)

   return inset_axes


def mark_inset(parent_axes, inset_axes, loc1, loc2, **kwargs):
   rect = TransformedBbox(inset_axes.viewLim, parent_axes.transData)

   pp = BboxPatch(rect, **kwargs)
   parent_axes.add_patch(pp)

   p1 = BboxConnector(inset_axes.bbox, rect, loc1=loc1, **kwargs)
   inset_axes.add_patch(p1)
   p1.set_clip_on(False)
   p2 = BboxConnector(inset_axes.bbox, rect, loc1=loc2, **kwargs)
   inset_axes.add_patch(p2)
   p2.set_clip_on(False)

   return pp, p1, p2