import os
import asf_search as asf
from datetime import datetime
from getpass import getpass
from pathlib import Path

#1: Search for and download NISAR GCOV HDF5 products from ASF

#authenticate with Earthdata Login
#create an ASF session object (stores your login + settings for requests)
session = asf.ASFSession()
#prompt for username/password
session.auth_with_creds(input('EDL Username'), getpass('EDL Password'))

#2: Define the search window and area of interest 

#insert start and end dates. Note that dates and times are inclusive. 
start_date = datetime(2025, 12, 4)
end_date = datetime(2025, 12, 5)
# POINT or POLYGON as WKT (well-known-text) using lat/lon pairs 
area_of_interest = "POLYGON((40.9131 12.3904,41.8891 12.3904,41.8891 13.2454,40.9131 13.2454,40.9131 12.3904))" 

#3: Build search options 

#these filters narrow results to NISAR GCOV products over the AOI and time window
opts = asf.ASFSearchOptions(
    **{
        "maxResults": 250,                    # cap the number of returned results
        "intersectsWith": area_of_interest,   # spatial filter (WKT geometry)
        "start": start_date,                  # temporal filter (start)
        "end": end_date,                      # temporal filter (end)
        "processingLevel": ["GCOV"],          # product type / processing level
        "dataset": ["NISAR"],                 # mission/dataset name
        "productionConfiguration": ["PR"],
        "dataMaturity" : "PROVISIONAL",
        "session": session,                   # use authenticated session from above
    }
)

#4: Run the search for HDF5 files
#execute search
response = asf.search(opts=opts)

#exclude QA_STATE files 
pattern = r'^(?!.*QA_STATS).*'

#extract files that end with .h5 
hdf5_files = response.find_urls(extension='.h5', pattern=pattern, directAccess=False)
hdf5_files

import re

pattern = re.compile(r"(NISAR_L2_PR_GCOV(?:_[^_]+){9})_(P\d{5})")

latest_version_dict = {}

for url in hdf5_files:
    m = pattern.search(url)
    if not m:
        continue

    product, crid = m.groups()

    if product not in latest_version_dict or crid > latest_version_dict[product][0]:
        latest_version_dict[product] = (crid, url)

hdf5_files = [i[1] for i in latest_version_dict.values()]
print(f"Retained {len(hdf5_files)} GCOV products:")
hdf5_files

data_dir = Path.home() / "GCOV_data_pixel_conversion_example"
data_dir.mkdir(exist_ok=True)
print(f'data_dir: {data_dir}')

#download files
asf.download_urls(hdf5_files, data_dir, session=session, processes=4)

import h5py

#path to the NISAR GCOV HDF5 file
h5_path = list(data_dir.glob("NISAR_L2_*_GCOV*.h5"))[0]

#GCOV data are stored under the "grids" groups
#select desired frequency: "frequencyA" or "frequencyB"
freq = "frequencyA"
grid_grp = f"/science/LSAR/GCOV/grids/{freq}"

#for this example, we will be using the HHHHH layer 
#for this GCOV file, different layers are available and listed in the output below.
layer = "HHHH"
layer_path = f"{grid_grp}/{layer}"

#print paths 
print("File:", h5_path)
print("Grid group:", grid_grp)
print("Current raster layer:", layer_path)

#examine available GCOV datasets within the file:
with h5py.File(h5_path, "r") as f:
    print("\nAvailable GCOV datasets:")
    for k in f[grid_grp].keys():
        print(" ", k)

#open the GCOV file (in read mode) to read coordinate and spatial information 
with h5py.File(h5_path, "r") as f:

    #in a GCOV file, the "projection" dataset stores coordinate systems
    #the output value will  be an EPSG code     
    epsg = int(f[f"{grid_grp}/projection"][()])

    #xCoordinates and yCoordinates are arrays that store the coordinates of each pixel center
    #these coordinates will later be used to center pixel footprints  
    x = f[f"{grid_grp}/xCoordinates"][:]
    y = f[f"{grid_grp}/yCoordinates"][:]

    #CoordinateSpacing values store the pixel size in meters 
    #this will later be used to map the pixel footprint
    dx = float(f[f"{grid_grp}/xCoordinateSpacing"][()])
    dy = float(f[f"{grid_grp}/yCoordinateSpacing"][()])

#print information
print("GCOV coordinate system (EPSG):", epsg)
print("Raster size:", len(y), "rows x", len(x), "columns")
print("Pixel size:", dx, "m (x) ×", dy, "m (y)")

#define sites
#these coordinates are in EPSG 4326
sites = {
    "Salt_Ponds":          (13.25889, 40.85111),
    "Hypersaline_Lake":    (13.26028, 40.89972),
    "Afdera_Volcano":      (13.08833, 40.85250),
    "Sand_Plains":         (13.27944, 40.80694),
    "Afdera_Lava_Field":   (13.09222, 40.72583),
    "Borawli_Caldera":     (13.30389, 40.98972),
    "Borawli_Lava_Field":  (13.21833, 40.98056),
}

from pyproj import Transformer
import numpy as np

#use transformer to convert ground truth data (EPSG 4326) coordinate system to GCOV coordinate system (EPSG 32637)
transformer = Transformer.from_crs(
    "EPSG:4326",
    f"EPSG:{epsg}",
    always_xy=True
)

site_pixels = {}

for name, (lat, lon) in sites.items():
    #convert the site location into GCOV map coordinates
    x_site, y_site = transformer.transform(lon, lat)

    #find the nearest GCOV pixel by comparing with the grid coordinates
    col = int(np.argmin(np.abs(x - x_site)))
    row = int(np.argmin(np.abs(y - y_site)))

    #store the GCOV pixel index for this site
    site_pixels[name] = (row, col)

#GCOV pixel indices corresponding to each site
site_pixels

import folium
from folium.raster_layers import ImageOverlay
from folium.plugins import MeasureControl, MousePosition
import matplotlib.pyplot as plt
import branca.colormap as cm


################## Display / quicklook settings
#downsample factor for faster visualization 
#the "step" variable tells Python to read every Nth pixel in both row and column directions 
step = 60

#overlay transparency (0 = invisible, 1 = opaque)
overlay_opacity = 0.55

################# Convert coordinates into lat/lon for Folium use
# Folium expects WGS84 lat/lon (EPSG:4326).
to_ll = Transformer.from_crs(f"EPSG:{epsg}", "EPSG:4326", always_xy=True)

################ Read downsampled arrays (HHHH + HVHV) for quicklooks
#note: this does NOT change the original data in the HDF5 file
with h5py.File(h5_path, "r") as f:

    #read downsampled polarization arrays for display
    hh_q = np.array(f[f"{grid_grp}/HHHH"][::step, ::step], dtype=float)
    hv_q = np.array(f[f"{grid_grp}/HVHV"][::step, ::step], dtype=float)

    #read downsampled coordinate arrays (meters)
    #these are used to place the overlay in the correct geographic location
    xq = f[f"{grid_grp}/xCoordinates"][::step]
    yq = f[f"{grid_grp}/yCoordinates"][::step]

################# Stretch display settings to allow easier visualization 
#this affects visualization only — the original GCOV data are unchanged.
#using a single shared vmin/vmax ensures that both polarizations use the same
#brightness scale, preventing misleading visual comparisons between HHHH and HVHV
stack = np.concatenate([hh_q[np.isfinite(hh_q)], hv_q[np.isfinite(hv_q)]])
vmin = np.nanpercentile(stack, 2)
vmax = np.nanpercentile(stack, 98)

print("Stretched display range (2–98% percentile):")
print(f"vmin = {vmin}")
print(f"vmax = {vmax}\n")

################# Convert GCOV rasters to RGBA image 
#folium cannot overlay the GCOV rasters in their current format, so they must 
#be converted to RGBA images for visualization 
def arr_to_rgba(arr, vmin, vmax):
    #mask missing values
    finite = np.isfinite(arr)

    #clip to display range
    arr_clip = np.clip(arr, vmin, vmax)

    #normalize to 0–1 for colormap
    norm = (arr_clip - vmin) / (vmax - vmin + 1e-12)
    norm = np.nan_to_num(norm, nan=0.0)

    #map to RGBA using grayscale
    rgba = plt.cm.gray(norm)

    #set alpha channel to transparent for NaNs
    rgba[..., 3] = finite.astype(float) * 0.85
    return rgba

#convert both polarizations to RGBA using the same stretch
rgba_hh = arr_to_rgba(hh_q, vmin, vmax)
rgba_hv = arr_to_rgba(hv_q, vmin, vmax)

################# Compute overlay bounds (lat/lon) from GCOV x/y (meters)
#ImageOverlay requires SW/NE bounds in lat/lon
xmin, xmax = float(xq.min()), float(xq.max())
ymin, ymax = float(yq.min()), float(yq.max())

lon_sw, lat_sw = to_ll.transform(xmin, ymin)
lon_ne, lat_ne = to_ll.transform(xmax, ymax)

bounds = [[lat_sw, lon_sw], [lat_ne, lon_ne]]

################# Center Folium map over ground truth sites 
center_lat = sum(lat for lat, lon in sites.values()) / len(sites)
center_lon = sum(lon for lat, lon in sites.values()) / len(sites)

m = folium.Map(
    location=[center_lat, center_lon],
    zoom_start=8,
    tiles="Esri.WorldImagery"
)

################# Add toggleable overlays 
ImageOverlay(
    image=rgba_hh,
    bounds=bounds,
    opacity=overlay_opacity,
    name="HHHH",
    interactive=True,
    zindex=2
).add_to(m)

ImageOverlay(
    image=rgba_hv,
    bounds=bounds,
    opacity=overlay_opacity,
    name="HVHV",
    interactive=True,
    zindex=2
).add_to(m)

################# Add ground truth sites and GCOV pixel footprints to map
sites_fg = folium.FeatureGroup(name="Ground Truth Sites", show=True)
pix_fg   = folium.FeatureGroup(name="GCOV Pixel Footprints", show=True)

for name, (lat, lon) in sites.items():

    #groundn truth site marker (blue)
    folium.Marker(
        location=[lat, lon],
        popup=f"{name}",
        icon=folium.Icon(color="blue")
    ).add_to(sites_fg)

    #pixel footprint (yellow)
    row, col = site_pixels[name]

    #pixel center coordinates
    x_center = float(x[col])
    y_center = float(y[row])

    #pixel corners (in meters)
    corners_xy = [
        (x_center - dx/2, y_center - dy/2),
        (x_center + dx/2, y_center - dy/2),
        (x_center + dx/2, y_center + dy/2),
        (x_center - dx/2, y_center + dy/2),
        (x_center - dx/2, y_center - dy/2),
    ]

    #convert corners to lat/lon for folium
    corners_latlon = []
    for xc, yc in corners_xy:
        lonp, latp = to_ll.transform(xc, yc)
        corners_latlon.append([latp, lonp])

    folium.Polygon(
        locations=corners_latlon,
        color="yellow",
        weight=2,
        fill=False,
        popup=f"{name}"
    ).add_to(pix_fg)

sites_fg.add_to(m)
pix_fg.add_to(m)

################# Add Legend
colormap = cm.linear.Greys_09.scale(round(vmin, 5), round(vmax, 5))
colormap.caption = "GCOV backscatter (shared stretch)"
# colormap.add_to(m)

################# Add interactive tools 
MeasureControl(position="topleft").add_to(m)     #measure distance/area
MousePosition(position="bottomleft").add_to(m)  #display lat/lon for cursor 
m.add_child(folium.LatLngPopup())               #click map to get lat/lon popup 

#keep control panel open
folium.LayerControl(collapsed=False).add_to(m)

display(colormap)
display(m)