import os
os.environ['USE_PYGEOS'] = '0'
import os.path
from pathlib import Path
# Data objects
import geopandas as gpd
import pandas as pd
import rioxarray as rxr
import xarray as xr

# GDPTools objects
from gdptools import ZonalGen
from gdptools.data.user_data import UserTiffData

# Plotting objects
import matplotlib.pyplot as plt

# set some input variables
data_dir = Path.home() / 'data'
output_dir = data_dir / 'results'
if not output_dir.exists():
    os.mkdir(output_dir)
# path to hru dataset
# the_hrus = os.path.join(data_dir, "NHM_20_810.gdb")
the_hrus = data_dir / "NHM_20_810.gdb"
# hrus dataset hru id field name 
hru_id = "hru_id"
print(f"the_hrus exists: {the_hrus.exists()}")
# NHM parameter name
param_name = "hru_elev"
# The statistic used to define the parameter value
function = "median"
in_raster = "HI_DataLayersforth/dem/dem.tif"
# dem raster
# the_raster = os.path.join(data_dir, in_raster)
the_raster = data_dir / in_raster
print(f"the_raster exists: {the_raster.exists()}")

# Use geopandas to read shape file,
# Best if shape file only has feature id, and geometry if possible.
gdf = gpd.read_file(the_hrus, layer='nhru')
id_feature = hru_id
print(f"Number of features: {len(gdf.groupby(id_feature))}")
# Filter the geodataframe to just the feature id and the geometry columns
gdf = gdf[["hru_id", "geometry"]]
gdf

# Be sure to use rioxarray to read the tiff
rds_ras = rxr.open_rasterio(the_raster)

rds_ras

# These params are used to fill out the TiffAttributes class
tx_name= 'x'
ty_name = 'y'
band = 1
bname = 'band'
crs = 26904 #
varname = "elev" # not currently used
categorical = False # is the data categorical or not, if the data are integers it should
                   # probably be categorical.
data = UserTiffData(
    var=varname,
    ds=rds_ras,
    proj_ds=crs,
    x_coord=tx_name,
    y_coord=ty_name,
    band=band,
    bname=bname,
    f_feature=gdf,
    id_feature=id_feature,
    proj_feature=crs
)

zonal_gen = ZonalGen(
    user_data=data,
    zonal_engine="serial",
    zonal_writer="csv",
    out_path=".",
    file_prefix="Hi_elev_stats"
)
stats = zonal_gen.calculate_zonal(categorical=categorical)

stats.head()

# Function to output results to NHM parameter file format
def stats_to_paramcsv2(param_name, stats, categorical, function, hru, hru_id, output_dir):
    # purpose: to write result of gdptools to parameter file
    #          and to add param to hru geopandas df for plots
    # arguments:
    #   param_name - character - specifying the ParamDB parameter name, 
    #      for example:  hru_slope  hru_elev
    #   stats - pandas df, output from gdptools function
    #          (cols like a pandas describe)
    #   categorical - Boolean (it determines the stats field names)
    #   function - chararacter -(function valid values for categorical False: sum, mean, median min, max)
    #                      function for categorical True is majority (or top)
    #   hru - geopandas df 
    #   hru_id - character - hru id field name
    #    output_dir - character - parameter file output path
    outpth = os.path.join(output_dir, '%s.csv'%(param_name))
    #outpth = output_dir / (param_name + '.csv')
    if(categorical):
        thepar_df = stats[['top']].rename(columns = {'top':param_name})
        hru[param_name] = stats['top'].values
        # rename the id field parameter file    
        thepar_df.index.name = "$id" 
        thepar_df.to_csv(outpth, index = True)
    else:
        stats.rename(columns = {'50%':'median'}, inplace = True)
        function_list = ["mean", "min","max", "sum","median"]
        if (function in function_list):           
            # for slope div 100
            if param_name == "hru_slope":
                thepar_df = stats[[function]].rename(columns = {function:param_name})
                thepar_df[param_name] = thepar_df[param_name] / 100
                hru[param_name] = stats[function].values / 100
            else : 
                thepar_df = stats[[function]].rename(columns = {function:param_name})
                hru[param_name] = stats[function].values      
            thepar_df.index.name = "$id" 
            thepar_df.to_csv(outpth, index = True)
        else:
            print("ERROR: invalid function: {0}".format(function))
            print("No param file generated")
            thepar_df = pd.DataFrame()
    return(thepar_df)

def vis_the_map(the_df,param_name,the_cmap):
# histogram
    plt.hist(the_df[param_name], edgecolor = 'black')
    plt.xlabel(param_name)
    plt.ylabel('value count')
    plt.show()
    # map
    hru.plot(param_name, cmap = the_cmap, legend = True)
    plt.title(param_name)
    plt.show()
    return

hru = gdf.copy()
tst_df = stats_to_paramcsv2(param_name, stats, categorical, function, hru, hru_id, output_dir)

tst_df.head()

vis_the_map(tst_df,param_name,'viridis')

param_name = "hru_slope"
function = "mean"
in_raster = "HI_DataLayersforth/slope/slope.tif"
# dem raster
# the_raster = os.path.join(data_dir, in_raster)
the_raster = data_dir / in_raster
print(f"the_raster exists: {the_raster.exists()}")

# use rioxarray to read the tiff
rds_ras = rxr.open_rasterio(the_raster)
rds_ras

# These params are used to fill out the TiffAttributes class
tx_name= 'x'
ty_name = 'y'
band = 'band'
crs = 26904
varname = "slope" # not currently used
categorical = False # is the data categorical or not, if the data are integers it should
                   # probably be categorical.
data = UserTiffData(
    var=varname,
    ds=rds_ras,
    proj_ds=crs,
    x_coord=tx_name,
    y_coord=ty_name,
    band=1,
    bname=band,
    f_feature=gdf,
    id_feature=id_feature,
    proj_feature=crs
)

zonal_gen = ZonalGen(
    user_data=data,
    zonal_engine="serial",
    zonal_writer="csv",
    out_path=".",
    file_prefix="Hi_slope_stats"
)
stats = zonal_gen.calculate_zonal(categorical=categorical)

stats.head()

hru = gdf.copy()
# generate parameter file
tst_df = stats_to_paramcsv2(param_name, stats, categorical, function, hru, hru_id, output_dir)
tst_df.head()

vis_the_map(tst_df,param_name,'viridis')

param_name = "cov_type"
function = "top"
in_raster = "HI_DataLayersforth/LULC/LULC.tif"
# dem raster
# the_raster = os.path.join(data_dir, in_raster)
the_raster = data_dir / in_raster
print(f"the_raster exists: {the_raster.exists()}")

# use rioxarray to read the tiff
rds_ras = rxr.open_rasterio(the_raster)
rds_ras

# These params are used to fill out the TiffAttributes class
tx_name= 'x'
ty_name = 'y'
band = 'band'
crs = 26904
varname = "LULC" # not currently used
categorical = True # is the data categorical or not, if the data are integers it should
                   # probably be categorical.
data = UserTiffData(
    var=varname,
    ds=rds_ras,
    proj_ds=crs,
    x_coord=tx_name,
    y_coord=ty_name,
    band=1,
    bname=band,
    f_feature=gdf,
    id_feature=id_feature,
    proj_feature=crs
)

zonal_gen = ZonalGen(
    user_data=data,
    zonal_engine="serial",
    zonal_writer="csv",
    out_path=".",
    file_prefix="Hi_cov_type_stats"
)
stats = zonal_gen.calculate_zonal(categorical=categorical)

stats

hru = gdf.copy()
# generate parameter file
tst_df = stats_to_paramcsv2(param_name, stats, categorical, function, hru, hru_id, output_dir)
tst_df.head()

vis_the_map(tst_df,param_name,'viridis')