# Read in raw x, y, dhdt_slope and referencegroundtrack data into the GPU
cudf_raw: cudf.DataFrame = dask_cudf.read_parquet(
    path="ATLXI/df_dhdt_antarctica.parquet",
    columns=["x", "y", "dhdt_slope", "referencegroundtrack"],
    # filters=[[('dhdt_slope', '<', -0.105)], [('dhdt_slope', '>', 0.105)]],
)
# Filter to points with dhdt that is less than -0.105 m/yr or more than +0.105 m/yr
# Based on ICESat-2 ATL06's accuracy and precision of 3.3 ± 7.2cm from Brunt et al 2020
# See https://doi.org/10.1029/2020GL090572
cudf_many = cudf_raw.loc[abs(cudf_raw.dhdt_slope) > 0.105].compute()
print(f"Trimmed {len(cudf_raw)} -> {len(cudf_many)}")
if "cudf_raw" in globals():
    del cudf_raw

# Use point in polygon to label points according to the drainage basins they fall in
cudf_many["drainage_basin"]: cudf.Series = deepicedrain.point_in_polygon_gpu(
    points_df=cudf_many, poly_df=drainage_basins, poly_limit=16
)
X_many = cudf_many.dropna()  # drop points that are not in a drainage basin
print(f"Trimmed {len(cudf_many)} -> {len(X_many)}")

# Subglacial lake finder
activelakes: dict = {
    "basin_name": [],  # Antarctic drainage basin name
    "refgtracks": [],  # Pipe-delimited list of ICESat-2 reference ground tracks
    "num_points": [],  # Number of clustered data points
    "maxabsdhdt": [],  # Maximum absolute dhdt value inside of lake boundary
    "inner_dhdt": [],  # Median elev change over time (dhdt) inside of lake bounds
    "mean_dhdt": [],  # Mean elev change over time (dhdt) inside of lake bounds
    "outer_dhdt": [],  # Median elevation change over time (dhdt) outside of lake
    "outer_std": [],  # Standard deviation of dhdt outside of lake
    "outer_mad": [],  # Median absolute deviation of dhdt outside of lake
    "geometry": [],  # Shapely Polygon geometry holding lake boundary coordinates
}
basin_name: str = "Cook"  # Set a basin name here
basins = drainage_basins[drainage_basins.NAME == basin_name].index  # one specific basin
# basins = drainage_basins[
#     drainage_basins.NAME.isin(("Cook", "Whillans"))
# ].index  # some specific basins
basins: pd.core.indexes.numeric.Int64Index = drainage_basins.index  # run on all basins

eps: int = 3000  # ICESat-2 tracks are separated by ~3 km across track, with each laser pair ~90 m apart
min_samples: int = 300
for basin_index in tqdm.tqdm(iterable=basins):
    # Initial data cleaning, filter to rows that are in the drainage basin
    basin = drainage_basins.loc[basin_index]
    X_local = X_many.loc[X_many.drainage_basin == basin.NAME]  # .reset_index(drop=True)

    # Get points with dhdt_slope higher than 3x the median dhdt_slope for the basin
    # E.g. if median dhdt_slope is 0.30 m/yr, then we cluster points over 0.90 m/yr
    abs_dhdt = X_local.dhdt_slope.abs()
    tolerance: float = 3 * abs_dhdt.median()
    X = X_local.loc[abs_dhdt > tolerance]

    if len(X) <= 1000:  # don't run on too few points
        continue

    # Run unsupervised clustering separately on draining and filling lakes
    # Draining lake points have negative labels (e.g. -1, -2, 3),
    # Filling lake points have positive labels (e.g. 1, 2, 3),
    # Noise points have NaN labels (i.e. NaN)
    cluster_vars = ["x", "y", "dhdt_slope"]
    draining_lake_labels = -deepicedrain.find_clusters(
        X=X.loc[X.dhdt_slope < 0][cluster_vars],
        eps=eps,
        min_samples=min_samples,
        verbose=cuml.common.logger.level_error,
    )
    filling_lake_labels = deepicedrain.find_clusters(
        X=X.loc[X.dhdt_slope > 0][cluster_vars],
        eps=eps,
        min_samples=min_samples,
        verbose=cuml.common.logger.level_error,
    )
    lake_labels = cudf.concat(objs=[draining_lake_labels, filling_lake_labels])
    lake_labels: cudf.Series = lake_labels.sort_index()
    assert lake_labels.name == "cluster_id"

    # Checking all potential subglacial lakes in a basin
    clusters: cudf.Series = lake_labels.unique()
    for cluster_label in clusters.to_array():
        # Store attribute and geometry information of each active lake
        lake_points: cudf.DataFrame = X.loc[lake_labels == cluster_label]

        # More data cleaning, dropping clusters with too few points
        try:
            assert len(lake_points) > 100
        except AssertionError:
            lake_labels = lake_labels.replace(to_replace=cluster_label, value=None)
            continue

        multipoint: shapely.geometry.MultiPoint = shapely.geometry.MultiPoint(
            points=lake_points[["x", "y"]].as_matrix()
        )
        convexhull: shapely.geometry.Polygon = multipoint.convex_hull

        # Filter out (most) false positive subglacial lakes
        # Check that elevation change over time in lake is anomalous to outside
        # The 5000 m distance from lake boundary setting is empirically based on
        # Smith et al. 2009's methodology at https://doi.org/10.3189/002214309789470879
        outer_ring_buffer = convexhull.buffer(distance=5000) - convexhull
        X_local["in_donut_ring"] = deepicedrain.point_in_polygon_gpu(
            points_df=X_local,
            poly_df=gpd.GeoDataFrame({"name": True, "geometry": [outer_ring_buffer]}),
        )
        outer_points = X_local.dropna(subset="in_donut_ring")
        outer_dhdt: float = outer_points.dhdt_slope.median()

        outer_std: float = outer_points.dhdt_slope.std()
        outer_mad: float = scipy.stats.median_abs_deviation(
            x=outer_points.dhdt_slope.to_pandas()
        )

        mean_dhdt: float = lake_points.dhdt_slope.mean()
        inner_dhdt: float = lake_points.dhdt_slope.median()
        X_local = X_local.drop(labels="in_donut_ring", axis="columns")

        # If lake interior's median dhdt value is within 3 median absolute deviations
        # of the lake exterior's dhdt value, we remove the lake label
        # I.e. skip if above background change not significant enough
        # Inspired by Kim et al. 2016's methodology at https://doi.org/10.5194/tc-10-2971-2016
        if abs(inner_dhdt - outer_dhdt) < 3 * outer_mad:
            lake_labels = lake_labels.replace(to_replace=cluster_label, value=None)
            continue

        maxabsdhdt: float = (
            lake_points.dhdt_slope.max()
            if cluster_label > 0  # positive label = filling
            else lake_points.dhdt_slope.min()  # negative label = draining
        )
        refgtracks: str = "|".join(
            map(str, lake_points.referencegroundtrack.unique().to_pandas())
        )

        # Save key variables to dictionary that will later go into geodataframe
        activelakes["basin_name"].append(basin.NAME)
        activelakes["refgtracks"].append(refgtracks)
        activelakes["num_points"].append(len(lake_points))
        activelakes["maxabsdhdt"].append(maxabsdhdt)
        activelakes["inner_dhdt"].append(inner_dhdt)
        activelakes["mean_dhdt"].append(mean_dhdt)
        activelakes["outer_dhdt"].append(outer_dhdt)
        activelakes["outer_std"].append(outer_std)
        activelakes["outer_mad"].append(outer_mad)
        activelakes["geometry"].append(convexhull)

    # Calculate total number of lakes found for one drainage basin
    clusters: cudf.Series = lake_labels.unique()
    n_draining, n_filling = (clusters < 0).sum(), (clusters > 0).sum()
    if n_draining + n_filling > 0:
        print(f"{len(X)} rows at {basin.NAME} above ± {tolerance:.2f} m/yr")
        print(f"{n_draining} draining and {n_filling} filling lakes found")

if len(activelakes["geometry"]) >= 1:
    gdf = gpd.GeoDataFrame(activelakes, crs="EPSG:3031")
    basename = "antarctic_subglacial_lakes"  # f"temp_{basin_name.lower()}_lakes"  #
    gdf.to_file(filename=f"{basename}_3031.geojson", driver="GeoJSON")
    gdf.to_crs(crs={"init": "epsg:4326"}).to_file(
        filename=f"{basename}_4326.geojson", driver="GeoJSON"
    )

print(f"Total of {len(gdf)} subglacial lakes found")

# Choose one Antarctic active subglacial lake polygon with EPSG:3031 coordinates
lake_name: str = "Whillans IX"
lake_catalog = deepicedrain.catalog.subglacial_lakes()
lake_ids, transect_id = (
    pd.json_normalize(lake_catalog.metadata["lakedict"])
    .query("lakename == @lake_name")[["ids", "transect"]]
    .iloc[0]
)
lake = (
    lake_catalog.read()
    .loc[lake_ids]
    .dissolve(by=np.zeros(shape=len(lake_ids), dtype="int64"), as_index=False)
    .squeeze()
)

region = deepicedrain.Region.from_gdf(gdf=lake, name=lake_name)
draining: bool = lake.inner_dhdt < 0

print(lake)
lake.geometry

# Subset data to lake of interest
placename: str = region.name.lower().replace(" ", "_")
df_lake: cudf.DataFrame = region.subset(data=df_dhdt)
# Get all raw xyz points and one transect line dataframe
track_dict: dict = deepicedrain.split_tracks(df=df_lake.to_pandas())
track_points: pd.DataFrame = (
    pd.concat(track_dict.values())
    .groupby(by=["x", "y"])
    .mean()  # z value is mean h_corr over all cycles
    .reset_index()[["x", "y", "h_corr"]]
)
try:
    _rgt, _pt = transect_id.split("_")
    df_transect: pd.DataFrame = (
        track_dict[transect_id][["x", "y", "h_corr", "cycle_number"]]
        .groupby(by=["x", "y"])
        .max()  # z value is maximum h_corr over all cycles
        .reset_index()
    )
except AttributeError:
    pass

# Save lake outline to OGR GMT file format
outline_points: str = f"figures/{placename}/{placename}.gmt"
if not os.path.exists(path=outline_points):
    os.makedirs(name=f"figures/{placename}", exist_ok=True)
    lake_catalog.read().loc[list(lake_ids)].to_file(
        filename=outline_points, driver="OGR_GMT"
    )

# Generate gridded time-series of ice elevation over lake
cycles: tuple = (3, 4, 5, 6, 7, 8, 9)
os.makedirs(name=f"figures/{placename}", exist_ok=True)
ds_lake: xr.Dataset = deepicedrain.spatiotemporal_cube(
    table=df_lake.to_pandas(),
    placename=placename,
    cycles=cycles,
    folder=f"figures/{placename}",
)
ds_lake.to_netcdf(path=f"figures/{placename}/xyht_{placename}.nc", mode="w")

# Get 3D grid_region (xmin/xmax/ymin/ymax/zmin/zmax),
# and calculate normalized z-values as Elevation delta relative to Cycle 3
z_limits: tuple = (float(ds_lake.z.min()), float(ds_lake.z.max()))  # original z limits
grid_region: tuple = region.bounds() + z_limits

ds_lake_diff: xr.Dataset = ds_lake - ds_lake.sel(cycle_number=3).z
z_diff_limits: tuple = (float(ds_lake_diff.z.min()), float(ds_lake_diff.z.max()))
diff_grid_region: np.ndarray = np.append(arr=grid_region[:4], values=z_diff_limits)

print(f"Elevation limits are: {z_limits}")