#!/usr/bin/env python3
"""
WRF Precipitation Extraction Script
Extracts RAINNC, XLAT, XLONG from WRF output files and saves to consolidated NetCDF files.
This makes subsequent analysis much faster.

Updated: TAMU_WRF now points to the new rerun at
  /global/cfs/projectdirs/m4486/www/TRACERMIP/TAMU_WRF_FSBM2
"""
import argparse
import os
import glob
import warnings
from datetime import datetime
import numpy as np
import xarray as xr

warnings.filterwarnings('ignore', category=FutureWarning)

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

# Model configurations with paths for each case and experiment
MODEL_CONFIG = {
    'ANL_WRF': {
        'june17': {
            'control': '/global/cfs/projectdirs/m4486/www/TRACERMIP/WRF_FSBM_ANL/17June_500m/RefCCN',
            'high_ccn': '/global/cfs/projectdirs/m4486/www/TRACERMIP/WRF_FSBM_ANL/17June_500m/HighCCN',
            'low_ccn': '/global/cfs/projectdirs/m4486/www/TRACERMIP/WRF_FSBM_ANL/17June_500m/LowCCN',
        },
        'aug7': {
            'control': '/global/cfs/projectdirs/m4486/www/TRACERMIP/WRF_FSBM_ANL/07Aug_500m/RefCCN',
            'high_ccn': '/global/cfs/projectdirs/m4486/www/TRACERMIP/WRF_FSBM_ANL/07Aug_500m/HighCCN',
            'low_ccn': '/global/cfs/projectdirs/m4486/www/TRACERMIP/WRF_FSBM_ANL/07Aug_500m/LowCCN',
        },
        'output_dir': os.path.join(OUTPUT_BASE, 'ANL_WRF'),
    },
    'NU_WRF': {
        'june17': {
            'control': '/global/cfs/projectdirs/m4486/Haochen/NU_WRF_Combined/june17/control',
            'high_ccn': '/global/cfs/projectdirs/m4486/Haochen/NU_WRF_Combined/june17/high_ccn',
            'low_ccn': '/global/cfs/projectdirs/m4486/Haochen/NU_WRF_Combined/june17/low_ccn',
        },
        'aug7': {
            'control': '/global/cfs/projectdirs/m4486/Haochen/NU_WRF_Combined/aug7/control',
            'high_ccn': '/global/cfs/projectdirs/m4486/Haochen/NU_WRF_Combined/aug7/high_ccn',
            'low_ccn': '/global/cfs/projectdirs/m4486/Haochen/NU_WRF_Combined/aug7/low_ccn',
        },
        'output_dir': os.path.join(OUTPUT_BASE, 'NU_WRF'),
    },
    'TAMU_WRF': {
        'june17': {
            'control': '/global/cfs/projectdirs/m4486/www/TRACERMIP/TAMU_WRF_FSBM2/New_runs_April2026/20220617/ctrl_run',
            'high_ccn': '/global/cfs/projectdirs/m4486/www/TRACERMIP/TAMU_WRF_FSBM2/New_runs_April2026/20220617/aerosol_3x_run',
            'low_ccn': '/global/cfs/projectdirs/m4486/www/TRACERMIP/TAMU_WRF_FSBM2/New_runs_April2026/20220617/aerosol_0.3x_run',
        },
        'aug7': {
            'control': '/global/cfs/projectdirs/m4486/www/TRACERMIP/TAMU_WRF_FSBM2/New_runs_April2026/20220807/control',
            'high_ccn': '/global/cfs/projectdirs/m4486/www/TRACERMIP/TAMU_WRF_FSBM2/New_runs_April2026/20220807/3x_aerosol',
            'low_ccn': '/global/cfs/projectdirs/m4486/www/TRACERMIP/TAMU_WRF_FSBM2/New_runs_April2026/20220807/0.3x_aerosol',
        },
        'output_dir': os.path.join(OUTPUT_BASE, 'TAMU_WRF'),
    },
}

# Case time ranges
CASE_CONFIG = {
    'june17': {
        'start_datetime': datetime(2022, 6, 17, 6, 0),
        'end_datetime': datetime(2022, 6, 18, 6, 0),
        'date_label': 'June 17, 2022'
    },
    'aug7': {
        'start_datetime': datetime(2022, 8, 7, 6, 0),
        'end_datetime': datetime(2022, 8, 8, 6, 0),
        'date_label': 'August 7, 2022'
    }
}


# ============================================================================
# EXTRACTION FUNCTIONS
# ============================================================================
def find_wrf_files(data_path, case_config):
    """Find all WRF output files in the given path within the time range."""
    # Try different file patterns
    patterns = [
        os.path.join(data_path, 'wrfout_d02*'),
        os.path.join(data_path, '2min_output', 'wrfout_d02*'),
        os.path.join(data_path, '10min_output', 'wrfout_d02*'),
    ]
    
    all_files = []
    for pattern in patterns:
        files = sorted(glob.glob(pattern))
        all_files.extend(files)
    
    if not all_files:
        print("  Warning: No WRF files found in {0}".format(data_path))
        return []
    
    # Remove duplicates and sort
    all_files = sorted(list(set(all_files)))
    
    # Filter by time range
    start_dt = case_config['start_datetime']
    end_dt = case_config['end_datetime']
    
    filtered_files = []
    for fpath in all_files:
        try:
            # Extract timestamp from filename
            fname = os.path.basename(fpath)
            timestamp_str = fname.split("_d02_")[-1]
            timestamp = datetime.strptime(timestamp_str, "%Y-%m-%d_%H:%M:%S")
            if start_dt <= timestamp <= end_dt:
                filtered_files.append((fpath, timestamp))
        except Exception as e:
            continue
    
    # Sort by timestamp
    filtered_files.sort(key=lambda x: x[1])
    
    return filtered_files


def extract_precipitation(file_list, output_file):
    """
    Extract RAINNC, XLAT, XLONG from WRF files and save to NetCDF.
    
    Parameters:
    -----------
    file_list : list of tuples
        List of (filepath, timestamp) tuples
    output_file : str
        Path to output NetCDF file
    """
    if not file_list:
        print("  No files to process")
        return False
    
    print("  Processing {0} files...".format(len(file_list)))
    
    # Lists to store data
    rainnc_list = []
    time_list = []
    xlat = None
    xlong = None
    
    for i, (fpath, timestamp) in enumerate(file_list):
        try:
            ds = xr.open_dataset(fpath, decode_timedelta=False)
            
            # Get RAINNC - ensure we get a 2D numpy array
            if 'RAINNC' in ds:
                rainnc = ds['RAINNC']
                
                # Remove Time dimension if present
                if 'Time' in rainnc.dims:
                    rainnc = rainnc.isel(Time=0)
                
                # Get the numpy array values
                rainnc_values = np.array(rainnc.values)
                
                # Ensure 2D
                if rainnc_values.ndim > 2:
                    rainnc_values = rainnc_values.squeeze()
                
                rainnc_list.append(rainnc_values)
                time_list.append(timestamp)
            
            # Get coordinates (only once)
            if xlat is None:
                if 'XLAT' in ds:
                    xlat_var = ds['XLAT']
                    if 'Time' in xlat_var.dims:
                        xlat_var = xlat_var.isel(Time=0)
                    xlat = np.array(xlat_var.values)
                    if xlat.ndim > 2:
                        xlat = xlat.squeeze()
                        
                if 'XLONG' in ds:
                    xlong_var = ds['XLONG']
                    if 'Time' in xlong_var.dims:
                        xlong_var = xlong_var.isel(Time=0)
                    xlong = np.array(xlong_var.values)
                    if xlong.ndim > 2:
                        xlong = xlong.squeeze()
            
            ds.close()
            
            # Progress indicator
            if (i + 1) % 50 == 0 or (i + 1) == len(file_list):
                print("    Processed {0}/{1} files...".format(i + 1, len(file_list)))
                
        except Exception as e:
            print("    Error reading {0}: {1}".format(os.path.basename(fpath), e))
            continue
    
    if not rainnc_list:
        print("  No valid data extracted")
        return False
    
    # Create xarray Dataset
    print("  Creating output dataset with {0} timesteps...".format(len(rainnc_list)))
    
    # Stack RAINNC along time dimension
    rainnc_array = np.stack(rainnc_list, axis=0)
    time_array = np.array(time_list, dtype='datetime64[ns]')
    
    # Get spatial dimensions
    ny, nx = rainnc_array.shape[1], rainnc_array.shape[2]
    
    print("  RAINNC shape: {0}".format(rainnc_array.shape))
    print("  XLAT shape: {0}".format(xlat.shape))
    print("  XLONG shape: {0}".format(xlong.shape))
    
    # Create Dataset using xr.Dataset constructor with data_vars dict
    ds_out = xr.Dataset(
        data_vars={
            'RAINNC': xr.DataArray(
                data=rainnc_array,
                dims=['Time', 'south_north', 'west_east'],
                attrs={
                    'description': 'Accumulated total grid scale precipitation',
                    'units': 'mm',
                }
            ),
            'XLAT': xr.DataArray(
                data=xlat,
                dims=['south_north', 'west_east'],
                attrs={
                    'description': 'Latitude',
                    'units': 'degrees_north',
                }
            ),
            'XLONG': xr.DataArray(
                data=xlong,
                dims=['south_north', 'west_east'],
                attrs={
                    'description': 'Longitude', 
                    'units': 'degrees_east',
                }
            ),
        },
        coords={
            'Time': time_array,
        },
        attrs={
            'description': 'Extracted WRF precipitation data',
            'source': 'WRF output files',
            'created': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
        }
    )
    
    # Save to NetCDF with compression
    print("  Saving to {0}...".format(output_file))
    encoding = {
        'RAINNC': {'zlib': True, 'complevel': 4},
        'XLAT': {'zlib': True, 'complevel': 4},
        'XLONG': {'zlib': True, 'complevel': 4},
    }
    ds_out.to_netcdf(output_file, encoding=encoding)
    
    # Print summary
    print("  Done! Output shape: {0}".format(rainnc_array.shape))
    print("  Time range: {0} to {1}".format(time_array[0], time_array[-1]))
    print("  File size: {0:.1f} MB".format(os.path.getsize(output_file) / 1e6))
    
    return True


def process_model(model_name, case_name, experiment):
    """Process a single model/case/experiment combination."""
    model_config = MODEL_CONFIG[model_name]
    case_config = CASE_CONFIG[case_name]
    
    # Get input path
    if case_name not in model_config or experiment not in model_config[case_name]:
        print("  Skipping {0}/{1}/{2} - path not configured".format(model_name, case_name, experiment))
        return False
    
    input_path = model_config[case_name][experiment]
    
    # Create output directory
    output_dir = model_config['output_dir']
    os.makedirs(output_dir, exist_ok=True)
    
    # Output filename
    output_file = os.path.join(output_dir, "{0}_{1}_{2}.nc".format(model_name, case_name, experiment))
    
    print("\n" + "=" * 60)
    print("Processing: {0} / {1} / {2}".format(model_name, case_name, experiment))
    print("=" * 60)
    print("  Input: {0}".format(input_path))
    print("  Output: {0}".format(output_file))
    
    # Check if output already exists
    if os.path.exists(output_file):
        print("  Output file already exists. Skipping...")
        print("  (Delete the file to re-extract)")
        return True
    
    # Find files
    file_list = find_wrf_files(input_path, case_config)
    print("  Found {0} files in time range".format(len(file_list)))
    
    if not file_list:
        return False
    
    # Extract and save
    success = extract_precipitation(file_list, output_file)
    
    return success


# ============================================================================
# MAIN
# ============================================================================
def main():
    parser = argparse.ArgumentParser(
        description='Extract WRF precipitation data and save to NetCDF'
    )
    parser.add_argument('--model', choices=['ANL_WRF', 'NU_WRF', 'TAMU_WRF', 'all'], 
                        default='all', help='Model to process')
    parser.add_argument('--case', choices=['june17', 'aug7', 'both'], 
                        default='both', help='Case to process')
    parser.add_argument('--experiment', choices=['control', 'high_ccn', 'low_ccn', 'all'],
                        default='all', help='Experiment to process')
    args = parser.parse_args()
    
    # Determine what to process
    models = list(MODEL_CONFIG.keys()) if args.model == 'all' else [args.model]
    cases = ['june17', 'aug7'] if args.case == 'both' else [args.case]
    experiments = ['control', 'high_ccn', 'low_ccn'] if args.experiment == 'all' else [args.experiment]
    
    print("=" * 60)
    print("WRF Precipitation Data Extraction")
    print("=" * 60)
    print("Models: {0}".format(models))
    print("Cases: {0}".format(cases))
    print("Experiments: {0}".format(experiments))
    print("Output base: {0}".format(OUTPUT_BASE))
    
    # Process each combination
    success_count = 0
    fail_count = 0
    
    for model in models:
        for case in cases:
            for experiment in experiments:
                try:
                    result = process_model(model, case, experiment)
                    if result:
                        success_count += 1
                    else:
                        fail_count += 1
                except Exception as e:
                    print("  Error: {0}".format(e))
                    fail_count += 1
    
    # Summary
    print("\n" + "=" * 60)
    print("SUMMARY")
    print("=" * 60)
    print("Successful: {0}".format(success_count))
    print("Failed: {0}".format(fail_count))
    print("Total: {0}".format(success_count + fail_count))
    print("\nDone!")


if __name__ == '__main__':
    main()