import os
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import numpy.ma as ma
import pandas as pd
import rioxarray as rxr
from rasterio.plot import plotting_extent
import geopandas as gpd
# Rasterstats contains the zonalstatistics function
# that you will use to extract raster values
import rasterstats as rs
import earthpy as et
import earthpy.plot as ep

# Set consistent plotting style
sns.set_style("white")
sns.set(font_scale=1.5)

# Download data and set working directory
data = et.data.get_data("spatial-vector-lidar")
os.chdir(os.path.join(et.io.HOME, 
                      'earth-analytics', 
                      'data'))

# Load & plot the data
sjer_lidar_chm_path = os.path.join("spatial-vector-lidar",
                                   "california", 
                                   "neon-sjer-site",
                                   "2013", 
                                   "lidar", 
                                   "SJER_lidarCHM.tif")

sjer_chm_data = rxr.open_rasterio(sjer_lidar_chm_path, masked=True).squeeze()

# Explore the data by plotting a histogram with earthpy
ax = ep.hist(sjer_chm_data.values,
             figsize=(8, 8),
             colors="purple",
             xlabel="Lidar Estimated Tree Height (m)",
             ylabel="Total Pixels",
             title="Distribution of Pixel Values \n Lidar Canopy Height Model")

# Turn off scientific notation
ax[1].ticklabel_format(useOffset=False,
                       style='plain')

# EXPLORE: View summary statistics of canopy height model
# Notice the mean value with 0's included in the data
print('Mean:', sjer_chm_data.mean().values)
print('Max:', sjer_chm_data.max().values)
print('Min:', sjer_chm_data.min().values)

# CLEANUP: Set CHM values of 0 to NAN (no data or not a number)
sjer_chm_data_no_zeros = sjer_chm_data.where(sjer_chm_data != 0, np.nan)
# View summary statistics of canopy height model after cleaning up the data
print('Mean:', sjer_chm_data_no_zeros.mean().values)
print('Max:', sjer_chm_data_no_zeros.max().values)
print('Min:', sjer_chm_data_no_zeros.min().values)

# Explore the data by plotting a histogram with earthpy
ax = ep.hist(sjer_chm_data_no_zeros.values,
             figsize=(8, 8),
             colors="purple",
             xlabel="Lidar Estimated Tree Height (m)",
             ylabel="Total Pixels",
             title="Distribution of Pixel Values \n Lidar Canopy Height Model")

# Turn off scientific notation
ax[1].ticklabel_format(useOffset=False,
                       style='plain')

sjer_centroids_path = os.path.join("spatial-vector-lidar",
                                   "california", 
                                   "neon-sjer-site",
                                   "vector_data", 
                                   "SJER_plot_centroids.shp")

sjer_plots_points = gpd.read_file(sjer_centroids_path)

type(sjer_plots_points)

# Ensure this is a points layer as you think it is
sjer_plots_points.geom_type.head()

fig, ax = plt.subplots(figsize=(10, 10))

# We plot with the zeros in the data so the CHM can be better represented visually
ep.plot_bands(sjer_chm_data,
              extent=plotting_extent(sjer_chm_data,
                                     sjer_chm_data.rio.transform()),  # Set spatial extent
              cmap='Greys',
              title="San Joachin Field Site \n Vegetation Plot Locations",
              scale=False,
              ax=ax)

sjer_plots_points.plot(ax=ax,
                       marker='s',
                       markersize=45,
                       color='purple')
ax.set_axis_off()
plt.show()

# Create a buffered polygon layer from your plot location points
sjer_plots_poly = sjer_plots_points.copy()

# Buffer each point using a 20 meter circle radius
# and replace the point geometry with the new buffered geometry
sjer_plots_poly["geometry"] = sjer_plots_points.geometry.buffer(20)
sjer_plots_poly.head()

# If the dir does not exist, create it
output_path = os.path.join("spatial-vector-lidar", 
                           "outputs")

if not os.path.isdir(output_path):
    os.mkdir(output_path)

# Export the buffered point layer as a shapefile to use in zonal stats
plot_buffer_path = os.path.join(output_path, 
                                "plot_buffer.shp")

sjer_plots_poly.to_file(plot_buffer_path)

# Extract zonal stats
sjer_tree_heights = rs.zonal_stats(plot_buffer_path,
                                   sjer_chm_data_no_zeros.values,
                                   nodata=-999,
                                   affine=sjer_chm_data_no_zeros.rio.transform(),
                                   geojson_out=True,
                                   copy_properties=True,
                                   stats="count min mean max median")

# View object type
type(sjer_tree_heights)

# Turn extracted data into a pandas geodataframe
sjer_lidar_height_df = gpd.GeoDataFrame.from_features(sjer_tree_heights)
sjer_lidar_height_df.head()

fig, ax = plt.subplots(figsize=(10, 5))

ax.bar(sjer_lidar_height_df['Plot_ID'],
       sjer_lidar_height_df['max'],
       color="purple")

ax.set(xlabel='Plot ID', 
       ylabel='Max Height',
       title='Maximum LIDAR Derived Tree Heights')

plt.setp(ax.get_xticklabels(), rotation=45, horizontalalignment='right')
plt.show()

# Extract zonal stats but retain the individual pixel values
sjer_tree_heights_ras = rs.zonal_stats(plot_buffer_path,
                                       sjer_chm_data_no_zeros.values,
                                       nodata=-999,
                                       affine=sjer_chm_data_no_zeros.rio.transform(),
                                       geojson_out=True,
                                       raster_out=True,
                                       copy_properties=True,
                                       stats="count min mean max median")
# Convert to geodataframe
sjer_lidar_height_df_ras = gpd.GeoDataFrame.from_features(
    sjer_tree_heights_ras)

# View subset of the dataframe
sjer_lidar_height_df_ras[["Plot_ID", 
                          "count", 
                          "geometry",
                          "mini_raster_affine", 
                          "mini_raster_array"]].head()

# Get list of sites
site_names = list(sjer_lidar_height_df_ras["Plot_ID"])

# Convert data in dataframe to a numpy array
arr = np.stack(sjer_lidar_height_df_ras['mini_raster_array'])

# Plot using earthpy
ep.hist(arr,
        bins=[0, 5, 10, 15, 20, 25],
        cols=3,
        title=site_names, 
        figsize=(15, 30))

plt.show()