%load_ext watermark

%load_ext lab_black

# all imports should go here

import pandas as pd
import numpy as np

import shapely.geometry as sg

import cartopy
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from cartopy.io.shapereader import Reader

import matplotlib.pyplot as plt
import matplotlib

from collections import namedtuple

import rasterio
from rasterio.plot import show
from rasterio.enums import Resampling

URL_E150_S30 = "C:\\Users\\donrc\\Downloads\\JRC_GFC2020_V2_S30_E150.tif"
URL_E140_S30 = "C:\\Users\\donrc\\Downloads\\JRC_GFC2020_V2_S30_E140.tif"
URL_W10_N60 = "C:\\Users\\donrc\\Downloads\\JRC_GFC2020_V2_N60_W10.tif"
URL_E0_N60 = (
    "C:\\Users\\donrc\\Downloads\\JRC_GFC2020_V2_N60_E0.tif"
)

URLS = [
    URL_E150_S30,
    URL_E140_S30,
    URL_W10_N60,
    URL_E0_N60,
]

def get_raster_data(dsr, scale_factor: float = 1.0 / 32):
    '''
    get_raster_data: read downscaled raster given DataSetReader

    Design:
    We usually expect scale_factor to be < 1, so as to reduce the
    size of the raster returned.

    Parameters:
    dsr: rasterio DataSetReader

    scale_factor: factor to inflate or deflate raster (default = 1/32)

    Returns:
    result of DataSetReader read() operation

    Limitations:
    No exception handling is done

    '''
    data2 = dsr.read(
        out_shape=(
            dsr.count,
            int(dsr.height * upscale_factor),
            int(dsr.width * upscale_factor),
        ),
        resampling=Resampling.nearest,
    )
    return data2


# end get_raster_data

#
# scaling factor to reduce raster size without affecting image quality (much)
#
upscale_factor = 1.0 / 32

Point = namedtuple('Point', ['lon', 'lat'])

Raster = namedtuple(
    'Raster', ['data', 'top_lat', 'left_lon']
)

# define the centers of our two maps
sydney = Point(lat=-33.86, lon=151.17)
london = Point(lat=51.52, lon=-0.09)

def get_forest_figure(
    centre, forest_rasters, ocean_url: str
):
    '''
    get_forest_figure: generate a matplotlib Figure, 1 degree square with forest cover

    Parameters:

    centre: namedtuple(Point, [lon, lat])

    forest_rasters: [namedtuple(Raster [data, top_lat, left_lon]) ]

    ocean_url: path to shapefile holding ocean data: set to None if not relevant

    Returns:
    matplotlib Figure object

    Side Effects:
    prints statement of rasters plotted

    Design:

    If a shapefile for ocean is available, use that else just draw cartopy coastlines (if any)

    We expect the raster to be binary (Forest at the pixel, Yes=1/No=0, so use
    custom colormap where 0 -> fully transparent

    We use PlateCarrree for the map projection

    We expect a list of rasters to draw, because the 10 degree grid of forest cover datasets doesnt
    perfectly align with the 1 degree square we are mapping (ie we need to plot two rasters
    to cover map extent

    Limitations:
    1 Degree square extent hard coded

    Knowledge of 10 degree extent of rasters hard-coded

    Green as forest color in colormap hard-coded

    '''
    my_cmap = matplotlib.colors.ListedColormap(
        ["none", "green"]
    )

    fig, ax = plt.subplots(
        figsize=(12, 12),
        subplot_kw={'projection': ccrs.PlateCarree()},
    )

    # define the extent of the plot
    # currently at 1 degree square
    ax.set_extent(
        [
            centre.lon - 0.5,
            centre.lon + 0.5,
            centre.lat - 0.5,
            centre.lat + 0.5,
        ],
        crs=ccrs.PlateCarree(),
    )
    if ocean_url != None:
        # draw blue ocean (if any)
        ax.add_geometries(
            Reader(ocean_url).geometries(),
            ccrs.PlateCarree(),
            facecolor=cfeature.COLORS['water'],
        )
    else:

        # draw coastlines
        ax.coastlines(
            resolution='10m',
            alpha=0.6,
            color='gray',
        )
    # end if

    # plot the array of forest cover)
    # we know that the rasters are on 10 degree squares
    for raster in forest_rasters:
        print(
            f'plotting raster {raster.left_lon=:5.1f}, {raster.top_lat=:5.1f}'
        )
        pcm = ax.imshow(
            raster.data,
            cmap=my_cmap,
            extent=[
                raster.left_lon,
                raster.left_lon + 10,
                raster.top_lat - 10,
                raster.top_lat,
            ],
        )

    # plot marker at centre of 1 degree square
    ax.plot(
        [centre.lon],
        [centre.lat],
        color='red',
        marker='o',
        markersize=10,
        zorder=9,
        linewidth=0,
        transform=ccrs.PlateCarree(),
    )

    # configure gridlines:  we dont want labels outside the black image
    # alpha=0.4 makes for dim lines
    gl = ax.gridlines(
        color='gray',
        alpha=0.4,
        draw_labels=False,
        x_inline=False,
        y_inline=False,
    )

    return fig


# end

ldn_data = get_raster_data(forest_dsrs[2])
ldn_forest1 = Raster(
    ldn_data[0].copy(),
    left_lon=forest_dsrs[2].bounds.left,
    top_lat=forest_dsrs[2].bounds.top,
)
ldn_data = get_raster_data(forest_dsrs[3])
ldn_forest2 = Raster(
    ldn_data[0].copy(),
    left_lon=forest_dsrs[3].bounds.left,
    top_lat=forest_dsrs[3].bounds.top,
)

syd_data = get_raster_data(forest_dsrs[0])
syd_forest1 = Raster(
    syd_data[0].copy(),
    left_lon=forest_dsrs[0].bounds.left,
    top_lat=forest_dsrs[0].bounds.top,
)