#!/usr/bin/env python3
"""
Detailed diagnostic to understand NaN pattern in E3SM data.
"""

import xarray as xr
import numpy as np

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

print("="*70)
print("Detailed NaN Pattern Diagnostic")
print("="*70)

ds = xr.open_dataset(e3sm_file)
pr = ds['precip_rate'].values

print(f"\nprecip_rate shape: {pr.shape}")
print(f"Total elements: {pr.size}")

# Check NaN pattern across time
print(f"\nNaN count per timestep (first 10):")
for t in range(min(10, pr.shape[0])):
    nan_count = np.sum(np.isnan(pr[t]))
    total = pr[t].size
    pct = 100 * nan_count / total
    print(f"  t={t}: {nan_count}/{total} NaN ({pct:.1f}%)")

# Check if ALL values at any timestep are NaN
all_nan_timesteps = [t for t in range(pr.shape[0]) if np.all(np.isnan(pr[t]))]
print(f"\nTimesteps where ALL values are NaN: {len(all_nan_timesteps)}")
if all_nan_timesteps:
    print(f"  First few: {all_nan_timesteps[:5]}")

# Check valid data
valid_mask = ~np.isnan(pr)
print(f"\nTotal valid (non-NaN) values: {np.sum(valid_mask)} / {pr.size}")
print(f"Any valid data at all? {np.any(valid_mask)}")

# Check if there's a consistent land/ocean mask
print(f"\nChecking spatial mask consistency:")
first_valid_t = None
for t in range(pr.shape[0]):
    if not np.all(np.isnan(pr[t])):
        first_valid_t = t
        break

if first_valid_t is not None:
    print(f"  First timestep with some valid data: t={first_valid_t}")
    mask_t0 = ~np.isnan(pr[first_valid_t])
    print(f"  Valid points at t={first_valid_t}: {np.sum(mask_t0)}")
    
    # Check if valid points have actual precipitation
    valid_data = pr[first_valid_t][mask_t0]
    print(f"  At valid points - min: {np.min(valid_data):.6e}, max: {np.max(valid_data):.6e}")
else:
    print("  NO timestep has any valid data!")

# Check a specific spatial point that should have data
mid_lat = pr.shape[1] // 2
mid_lon = pr.shape[2] // 2
print(f"\nTime series at center point [{mid_lat}, {mid_lon}]:")
center_ts = pr[:, mid_lat, mid_lon]
print(f"  NaN count: {np.sum(np.isnan(center_ts))}/{len(center_ts)}")
print(f"  First 10 values: {center_ts[:10]}")

# Check max precip location
if np.any(~np.isnan(pr)):
    max_idx = np.unravel_index(np.nanargmax(pr), pr.shape)
    print(f"\nMax precip location: time={max_idx[0]}, lat_idx={max_idx[1]}, lon_idx={max_idx[2]}")
    print(f"Max value: {pr[max_idx]:.6e} mm/s = {pr[max_idx]*3600:.2f} mm/hr")
    
    # Time series at max location
    max_ts = pr[:, max_idx[1], max_idx[2]]
    print(f"Time series at max location - NaN count: {np.sum(np.isnan(max_ts))}/{len(max_ts)}")

ds.close()
print("\n" + "="*70)