#!/usr/bin/env python3
"""
Fix precip_accum in existing E3SM extracted file.
The precip_rate is valid but precip_accum has NaN propagation issues.
"""

import xarray as xr
import numpy as np
import os

e3sm_file = "/global/cfs/projectdirs/m4486/Haochen/Extracted_Data/Hourly_Precip/E3SM_THREAD_New/E3SM_THREAD_aug7_control.nc"
output_file = "/global/cfs/projectdirs/m4486/Haochen/Extracted_Data/Hourly_Precip/E3SM_THREAD_New/E3SM_THREAD_aug7_control_fixed.nc"

print("="*70)
print("Fixing E3SM precip_accum")
print("="*70)

ds = xr.open_dataset(e3sm_file)
pr = ds['precip_rate'].values  # mm/s

print(f"precip_rate shape: {pr.shape}")
print(f"precip_rate max: {np.nanmax(pr):.6e} mm/s = {np.nanmax(pr)*3600:.2f} mm/hr")

# dt = 10 min = 600 sec
dt_sec = 600.0

# Method: Fill NaN with 0 for cumsum, keep result (no NaN restoration for ocean)
# This way land points accumulate correctly
pr_filled = np.nan_to_num(pr, nan=0.0)

print(f"\nAfter filling NaN with 0:")
print(f"  pr_filled max: {np.max(pr_filled):.6e}")
print(f"  pr_filled min: {np.min(pr_filled):.6e}")

# Compute cumulative sum
precip_accum = np.cumsum(pr_filled * dt_sec, axis=0)

print(f"\nprecip_accum computed:")
print(f"  Shape: {precip_accum.shape}")
print(f"  Min: {np.min(precip_accum):.2f}")
print(f"  Max: {np.max(precip_accum):.2f}")
print(f"  Any > 0? {np.any(precip_accum > 0)}")

# Check final timestep
print(f"\nFinal timestep (t=-1):")
print(f"  Max accum: {np.max(precip_accum[-1]):.2f} mm")
print(f"  Mean accum: {np.mean(precip_accum[-1]):.2f} mm")

# Create new dataset
ds_out = xr.Dataset({
    'precip_rate': (['time', 'lat', 'lon'], pr),
    'precip_accum': (['time', 'lat', 'lon'], precip_accum.astype(np.float32)),
}, coords={
    'time': ds['time'].values,
    'lat': ds['lat'].values,
    'lon': ds['lon'].values
})

# Copy attributes
ds_out.attrs = ds.attrs
ds_out.attrs['fix_applied'] = 'Recomputed precip_accum with NaN filled as 0'

ds_out['precip_rate'].attrs = ds['precip_rate'].attrs
ds_out['precip_accum'].attrs = {
    'long_name': 'Accumulated Precipitation',
    'units': 'mm',
    'description': 'Time-integrated total precipitation (NaN treated as 0)'
}

# Save
print(f"\nSaving to: {output_file}")
encoding = {
    'precip_rate': {'zlib': True, 'complevel': 4, 'dtype': 'float32'},
    'precip_accum': {'zlib': True, 'complevel': 4, 'dtype': 'float32'},
}
ds_out.to_netcdf(output_file, encoding=encoding)

# Verify
print("\nVerifying saved file...")
ds_check = xr.open_dataset(output_file)
pa_check = ds_check['precip_accum'].values
print(f"  precip_accum max: {np.nanmax(pa_check):.2f} mm")
print(f"  precip_accum any > 0? {np.any(pa_check > 0)}")
ds_check.close()

ds.close()

# Optionally replace the original
print(f"\nTo replace original file, run:")
print(f"  mv {output_file} {e3sm_file}")

print("\n" + "="*70)
print("Done!")
print("="*70)