#!/usr/bin/env python3

import os
import glob
import numpy as np
import xarray as xr
from datetime import datetime

# ============================================================================
# PATH CONFIGURATION
# ============================================================================
MRMS_BASE_PATH = "/global/cfs/projectdirs/m4486/Haochen/mrms_data"
OUTPUT_BASE = "/global/cfs/projectdirs/m4486/Haochen/Extracted_Data/Hourly_Precip/MRMS"

# Domain to extract (covers both red boxes and large domain with buffer)
# Red box June17: lon -96.25 to -95.35, lat 29.6 to 30.2
# Red box Aug7: lon -95.92 to -95.1, lat 29.5 to 30.17  
# Large domain: lon -96.1 to -94.0, lat 28.5 to 30.45
# Extract domain (with buffer): lon -97.0 to -93.5, lat 28.0 to 31.0
EXTRACT_DOMAIN = {
    'lon_min': -97.0,
    'lon_max': -93.5,
    'lat_min': 28.0,
    'lat_max': 31.0
}

CASE_CONFIG = {
    'june17': {
        'mrms_dir': 'june17/mrms_2mins',
        'start_datetime': datetime(2022, 6, 17, 12, 0),
        'end_datetime': datetime(2022, 6, 18, 6, 0),
        'date_label': 'June 17, 2022'
    },
    'aug7': {
        'mrms_dir': 'aug7/mrms_2mins',
        'start_datetime': datetime(2022, 8, 7, 12, 0),
        'end_datetime': datetime(2022, 8, 8, 6, 0),
        'date_label': 'August 7, 2022'
    }
}


def parse_mrms_filename(filename):
    """Parse datetime from MRMS filename.
    
    Example: MRMS_PrecipRate_00.00_20220617-120000.grib2
    """
    basename = os.path.basename(filename)
    # Extract datetime string: YYYYMMDD-HHMMSS
    parts = basename.split('_')
    if len(parts) >= 4:
        dt_str = parts[3].replace('.grib2', '')  # 20220617-120000
        try:
            return datetime.strptime(dt_str, '%Y%m%d-%H%M%S')
        except ValueError:
            pass
    return None


def extract_mrms_data(case_name):
    """
    Extract MRMS precipitation rate data from GRIB2 files.
    Subsets to Houston region and processes in batches to reduce memory usage.
    
    Parameters:
    -----------
    case_name : str
        Case name (june17 or aug7)
    """
    import gc
    
    config = CASE_CONFIG[case_name]
    
    print(f"\n{'='*70}")
    print(f"Processing MRMS: {config['date_label']}")
    print(f"{'='*70}")
    
    # Build file path
    mrms_path = os.path.join(MRMS_BASE_PATH, config['mrms_dir'], 'MRMS_PrecipRate_00.00_*.grib2')
    
    # Output file
    os.makedirs(OUTPUT_BASE, exist_ok=True)
    output_file = os.path.join(OUTPUT_BASE, f"MRMS_{case_name}_10min.nc")
    
    print(f"Input path: {os.path.dirname(mrms_path)}")
    print(f"Output file: {output_file}")
    print(f"Extract domain: lat {EXTRACT_DOMAIN['lat_min']}-{EXTRACT_DOMAIN['lat_max']}°N, "
          f"lon {EXTRACT_DOMAIN['lon_min']}-{EXTRACT_DOMAIN['lon_max']}°E")
    
    # Check if output exists
    if os.path.exists(output_file):
        print(f"Output file already exists. Skipping... (delete to re-extract)")
        return True
    
    # Find all GRIB2 files
    mrms_files = sorted(glob.glob(mrms_path))
    
    if not mrms_files:
        print(f"ERROR: No GRIB2 files found at {mrms_path}")
        return False
    
    print(f"Found {len(mrms_files)} GRIB2 files")
    
    # Filter files by time range
    start_dt = config['start_datetime']
    end_dt = config['end_datetime']
    
    filtered_files = []
    for f in mrms_files:
        file_dt = parse_mrms_filename(f)
        if file_dt and start_dt <= file_dt <= end_dt:
            filtered_files.append((file_dt, f))
    
    filtered_files.sort(key=lambda x: x[0])
    print(f"Files in time range: {len(filtered_files)}")
    
    if not filtered_files:
        print("ERROR: No files in the specified time range")
        return False
    
    # First pass: get coordinates from first file
    print("Getting coordinates from first file...")
    first_file = filtered_files[0][1]
    ds_first = xr.open_dataset(first_file, engine='cfgrib')
    
    lat_vals = ds_first['latitude'].values
    lon_vals = ds_first['longitude'].values
    lon_vals_converted = np.where(lon_vals > 180, lon_vals - 360, lon_vals)
    
    # Find indices for the subset domain
    lat_idx = np.where((lat_vals >= EXTRACT_DOMAIN['lat_min']) & 
                       (lat_vals <= EXTRACT_DOMAIN['lat_max']))[0]
    lon_idx = np.where((lon_vals_converted >= EXTRACT_DOMAIN['lon_min']) & 
                       (lon_vals_converted <= EXTRACT_DOMAIN['lon_max']))[0]
    
    lat_subset = lat_vals[lat_idx]
    lon_subset = lon_vals_converted[lon_idx]
    
    print(f"  Original grid size: {len(lat_vals)} x {len(lon_vals)}")
    print(f"  Subset grid size: {len(lat_subset)} x {len(lon_subset)}")
    print(f"  Subset lat range: {lat_subset.min():.2f} to {lat_subset.max():.2f}")
    print(f"  Subset lon range: {lon_subset.min():.2f} to {lon_subset.max():.2f}")
    
    ds_first.close()
    del ds_first, lat_vals, lon_vals, lon_vals_converted
    gc.collect()
    
    # Pre-allocate array for all data
    n_times = len(filtered_files)
    n_lat = len(lat_subset)
    n_lon = len(lon_subset)
    
    print(f"  Pre-allocating array: {n_times} x {n_lat} x {n_lon} = {n_times * n_lat * n_lon * 4 / 1e6:.1f} MB")
    
    # Use float32 to save memory
    precip_rate_arr = np.full((n_times, n_lat, n_lon), np.nan, dtype=np.float32)
    time_list = []
    
    # Process files
    print("Reading MRMS GRIB2 files...")
    successful_reads = 0
    
    for i, (file_dt, fpath) in enumerate(filtered_files):
        if i % 50 == 0:
            print(f"  Processing file {i+1}/{len(filtered_files)}: {os.path.basename(fpath)}")
            gc.collect()  # Force garbage collection every 50 files
        
        try:
            ds = xr.open_dataset(fpath, engine='cfgrib')
            
            # cfgrib reads MRMS PrecipRate as 'unknown'
            if 'unknown' in ds:
                prate = ds['unknown'].values
            else:
                var_names = list(ds.data_vars)
                if var_names:
                    prate = ds[var_names[0]].values
                else:
                    ds.close()
                    time_list.append(file_dt)
                    continue
            
            # Squeeze and subset
            prate = np.squeeze(prate)
            prate_subset = prate[np.ix_(lat_idx, lon_idx)].astype(np.float32)
            
            # Replace missing/fill values with NaN
            prate_subset[(prate_subset < 0) | (prate_subset > 1e10)] = np.nan
            
            # Store directly in pre-allocated array
            precip_rate_arr[i, :, :] = prate_subset
            time_list.append(file_dt)
            successful_reads += 1
            
            # Explicitly close and delete
            ds.close()
            del ds, prate, prate_subset
            
        except Exception as e:
            if "index file" not in str(e):
                print(f"  Error reading {os.path.basename(fpath)}: {e}")
            time_list.append(file_dt)
            continue
    
    print(f"  Successfully read: {successful_reads}/{len(filtered_files)} files")
    gc.collect()
    
    if successful_reads == 0:
        print("ERROR: No data extracted from MRMS files")
        return False
    
    # Create xarray dataset
    print("Creating xarray dataset...")
    time_arr = np.array(time_list, dtype='datetime64[ns]')
    
    ds_2min = xr.Dataset({
        'precip_rate': (['time', 'lat', 'lon'], precip_rate_arr)
    }, coords={
        'time': time_arr,
        'lat': lat_subset,
        'lon': lon_subset
    })
    
    # Free the numpy array
    del precip_rate_arr
    gc.collect()
    
    # Resample to 10-minute intervals
    print("Resampling to 10-minute intervals...")
    ds_10min = ds_2min.resample(time='10min').mean()
    
    # Free 2-min dataset
    ds_2min.close()
    del ds_2min
    gc.collect()
    
    print(f"  After resampling: {len(ds_10min['time'])} timesteps")
    
    # Add metadata
    ds_10min.attrs = {
        'title': 'MRMS Precipitation Rate Data (10-min) - Houston Region',
        'case': case_name,
        'description': 'MRMS radar precipitation rate resampled from 2-min to 10-min, subset to Houston region',
        'source': 'MRMS (Multi-Radar Multi-Sensor)',
        'original_resolution': '2-minute',
        'resampled_resolution': '10-minute',
        'resampling_method': 'mean',
        'domain': f"lat {EXTRACT_DOMAIN['lat_min']}-{EXTRACT_DOMAIN['lat_max']}N, lon {EXTRACT_DOMAIN['lon_min']}-{EXTRACT_DOMAIN['lon_max']}E",
        'creation_date': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
        'creator': 'Haochen Zhang',
        'n_files_processed': successful_reads
    }
    
    ds_10min['precip_rate'].attrs = {
        'long_name': 'Precipitation Rate',
        'units': 'mm/hr',
        'description': 'MRMS radar precipitation rate'
    }
    
    ds_10min['lat'].attrs = {'long_name': 'Latitude', 'units': 'degrees_north'}
    ds_10min['lon'].attrs = {'long_name': 'Longitude', 'units': 'degrees_east'}
    
    # Save to NetCDF
    print(f"Writing to output file...")
    encoding = {
        'precip_rate': {'zlib': True, 'complevel': 4, 'dtype': 'float32'},
        'lat': {'zlib': True, 'complevel': 4, 'dtype': 'float32'},
        'lon': {'zlib': True, 'complevel': 4, 'dtype': 'float32'}
    }
    
    ds_10min.to_netcdf(output_file, encoding=encoding)
    
    # Get file size
    file_size_mb = os.path.getsize(output_file) / (1024**2)
    
    print(f"\n{'='*70}")
    print(f"SUCCESS: MRMS data extracted!")
    print(f"  Output file: {output_file}")
    print(f"  Time steps: {len(ds_10min['time'])}")
    print(f"  Grid: {n_lat} lat x {n_lon} lon")
    print(f"  File size: {file_size_mb:.2f} MB")
    print(f"{'='*70}")
    
    ds_10min.close()
    gc.collect()
    
    return True


def main():
    """Main function to process all MRMS cases."""
    
    print("\n" + "="*70)
    print("MRMS Precipitation Rate Extraction Script")
    print("="*70)
    print(f"Start time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    
    results = []
    
    for case_name in CASE_CONFIG.keys():
        success = extract_mrms_data(case_name)
        results.append({
            'case': case_name,
            'success': success
        })
    
    # Print final summary
    print("\n" + "="*70)
    print("FINAL SUMMARY")
    print("="*70)
    print(f"Total cases: {len(results)}")
    print(f"Successful: {sum(1 for r in results if r['success'])}")
    print(f"Failed: {sum(1 for r in results if not r['success'])}")
    print("\nDetailed results:")
    
    for result in results:
        status = "✓ SUCCESS" if result['success'] else "✗ FAILED"
        print(f"  {status}: MRMS - {result['case']}")
    
    print(f"\nOutput directory: {OUTPUT_BASE}")
    print(f"End time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("="*70 + "\n")


if __name__ == "__main__":
    main()