#!/usr/bin/env python3
import yaml, os, json, sys
from osgeo import gdal, ogr, osr
from netCDF4 import Dataset
import geopandas as gpd
import matplotlib.pyplot as plt
import geojsoncontour as gjc
import time

LEVELS = [0.3,0.5,1,2,5]

rTITLES = gjc.utilities.multipoly.get_contourf_levels(LEVELS,extend='max')

TITLES = ["{} {}".format(x, '') for x in rTITLES]

def main(elemant_map_path:str, inundation_tif_path:str, out_dir:str):
    elemants = yaml.safe_load(open(elemant_map_path,'r'))
    
    # convert geotiff to netCDF
    if not (inundation_tif_path.lower().endswith('.tiff') or  \
        inundation_tif_path.lower().endswith('.geotiff') or \
            inundation_tif_path.lower().endswith('.tif') ):

        raise Exception('invalid filename of raster')

    nc_name, geojson_name = inundation_tif_path+'.nc4', inundation_tif_path+'.geojson'

    print('converting geotiff to netcdf4..')
    ds = gdal.Translate(nc_name, inundation_tif_path, format='NetCDF')
    
    del ds

    # read converted netcdf
    # >>
    nf  =  Dataset(nc_name,'r')

    lats = nf.variables['lat'][:]
    lons = nf.variables['lon'][:]

    depth = nf.variables['Band1'][:]

    print('vectorizing depth map...')
    contours = plt.contourf(lons,lats, depth, levels = LEVELS, extend='max' )

    # plt.show()
    cont_geojson_str = gjc.contourf_to_geojson(contours, geojson_filepath=geojson_name, unit='')

    #cont_geojson = json.loads(cont_geojson_str)

    # remove objects
    del cont_geojson_str
    del contours

    cont_df = gpd.read_file( geojson_name )
    cont_df.set_crs('EPSG:32642', inplace=True, allow_override=True)

    print('buffering inundation map to reduce invalid geometry...')
    cont_df['geometry'] = cont_df.buffer(0)

    # <<
    print('in vector crs', cont_df.crs)

    idata = {
        'titles': TITLES,
        'exp': dict(),
    }

    if not os.path.exists( os.path.join(out_dir,'exp_data') ):
        os.makedirs(os.path.join(out_dir,'exp_data'))

    for elem in elemants.keys():
        print(f'Processing {elem}')
        idata['exp'][elem] = {
            'unit': elemants[elem]['unit'],
            'vals': []
        }

        # Read the element shapefile and reproject
        path = elemants[elem]['path']
        sdf = gpd.read_file(path).to_crs('EPSG:32642')

        # Perform spatial overlay (intersection) between the element shapefile and the polygonized raster
        insec_df = gpd.overlay(sdf, cont_df, how='intersection')

        if not insec_df.empty:
            insec_df.to_file(os.path.join(out_dir, f'exp_data/{elem}_exp.geojson'), driver='GeoJSON')

        for title in TITLES:
            df_r = insec_df.loc[insec_df['title'] == title]
            if elemants[elem]['type'] == 'length':
                idata['exp'][elem]['vals'].append(df_r.length.sum() * elemants[elem]['factor'])

            if elemants[elem]['type'] == 'area':
                idata['exp'][elem]['vals'].append(df_r.area.sum() * elemants[elem]['factor'])

            if elemants[elem]['type'] == 'count':
                idata['exp'][elem]['vals'].append(len(df_r) * elemants[elem]['factor'])

    # Write output data to JSON file
    with open(os.path.join(out_dir, 'exp_vals.json'), 'w') as wf:
        json.dump(idata, wf)

if __name__ == '__main__':
    print('Started processing exposure...')
    t0 = time.perf_counter()
    main(sys.argv[1], sys.argv[2], sys.argv[3])
    t1 = time.perf_counter()
    print(f'Time taken: {(t1 - t0) / 60:0.2f} minutes')