import dataclasses
import math
import os
import typing

os.environ["CUDA_VISIBLE_DEVICES"] = ""

import xarray as xr
import salem

import geopandas as gpd
import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d
import numpy as np
import pandas as pd
import pygmt as gmt
import quilt
import rasterio
import skimage

import comet_ml
import chainer
import cupy

from features.environment import _load_ipynb_modules, _download_model_weights_from_comet

data_prep = _load_ipynb_modules("data_prep.ipynb")

def get_image_with_bounds(filepaths: list, indexers: dict = None) -> xr.DataArray:
    """
    Retrieve raster image in xarray.DataArray format patched
    with projected coordinate bounds as (xmin, ymin, xmax, ymax)

    Note that if more than one filepath is passed in,
    the output groundtruth image array will not be valid
    (see https://github.com/pydata/xarray/issues/2159),
    but the window_bound extents will be correct
    """

    with xr.open_mfdataset(
        paths=filepaths, combine="nested", concat_dim=None
    ) as dataset:
        # Retrieve dataarray from NetCDF datasets
        dataarray = dataset.z.isel(indexers=indexers)

    # Patch projection information into xarray grid
    dataarray.attrs["pyproj_srs"] = "epsg:3031"
    sgrid = dataarray.salem.grid.corner_grid
    assert sgrid.origin == "lower-left"  # should be "lower-left", not "upper-left"

    # Patch bounding box extent into xarray grid
    if len(filepaths) == 1:
        left, right, bottom, top = sgrid.extent
    elif len(filepaths) > 1:
        print("WARN: using multiple inputs, output groundtruth image may look funny")
        x_offset, y_offset = sgrid.dx / 2, sgrid.dy / 2
        left, right = (
            float(dataarray.x[0] - x_offset),
            float(dataarray.x[-1] + x_offset),
        )
        assert sgrid.x0 == left
        bottom, top = (
            float(dataarray.y[0] - y_offset),
            float(dataarray.y[-1] + y_offset),
        )
        assert sgrid.y0 == bottom  # dataarray.y.min()-y_offset

    # check that y-axis and x-axis lengths are divisible by 4
    try:
        shape = int((top - bottom) / sgrid.dy), int((right - left) / sgrid.dx)
        assert all(i % 4 == 0 for i in shape)
    except AssertionError:
        print(f"WARN: Image shape {shape} should be divisible by 4 for DeepBedMap")
    finally:
        dataarray.attrs["bounds"] = [left, bottom, right, top]

    return dataarray

def get_deepbedmap_model_inputs(
    window_bound: rasterio.coords.BoundingBox,
    padding: int = 1000,
    use_whole_rema: bool = False,
) -> (np.ndarray, np.ndarray, np.ndarray, np.ndarray):
    """
    Outputs one large tile for each of:
    BEDMAP2, REMA, MEASURES Ice Flow Velocity and Antarctic Snow Accumulation
    according to a given window_bound in the form of (xmin, ymin, xmax, ymax).
    """
    data_prep = _load_ipynb_modules("data_prep.ipynb")

    if window_bound == rasterio.coords.BoundingBox(
        left=-1_594_000.0, bottom=-166_500.0, right=-1_575_000.0, top=-95_500.0
    ):
        # Quickly pull from cached quilt storage if using (hardcoded) test region
        quilt.install(package="weiji14/deepbedmap/model/test", force=True)
        pkg = quilt.load(pkginfo="weiji14/deepbedmap/model/test")
        X_tile = pkg.X_tile()
        W1_tile = pkg.W1_tile()
        W2_tile = pkg.W2_tile()
        W3_tile = pkg.W3_tile()
    else:
        X_tile = data_prep.selective_tile(
            filepath="lowres/bedmap2_bed.tif",
            window_bounds=[[*window_bound]],
            padding=padding,
            gapfiller=-5000.0,
        )
        W3_tile = data_prep.selective_tile(
            filepath="misc/Arthern_accumulation_bedmap2_grid1.tif",
            window_bounds=[[*window_bound]],
            padding=padding,
            gapfiller=0.0,
        )
        W2_tile = np.concatenate(
            [
                data_prep.selective_tile(
                    filepath="netcdf:misc/antarctic_ice_vel_phase_map_v01.nc:VX",
                    window_bounds=[[*window_bound]],
                    resolution=500,
                    padding=padding,
                ),
                data_prep.selective_tile(
                    filepath="netcdf:misc/antarctic_ice_vel_phase_map_v01.nc:VY",
                    window_bounds=[[*window_bound]],
                    resolution=500,
                    padding=padding,
                ),
            ],
            axis=1,
        )
        if not use_whole_rema:
            W1_tile = data_prep.selective_tile(
                filepath="misc/REMA_100m_dem_filled.tif",
                window_bounds=[[*window_bound]],
                padding=padding,
            )
        elif use_whole_rema:
            print("Getting: misc/REMA_100m_dem_filled.tif ...", end="")
            with xr.open_rasterio("misc/REMA_100m_dem_filled.tif") as ds:
                W1_tile = np.expand_dims(a=ds.values, axis=0)
                # special zero padding for REMA, 10pixels on top, bottom, left and right
                W1_tile = np.pad(
                    array=W1_tile,
                    pad_width=[(0, 0), (0, 0), (10, 10), (10, 10)],
                    mode="constant",
                )
            print("Done!")

    return X_tile, W1_tile, W2_tile, W3_tile

def subplot(directive: str, row: int = None, col: int = None, **kwargs):
    """Thin wrapper around https://docs.generic-mapping-tools.org/latest/subplot.html"""
    with gmt.clib.Session() as lib:
        rowcol = ""  # default is blank, e.g. when directive == "end"
        if row is not None and col is not None:
            if directive == "begin":
                rowcol = f"{row}x{col}"
            elif directive == "set":
                rowcol = f"{row},{col}"
        arg_str = " ".join(
            filter(None, [directive, rowcol, gmt.helpers.build_arg_string(kwargs)])
        )
        lib.call_module(module="subplot", args=arg_str)

def grdview(reliefgrid: str, **kwargs):
    """Thin wrapper around https://docs.generic-mapping-tools.org/latest/grdview.html"""
    kind = gmt.helpers.data_kind(reliefgrid, None, None)
    with gmt.clib.Session() as lib:
        if kind == "file":
            file_context = gmt.helpers.dummy_context(reliefgrid)
        elif kind == "grid":
            file_context = lib.virtualfile_from_grid(reliefgrid)
        with file_context as fname:
            arg_str = " ".join([fname, gmt.helpers.build_arg_string(kwargs)])
            lib.call_module(module="grdview", args=arg_str)

def plot_3d_view(
    img: str,  # filename (or xr.DataArray in the future) of the DEM to plot
    ax: tuple = None,  # optional (row, col) specifying gmt subplot position to use
    elev: int = 60,  # elevation angle above z-plane, value between 0-90
    azim: int = 202.5,  # azimuth from North, e.g. SouthEast is 135
    zmin: int = -1400,  # z-plane minimum elevation, in metres below or above datum
    cmap: str = "oleron",  # colormap to use
    title: str = None,  # title/label for the plot
    zlabel: str = None,  # z-axis label
):
    """
    Creates a 3D perspective view plot of an elevation surface using gmt grdview.
    The elevation (elev) and azimuth (azim) angle will need to be set accordingly,
    here it is looking from the SSW (202.5deg) at an angle 60deg above the z-plane.
    You can also provide a title (actually an annotation) and z-axis label. Note that
    there are several hardcoded defaults like the vertical exaggeration and axis labels.
    """
    if ax:
        row, col = ax  # split ax tuple into (row, col)
        subplot(directive="set", row=row, col=col, A=f'"{title}"')

    grdview(
        reliefgrid=img,
        B=f"SWZ",  # plot South, West axes, and Z-axis
        Bx='af+l"Polar Stereographic X (m)"',  # add x-axis annotations and minor ticks
        By='af+l"Polar Stereographic Y (m)"',  # add y-axis annotations and minor ticks
        Bz=f'af+l"{zlabel}"',  # add z-axis annotations, minor ticks and axis label
        C=cmap,  # colormap to use
        Jz=0.01,  # zscaling factor, default to 10x vertical exaggeration
        Q="sim",  # surface, image and mesh plot
        N=f"{zmin}",  # z-plane to plot on
        p=f"{azim}/{elev}",  # perspective using azimuth/elevation
    )
    return None  # should return instantiated PyGMT fig/Figure class

def load_trained_model(
    experiment_key: str = "055b697548e048b78202cfebb78d6d8c",  # or simply use "latest"
    model_weights_path: str = "model/weights/srgan_generator_model_weights.npz",
):
    """
    Returns a trained Generator DeepBedMap neural network model.

    The model's weights and hyperparameters settings are retrieved from
    https://comet.ml/weiji14/deepbedmap using an `experiment_key` setting
    which can be set to 'latest' or some 32-character alphanumeric string.
    """
    srgan_train = _load_ipynb_modules("srgan_train.ipynb")

    # Download either 'latest' model weights from Comet.ML or one using experiment_key
    # Will also get the hyperparameters "num_residual_blocks" and "residual_scaling"
    num_residual_blocks, residual_scaling = _download_model_weights_from_comet(
        experiment_key=experiment_key, download_path=model_weights_path
    )

    # Architect the model with appropriate "num_residual_blocks" and "residual_scaling"
    model = srgan_train.GeneratorModel(
        num_residual_blocks=num_residual_blocks, residual_scaling=residual_scaling
    )

    # Load trained neural network weights into model
    chainer.serializers.load_npz(file=model_weights_path, obj=model)

    return model

model = load_trained_model()

with chainer.using_config(name="enable_backprop", value=False):
    Y_hat = model.forward(x=X_tile, w1=W1_tile, w2=W2_tile, w3=W3_tile).array

# Save BEDMAP3 to GeoTiff and NetCDF format
deepbedmap3_grid = data_prep.save_array_to_grid(
    outfilepath="model/deepbedmap3",
    window_bound=window_bound,
    array=Y_hat[0, :, :, :],
    save_netcdf=True,
)
deepbedmap3_grid = xr.DataArray(
    data=np.flipud(cupy.asnumpy(Y_hat[0, 0, :, :])),
    dims=["y", "x"],
    coords={"y": deepbedmap3_grid.y, "x": deepbedmap3_grid.x},  # for multiple grids
    # coords={"y": groundtruth.y, "x": groundtruth.x},  # for single grid
)
deepbedmap3_grid = xr.open_dataarray("model/deepbedmap3.nc")

# Get elevation difference between DeepBedMap and CubicBedMap
elevdiffmap = deepbedmap3_grid - cubicbedmap_grid
elevdiffmap_grid = data_prep.save_array_to_grid(
    outfilepath="model/elevdiffmap",
    window_bound=window_bound,
    array=elevdiffmap,
    save_netcdf=True,
)

tracks = [data_prep.ascii_to_xyz(pipeline_file=f"{pf}.json") for pf in test_filepaths]
points: pd.DataFrame = pd.concat(objs=tracks)  # concatenate all tracks into one table

df_groundtruth = gmt.grdtrack(
    points=points, grid=groundtruth, newcolname="z_interpolated"
)
# df_deepbedmap3 = gmt.grdtrack(
#     points=points, grid=deepbedmap3_grid, newcolname="z_interpolated"
# )
df_deepbedmap3 = gmt.grdtrack(
    points=points, grid="model/deepbedmap3.nc", newcolname="z_interpolated"
)
df_cubicbedmap = gmt.grdtrack(
    points=points, grid="model/cubicbedmap.nc", newcolname="z_interpolated"
)
df_synthetichr = gmt.grdtrack(
    points=points, grid="model/synthetichr.nc", newcolname="z_interpolated"
)

try:
    X_tile = np.load(file="X_tile_big.npy")
    W1_tile = np.load(file="W1_tile_big.npy")
    W2_tile = np.load(file="W2_tile_big.npy")
    W3_tile = np.load(file="W3_tile_big.npy")
except FileNotFoundError:
    X_tile, W1_tile, W2_tile, W3_tile = get_deepbedmap_model_inputs(
        window_bound=window_bound_big, use_whole_rema=True
    )
    np.save(file="X_tile_big.npy", arr=X_tile)
    np.save(file="W1_tile_big.npy", arr=W1_tile)
    np.save(file="W2_tile_big.npy", arr=W2_tile)
    np.save(file="W3_tile_big.npy", arr=W3_tile)

# clip Ice Surface Elevation, Velocity and Accummulation values to above 0.0
W1_tile = np.clip(a=W1_tile, a_min=0.0, a_max=None)
W2_tile = np.clip(a=W2_tile, a_min=0.0, a_max=None)
W3_tile = np.clip(a=W3_tile, a_min=0.0, a_max=None)

@dataclasses.dataclass(frozen=True)
class Shape:
    y: int  # size in y-direction
    x: int  # size in x-direction

# The whole of Antarctica tile and predictor
with chainer.using_config(name="cudnn_deterministic", value=True):
    # Size are in kilometres, same as BEDMAP2's 1km resolution
    final_shape = Shape(y=18000, x=22000)  # 4x that of BEDMAP2
    ary_shape = Shape(y=1000, x=1000)  # 1000pixels * 250m = 250km
    stride = Shape(y=1000, x=1000)  # cut a tile every 1000 metres
    xtrapad = Shape(y=18, x=18)  # extra padding at borders (which will be clipped off)

    Y_hat = np.full(
        shape=(1, final_shape.y, final_shape.x), fill_value=np.nan, dtype=np.float32
    )

    steps = []
    for y_step in range(0, final_shape.y, stride.y):
        for x_step in range(0, final_shape.x, stride.x):
            steps.append(Shape(y=y_step, x=x_step))

    for step in steps:
        # plus `xtrapad.y` pixels and 1 pixel on bottom and top
        y0 = max(0, (step.y // 4) - xtrapad.y - 1)
        y1 = min(final_shape.y // 4, ((step.y + ary_shape.y) // 4) + xtrapad.y + 1)
        # plus `xtrapad.x` pixels and 1 pixel on left and right
        x0 = max(0, (step.x // 4) - xtrapad.x - 1)
        x1 = min(final_shape.x // 4, ((step.x + ary_shape.x) // 4) + xtrapad.x + 1)
        # print(str(x0).zfill(4), str(x1).zfill(4), str(y0).zfill(4), str(y1).zfill(4))

        # Hardcoded crops of BEDMAP2 (X), REMA (W1), MEaSUREs (W2), Accumulation (W3)
        X_tile_crop = model.xp.asarray(a=X_tile[:, :, y0:y1, x0:x1], dtype="float32")
        W1_tile_crop = model.xp.asarray(
            a=W1_tile[:, :, y0 * 10 : y1 * 10, x0 * 10 : x1 * 10], dtype="float32"
        )
        W2_tile_crop = model.xp.asarray(
            a=W2_tile[:, :, y0 * 2 : y1 * 2, x0 * 2 : x1 * 2], dtype="float32"
        )
        W3_tile_crop = model.xp.asarray(a=W3_tile[:, :, y0:y1, x0:x1], dtype="float32")

        # DeepBedMap terrain inference
        with chainer.using_config(name="enable_backprop", value=False):
            Y_pred = model.forward(
                x=X_tile_crop, w1=W1_tile_crop, w2=W2_tile_crop, w3=W3_tile_crop
            )

        try:
            y_slice = slice((y0 + xtrapad.y + 1) * 4, (y1 - xtrapad.y - 1) * 4)
            x_slice = slice((x0 + xtrapad.x + 1) * 4, (x1 - xtrapad.x - 1) * 4)
            Y_pred_uncut: np.ndarray = cupy.asnumpy(Y_pred.array)[0, :, :, :]
            Y_hat[:, y_slice, x_slice] = Y_pred_uncut[
                :, xtrapad.y * 4 : -xtrapad.y * 4, xtrapad.x * 4 : -xtrapad.x * 4
            ]
        except ValueError:
            raise
        finally:
            X_tile_crop = W1_tile_crop = W2_tile_crop = W3_tile_crop = None

# Save BEDMAP3 to GeoTiff and NetCDF format
# Using LZW compression and int16 instead of float32 to keep things smaller
_ = data_prep.save_array_to_grid(
    window_bound=window_bound_big,
    array=Y_hat.astype(dtype=np.int16),
    outfilepath="model/deepbedmap3_big_int16",
    dtype=np.int16,
    tiled=True,
    compression=rasterio.enums.Compression.lzw.value,  # Lempel-Ziv-Welch, lossless
)