from collections import OrderedDict
import warnings
import weakref

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.artist
import matplotlib.collections
import mpl_toolkits.axisartist.floating_axes as floating_axes
from mpl_toolkits.axisartist.floating_axes import FloatingSubplot

import cartopy.crs as ccrs
import cartopy.mpl.patch as cpatch
from cartopy.mpl.feature_artist import FeatureArtist, _freeze, _GeomKey
from cartopy.mpl.geoaxes import InterProjectionTransform

class AnAlternativeGeoAxes:
    def __init__(self, fig, rect, proj, region):
        """use ax_lonlat for (lon, lat) data
        use ax_proj for (x, y) data"""
        tr = InterProjectionTransform(ccrs.PlateCarree(),
                                             proj)
        grid_helper = floating_axes.GridHelperCurveLinear(
            tr,
            extremes=region,
            grid_locator1=None,
            grid_locator2=None,
            tick_formatter1=None,
            tick_formatter2=None)
        self.ax_proj = myFloatingSubplot(fig, proj, tr, rect, grid_helper=grid_helper)
        fig.add_subplot(self.ax_proj)
        
        self.ax_lonlat = self.ax_proj.get_aux_axes(tr)
        self.ax_lonlat.patch = self.ax_proj.patch  # for aux_ax to have a clip path as in ax
        self.ax_proj.patch.zorder = 0.9  # but this has a side effect that the patch is
        # drawn twice, and possibly over some other
        # artists. So, we decrease the zorder a bit to
        # prevent this.
        self.tr = tr
        self.target_crs = proj
        
    def scatter(self, *args, **kwargs):
        self.ax_lonlat.scatter(*args, **kwargs)
        
    def grid(self, *args, **kwargs):
        self.ax_proj.grid(*args, **kwargs)

    def set_xlabel(self, *args, **kwargs):
        self.ax_proj.set_xlabel(*args, **kwargs)

    def set_ylabel(self, *args, **kwargs):
        self.ax_proj.set_ylabel(*args, **kwargs)

    def set_title(self, *args, **kwargs):
        self.ax_proj.set_title(*args, **kwargs)

    def add_cartopy_feature(self, feature, zorder=1, **kwargs):
        artist = myFeatureArtist(feature, zorder=zorder, **kwargs)
        self.ax_proj.add_artist(artist)

class myFloatingSubplot(FloatingSubplot):
    def __init__(self, fig, proj, tr, *args, **kwargs):
        super().__init__(fig, *args, **kwargs)
        self.projection = proj
        self.tr = tr
    
    # I added 'get_extent' here to be called by myFeatureArtist
    def get_extent(self):
        xlim = self.get_xlim()
        ylim = self.get_ylim()
        extent = (xlim[0], xlim[1], ylim[0], ylim[1])
        return extent
    
class myFeatureArtist(FeatureArtist):
    def __init__(self, feature, **kwargs):
        super().__init__(feature, **kwargs)
    
    def draw(self, renderer, *args, **kwargs):
        """
        Draw the geometries of the feature that intersect with the extent of
        the :class:`cartopy.mpl.GeoAxes` instance to which this
        object has been added.

        """
        if not self.get_visible():
            return

        ax = self.axes
        feature_crs = self._feature.crs

        # Get geometries that we need to draw.
        extent = None
        try:
            extent = ax.get_extent(feature_crs)
        except ValueError:
            warnings.warn('Unable to determine extent. Defaulting to global.')
        except: # added this new get_extent() here
            extent = ax.get_extent()
        geoms = self._feature.intersecting_geometries(extent)

        # Combine all the keyword args in priority order.
        prepared_kwargs = dict(self._feature.kwargs)
        prepared_kwargs.update(self._kwargs)
        prepared_kwargs.update(kwargs)

        # Freeze the kwargs so that we can use them as a dict key. We will
        # need to unfreeze this with dict(frozen) before passing to mpl.
        prepared_kwargs = _freeze(prepared_kwargs)

        # Project (if necessary) and convert geometries to matplotlib paths.
        stylised_paths = OrderedDict()
        key = ax.projection
        # key = ccrs.LambertConformal(central_longitude=0)
        for geom in geoms:
            # As Shapely geometries cannot be relied upon to be
            # hashable, we have to use a WeakValueDictionary to manage
            # their weak references. The key can then be a simple,
            # "disposable", hashable geom-key object that just uses the
            # id() of a geometry to determine equality and hash value.
            # The only persistent, strong reference to the geom-key is
            # in the WeakValueDictionary, so when the geometry is
            # garbage collected so is the geom-key.
            # The geom-key is also used to access the WeakKeyDictionary
            # cache of transformed geometries. So when the geom-key is
            # garbage collected so are the transformed geometries.
            geom_key = _GeomKey(geom)
            FeatureArtist._geom_key_to_geometry_cache.setdefault(
                geom_key, geom)
            mapping = FeatureArtist._geom_key_to_path_cache.setdefault(
                geom_key, {})
            geom_paths = mapping.get(key)
            if geom_paths is None:
                if ax.projection != feature_crs:
                    projected_geom = ax.projection.project_geometry(
                        geom, feature_crs)
                else:
                    projected_geom = geom
                geom_paths = cpatch.geos_to_path(projected_geom)
                mapping[key] = geom_paths

            if not self._styler:
                style = prepared_kwargs
            else:
                # Unfreeze, then add the computed style, and then re-freeze.
                style = dict(prepared_kwargs)
                style.update(self._styler(geom))
                style = _freeze(style)

            stylised_paths.setdefault(style, []).extend(geom_paths)

        # transform = ax.projection._as_mpl_transform(ax)
        # transData here is a transform from projected coordinates to 
        # image coordinates (?).
        transform = ax.transData 

        # Draw one PathCollection per style. We could instead pass an array
        # of style items through to a single PathCollection, but that
        # complexity does not yet justify the effort.
        for style, paths in stylised_paths.items():
            # Build path collection and draw it.
            c = matplotlib.collections.PathCollection(
                    paths,
                    transform=transform,
                    **dict(style))
            c.set_clip_path(ax.patch)
            c.set_figure(ax.figure)
            c.draw(renderer)

        # n.b. matplotlib.collection.Collection.draw returns None
        return None



if __name__ == '__main__':
    import cartopy.feature as cfeature
    land_50m = cfeature.NaturalEarthFeature('physical', 'land', '50m',
                                            edgecolor='k', linewidths=0.5,
                                            facecolor=cfeature.COLORS['land'])
    countries = cfeature.NaturalEarthFeature('cultural', 'admin_0_boundary_lines_land',
        '50m', edgecolor='k', linewidths=0.25, facecolor='none')


    fig = plt.figure(figsize=(8, 4), dpi=150)  
    ax = AnAlternativeGeoAxes(fig, 111, ccrs.LambertConformal(central_longitude=0), region=(-60, 60, 30, 70))

    ax.add_cartopy_feature(land_50m)
    ax.add_cartopy_feature(countries)
    
    ax.set_xlabel('Longitude (\xb0)')
    ax.set_ylabel('Latitude (\xb0)')
    ax.set_title('Lambert Conformal')
    ax.grid()

    plt.savefig('lambert.png', dpi=300)