# Data ETL imports
import numpy as np
import glob
import rioxarray

# GIS imports
import shapely.wkt as wkt
import rasterio
from rasterio.merge import merge
from rasterio.plot import show
from rasterio.warp import reproject, transform_bounds
from rasterio.crs import CRS
import folium
from folium import plugins
from rtc_utils import getbasemaps
from rasterio.warp import calculate_default_transform, Resampling
from rasterio.mask import mask
from shapely.geometry import box

# Plotting imports
import matplotlib as mpl
import matplotlib.pyplot as plt

# Misc imports
from pathlib import Path
import sys
from datetime import datetime
from functools import partial

sys.path.append('../../')

import warnings
warnings.filterwarnings('ignore')

import os
os.environ["AWS_NO_SIGN_REQUEST"] = "YES"

# set event date
event_date = datetime(year=2023, month=2, day=18)

data_dir = 's3://opera-pst-rs-pop1/products/RTC_S1/'
filepath = Path('../data/rtc_manifest.txt')

with open(filepath, 'r') as f:
    file_list = f.readlines()

# Given a filename, parse out the time stamp and return it as a datetime object
def parse_date_time(filename):
    yyyymmdd = filename.split("/")[-1].split('_')[4][:8]
    year, month, day = [int(x) for x in [yyyymmdd[:4], yyyymmdd[4:6], yyyymmdd[6:]]]
    return datetime(year=year, month=month, day=day)

# Given a filename, return if it is before the reference date
def return_is_preevent(filename, reference_date):
    return parse_date_time(filename) < reference_date

# Given a filename, return if it is on or after the reference date
def return_is_postevent(filename, reference_date):
    return parse_date_time(filename) >= reference_date

# From the manifest list, separate out the pre and post event filepaths
fullPath_preEvent = sorted([x.strip() for x in file_list if return_is_preevent(x, event_date)])
fullPath_postEvent = sorted([x.strip() for x in file_list if return_is_postevent(x, event_date)])

# Given the list of pre and post event paths, return the names of the VV and VH files
def return_s3_location(polarization, filepath):
     return f"{filepath}/{filepath.split('/')[-1]}_{polarization}.tif"

vv_paths_preevent = list(map(partial(return_s3_location, "VV"), fullPath_preEvent))
vv_paths_postevent = list(map(partial(return_s3_location, "VV"), fullPath_postEvent))

vh_paths_preevent = list(map(partial(return_s3_location, "VH"), fullPath_preEvent))
vh_paths_postevent = list(map(partial(return_s3_location, "VH"), fullPath_postEvent))

def filter_by_burstids(lines, burstids):
    return [line for line in lines if any([bid in line for bid in burstids])]

burstids = ['112912-IW2', '112913-IW2']

filtered_vv_list_preevent = filter_by_burstids(vv_paths_preevent, burstids)
filtered_vv_list_postevent = filter_by_burstids(vv_paths_postevent, burstids)

filtered_vh_list_preevent = filter_by_burstids(vh_paths_preevent, burstids)
filtered_vh_list_postevent = filter_by_burstids(vh_paths_postevent, burstids)

# location for saved rasters
output_path = Path("../data")
output_path.mkdir(exist_ok=True, parents=True)

def write_burst_mean(burst_id, burst_list, prefix, polarization):
    burst_list = [x for x in burst_list if burst_id in x] # filter list by burst_id
        
    # Open the first raster to get its properties
    with rasterio.open(burst_list[0]) as src:
        mean_profile = src.profile
        data_type = src.dtypes[0]
        shape = src.shape

    mean_array = np.zeros(shape, dtype=data_type)
    
    # Loop through each raster and accumulate pixel values
    for raster_path in burst_list:
        with rasterio.open(raster_path) as src:
            raster_data = src.read(1)
            mean_array += raster_data

    mean_array = mean_array / len(burst_list)

    # write out averaged burst
    output_file = output_path / f"{burst_id}_mean_{prefix}Event_{polarization}.tif"
    # burst_mean, burst_mean_profile = write_burst_mean([x for x in filtered_vv_list_preevent if burstids[0] in x])
    with rasterio.open(output_file, 'w', **mean_profile) as dst:
        dst.write(mean_array, 1)

    return output_file

# There are 2 polarizations, 2 burst IDs, and pre- and post- event lists. In all, there are 8 rasters that need to be computed and written out
# Averaging pre-event bursts
burst1_vv_preevent = write_burst_mean(burstids[0], filtered_vv_list_preevent, "pre", "VV")
burst2_vv_preevent = write_burst_mean(burstids[1], filtered_vv_list_preevent, "pre", "VV")

burst1_vh_preevent = write_burst_mean(burstids[0], filtered_vh_list_preevent, "pre", "VH")
burst2_vh_preevent = write_burst_mean(burstids[1], filtered_vh_list_preevent, "pre", "VH")

# Averaging post-event bursts
burst1_vv_postevent = write_burst_mean(burstids[0], filtered_vv_list_postevent, "post", "VV")
burst2_vv_postevent = write_burst_mean(burstids[1], filtered_vv_list_postevent, "post", "VV")

burst1_vh_postevent = write_burst_mean(burstids[0], filtered_vh_list_postevent, "post", "VH")
burst2_vh_postevent = write_burst_mean(burstids[1], filtered_vh_list_postevent, "post", "VH")

# Merge the pre-event datasets
mosaic_preEvent_VV, transform_preEvent_VV = merge([burst1_vv_preevent, burst2_vv_preevent], nodata=np.nan, method='first')
mosaic_preEvent_VH, transform_preEvent_VH = merge([burst1_vh_preevent, burst2_vh_preevent], nodata=np.nan, method='first')

# create a sample profile to be used later to write out the merged profiles
with rasterio.open(burst1_vv_preevent) as ds:
    out_profile = ds.profile
out_profile.update(transform=transform_preEvent_VV, width=mosaic_preEvent_VV.shape[2], height=mosaic_preEvent_VV.shape[1])

# Merge the post-event datasets
mosaic_postEvent_VV, transform_postEvent_VV = merge([burst1_vv_postevent, burst2_vv_postevent], nodata=np.nan, method='first')
mosaic_postEvent_VH, transform_postEvent_VH = merge([burst1_vh_postevent, burst2_vh_postevent], nodata=np.nan, method='first')

# Convert rasters to dB scale to help compress dynamic range
mosaic_preEvent_VV_db = 10*np.log10(mosaic_preEvent_VV)
mosaic_preEvent_VH_db = 10*np.log10(mosaic_preEvent_VH)

mosaic_postEvent_VV_db = 10*np.log10(mosaic_postEvent_VV)
mosaic_postEvent_VH_db = 10*np.log10(mosaic_postEvent_VH)

# helper function for folium visualization
def colorize(array=[], cmap='RdBu', cmin=[], cmax=[]):
    normed_data = (array - cmin) / (cmax - cmin)    
    cm = plt.cm.get_cmap(cmap)
    return cm(normed_data) 

mask_bounds = [421820.0, 7368620.0, 433490.0, 7376000.0]
mask_shape = box(*mask_bounds) # we are going to crop the raster to an area of interest
# mask_shape = world[world.name == 'Brazil'].iloc[0].geometry

# Load up VV data
input_file = output_path / "logRatio_VV.tif"
with rasterio.open(input_file) as ds:
    vv_img, transform = mask(ds, [mask_shape], nodata=np.nan)
    img_bounds = ds.bounds
    src_transform = ds.transform
    src_crs = ds.crs
mask_bounds_4326 = transform_bounds(src_crs, CRS.from_epsg(4326), *mask_bounds)
mask_shape_4326 = box(*mask_bounds_4326)

# Load up VV data
input_file = output_path / "logRatio_VH.tif"
with rasterio.open(input_file) as ds:
    vh_img, transform = mask(ds, [mask_shape], nodata=np.nan)

left, bottom, right, top = img_bounds
transform, width, height = calculate_default_transform(src_crs, CRS.from_epsg(4326), width = vv_img.shape[-1], height=vv_img.shape[-2], left=left, bottom=bottom, right=right, top=top)

vv_img_transformed = np.zeros((1, height, width))*np.nan
_ = reproject(vv_img, vv_img_transformed, src_crs = src_crs, dst_crs=CRS.from_epsg(4326), src_transform=src_transform, src_nodata=np.nan, dst_transform=transform)

vh_img_transformed = np.zeros((1, height, width))*np.nan
_ = reproject(vh_img, vh_img_transformed, src_crs = src_crs, dst_crs=CRS.from_epsg(4326), src_transform=src_transform, src_nodata=np.nan, dst_transform=transform)

# Let's add an ocean mask to aid with the visualization. We use a separately downloaded NASA DEM tile for this purpose
with rasterio.open('../data/s24w046.hgt') as ds:
    dem_img, dem_transform = mask(ds, [mask_shape_4326])
    dem_crs = ds.crs

dem_array = np.zeros_like(vv_img_transformed)

reproject(dem_img,
            dem_array,
            src_transform=dem_transform,
            src_crs=dem_crs,
            dst_transform=transform,
            dst_crs=CRS.from_epsg(4326),
            dst_nodata=np.nan,
            resampling=Resampling.nearest)

# We can use the fact that land will have non-zero elevation values to mask out the ocean
dem_array = np.where(dem_array[0]>0, 1, np.nan)