# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # open FIA data points

# +
# Mostly-standard imports
import os
import sys 
sys.path.append('/content')
import xarray as xr
import tempfile
import numpy as np
import shutil
import urllib
import matplotlib.pyplot as plt
import pandas as pd
import geopandas as gpd
from shapely.geometry import Point
from glob import glob
import utm

import src.azuretools as azt

#from mpl_toolkits.basemap import Basemap

# Less-common-but-still-pip-installable imports
from netCDF4 import Dataset

# pip install progressbar2, not progressbar
import progressbar


# This will contain just the .nc files
nasadem_file_list = None
                                   
temp_dir = os.path.join(tempfile.gettempdir(),'nasadem')
os.makedirs(temp_dir,exist_ok=True)

# %matplotlib inline
# %load_ext autoreload
# %autoreload 2

# +
blob = '/home/datablob'
fiafile = f'{blob}/fia_no_pltcn.csv'
plots = gpd.GeoDataFrame(pd.read_csv(fiafile))
plots.Latitude = plots.LAT.astype('float')
plots.Longitude = plots.LON.astype('float')
geometry = [Point(xy) for xy in zip(plots.Longitude, plots.Latitude)]
crs = {'init': 'epsg:4326'} 
geo_df = gpd.GeoDataFrame(plots, crs=crs, geometry=geometry)

#subset by STATECD
statecd = pd.read_csv(glob(f'{blob}/supp_data/*.csv')[0])
query = 'CA'

code = statecd[statecd['STATEAB']== query]['STATECD'].values[0]
subdf = geo_df[geo_df['STATECD']==code]
fields = ['INDEX', 'INVYR', 'LAT', 'LON']
subdf = subdf.astype({'LAT':'float', 'LON':'float', 'INVYR':'int32'})
subdf = gpd.GeoDataFrame(subdf, geometry=gpd.points_from_xy(subdf.LON, subdf.LAT))
subdf = subdf.set_crs(epsg=4326)
subdf = subdf.rename(columns={'Unnamed: 0':'INDEX'})

#statebounds for reducing query on cloud
statebounds = glob(f'{blob}/supp_data/shp/*.shp')[0]
states = gpd.read_file(statebounds)
state = states[states['STPOSTAL']==query]
state_shape = state.geometry
# -

'''
run this if you are hoping to build your URL list
'''
url_list = []
for i,s in subdf.iterrows():
    tile_of_interest = [s['geometry'].y,s['geometry'].x]
    tile_name = lat_lon_to_nasadem_tile(tile_of_interest[0],tile_of_interest[1])
    url = azt.nasadem_blob_root + tile_name
    url_list.append(url)
subdf['DEM_URL'] = url_list

# # selecting by number of cells

# +
'''
if we want to select by number of cells around the center point
'''
cell_buffer = 10
if not os.path.exists(f'{blob}/training_tiles/NASADEM/{query}'):
    os.makedirs(f'{blob}/training_tiles/NASADEM/{query}')
    
for sample_index in range(len(subdf)):
    sample_point = gpd.GeoDataFrame(subdf.iloc[sample_index]).T
    sample_point.crs = {'init':'epsg:4269'}
    tile_of_interest = [sample_point.geometry.y.values[0], sample_point.geometry.x.values[0]]
    tile_name = azt.lat_lon_to_nasadem_tile(tile_of_interest[0],\
                                            tile_of_interest[1],\
                                            nasadem_file_list)
    url = azt.nasadem_blob_root + tile_name
    fn = azt.download_url(url,progress_updater = azt.DownloadProgressBar())
    fh = xr.open_dataset(fn, engine='h5netcdf')
    yi = np.argmin(np.abs(fh.lat - float(sample_point['LAT'])).values)
    xi = np.argmin(np.abs(fh.lon - float(sample_point['LON'])).values)
    #set the area around those point defined as c
    data_chunk = fh.isel(lat=slice(yi-cell_buffer,yi+cell_buffer),\
                       lon=slice(xi-cell_buffer,xi+cell_buffer))
    
    
    outfilename = f'{blob}/training_tiles/NASADEM/{query}/{sample_point.INDEX.values[0]}.nc'
    data_chunk.to_netcdf(outfilename)
# -

# # selecting by distance

# +
'''
if we want to select by distance around center point
'''
if not os.path.exists(f'{blob}/training_tiles/NASADEM/{query}'):
    os.makedirs(f'{blob}/training_tiles/NASADEM/{query}')
    
#given we know the utm values can we go look around 1km?
buffer_dist = 500 #500 meters both sides of the center for 1km image

for sample_index in range(len(subdf)):
    sample_point = gpd.GeoDataFrame(subdf.iloc[sample_index]).T
    sample_point.crs = {'init':'epsg:4269'}
    tile_of_interest = [sample_point.geometry.y.values[0], sample_point.geometry.x.values[0]]
    tile_name = azt.lat_lon_to_nasadem_tile(tile_of_interest[0],\
                                            tile_of_interest[1],\
                                            nasadem_file_list)
    url = azt.nasadem_blob_root + tile_name
    fn = azt.download_url(url,progress_updater = azt.DownloadProgressBar())
    fh = xr.open_dataset(fn, engine='h5netcdf')
    
    e_utm_values, n_utm_values, zone, hemi = utm.from_latlon(latitude=float(sample_point.geometry.y), \
                                                         longitude=float(sample_point.geometry.x))
    bounds = {'miny':n_utm_values-buffer_dist, \
              'maxy':n_utm_values+buffer_dist, \
              'minx':e_utm_values-buffer_dist, \
              'maxx':e_utm_values+buffer_dist}
    
    bbox = {'lowerleft':(bounds['miny'], bounds['minx'], zone, hemi),\
            'upperleft':(bounds['maxy'], bounds['minx'], zone, hemi),\
            'upperright':(bounds['maxy'], bounds['maxx'], zone, hemi),\
            'lowerright':(bounds['miny'], bounds['maxx'], zone, hemi)}
    #transform bounds to lat lon
    
    bbox_latlon = {k:utm.to_latlon(v[1], v[0], zone_number=v[2], zone_letter=v[3]) for k,v in bbox.items()}
    minx, maxx, miny, maxy = bbox_latlon['lowerleft'][1],bbox_latlon['lowerright'][1],\
                            bbox_latlon['lowerleft'][0],bbox_latlon['upperright'][0]
    minyi = np.argmin(np.abs(fh.lat - miny).values)
    maxyi = np.argmin(np.abs(fh.lat - maxy).values)
    minxi = np.argmin(np.abs(fh.lon - minx).values)
    maxxi = np.argmin(np.abs(fh.lon - maxx).values)


    #set the area around those point defined as c
    data_chunk = fh.isel(lat=slice(maxyi, minyi),\
                   lon=slice(minxi, maxxi))

    
    
    outfilename = f'{blob}/training_tiles/NASADEM/{query}/{sample_point.INDEX.values[0]}.nc'
    data_chunk.to_netcdf(outfilename)

# +
fh = xr.open_dataset(fn, engine='h5netcdf')
heights = fh['NASADEM_HGT'][:]
lons = fh.variables['lon'][:]
lats = fh.variables['lat'][:]

min_height = np.min(heights)
max_height = np.max(heights)
height_units = fh['NASADEM_HGT'].units

fh.close()

print('Height ranges from {} {} to {} {}'.format(min_height,height_units,
      max_height,height_units))

extent = [np.min(lons), np.max(lons), np.min(lats), np.max(lats)]
plt.imshow(heights,extent=extent)
plt.xlabel('Longitude')
plt.ylabel('Latitude')
cb = plt.colorbar()
cb.set_label('Height ({})'.format(height_units))

StarcoderdataJupyterScriptsDedupFiltered