#!/usr/bin/env python3
import argparse
import os
import warnings
from datetime import datetime
import matplotlib
matplotlib.use('Agg')
warnings.filterwarnings('ignore', category=FutureWarning)
warnings.filterwarnings('ignore', category=UserWarning)
import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
import matplotlib.dates as mdates

# ============================================================================
# PATH CONFIGURATION
# ============================================================================
EXTRACTED_DATA_BASE = "/global/cfs/projectdirs/m4486/Haochen/Extracted_Data/Hourly_Precip"
OUTPUT_DIR = "/global/cfs/projectdirs/m4486/Haochen/Precip/output_figures"

# E3SM_THREAD paths - different for June vs August
E3SM_THREAD_DIR_JUNE = os.path.join(EXTRACTED_DATA_BASE, 'E3SM_THREAD')      # Original path for June
E3SM_THREAD_DIR_AUG = os.path.join(EXTRACTED_DATA_BASE, 'E3SM_THREAD_New')   # New path for August

# CM1 domain center point for coordinate conversion
CM1_CENTER_LON = -95.0792
CM1_CENTER_LAT = 29.4719

# Large domain analysis box (same for both cases)
LARGE_DOMAIN_BOX = {'lon_min': -96.1, 'lon_max': -94.0, 'lat_min': 28.5, 'lat_max': 30.45}

# Model configurations (WRF models)
MODEL_CONFIG = {
    'ANL_WRF': {
        'extracted_dir': os.path.join(EXTRACTED_DATA_BASE, 'ANL_WRF'),
        'color': 'black',
        'linestyle': '-',
        'label': 'ANL_WRF'
    },
    'NU_WRF': {
        'extracted_dir': os.path.join(EXTRACTED_DATA_BASE, 'NU_WRF'),
        'color': 'blue',
        'linestyle': '-',
        'label': 'NU_WRF'
    },
    'TAMU_WRF': {
        'extracted_dir': os.path.join(EXTRACTED_DATA_BASE, 'TAMU_WRF'),
        'color': 'red',
        'linestyle': '-',
        'label': 'TAMU_WRF'
    },
    'WRF_DRI': {
        'extracted_dir': os.path.join(EXTRACTED_DATA_BASE, 'WRF_DRI'),
        'color': '#00CED1',  # Cyan/teal
        'linestyle': '-',
        'label': 'WRF_DRI'
    }
}

# CM1 configuration
CM1_CONFIG = {
    'extracted_dir': os.path.join(EXTRACTED_DATA_BASE, 'CM1'),
    'color': 'green',
    'linestyle': '-',
    'label': 'CM1'
}

# ICON_KIT configuration
ICONKIT_CONFIG = {
    'extracted_dir': os.path.join(EXTRACTED_DATA_BASE, 'ICON_KIT'),
    'color': 'purple',
    'linestyle': '-',
    'label': 'ICON_KIT'
}

# E3SM_THREAD configuration (directory will be set dynamically based on case)
E3SM_CONFIG = {
    'extracted_dir': None,  # Will be set per case
    'color': 'orange',
    'linestyle': '-',
    'label': 'E3SM_THREAD'
}

# RAMS configuration
RAMS_CONFIG = {
    'extracted_dir': os.path.join(EXTRACTED_DATA_BASE, 'RAMS'),
    'color': 'brown',
    'linestyle': '-',
    'label': 'RAMS'
}

# MRMS configuration
MRMS_CONFIG = {
    'extracted_dir': os.path.join(EXTRACTED_DATA_BASE, 'MRMS'),
    'color': 'gray',
    'linestyle': '-',
    'linewidth': 2.5,
    'label': 'MRMS'
}

CASE_CONFIG = {
    'june17': {
        'start_datetime': datetime(2022, 6, 17, 12, 0),
        'end_datetime': datetime(2022, 6, 18, 6, 0),
        'date_label': 'June 17, 2022',
        'rate_ylim': (0, 0.75),  # Y-axis limits for precipitation rate plot
        'e3sm_dir': E3SM_THREAD_DIR_JUNE  # Use original E3SM path
    },
    'aug7': {
        'start_datetime': datetime(2022, 8, 7, 12, 0),
        'end_datetime': datetime(2022, 8, 8, 6, 0),
        'date_label': 'August 7, 2022',
        'rate_ylim': (0, 0.4),  # Y-axis limits for precipitation rate plot
        'accum_ylim': (0, 1),   # Y-axis limits for accumulated precipitation plot
        'e3sm_dir': E3SM_THREAD_DIR_AUG   # Use new E3SM path
    }
}


# ============================================================================
# HELPER FUNCTION TO GET E3SM DIRECTORY
# ============================================================================
def get_e3sm_dir(case_name):
    """Get the appropriate E3SM_THREAD directory based on case."""
    return CASE_CONFIG[case_name]['e3sm_dir']


# ============================================================================
# CM1 COORDINATE CONVERSION
# ============================================================================
def km_to_latlon(xh_km, yh_km, center_lon=CM1_CENTER_LON, center_lat=CM1_CENTER_LAT):
    """Convert CM1 km coordinates to lat/lon."""
    km_per_deg_lat = 111.0
    km_per_deg_lon = 111.0 * np.cos(np.radians(center_lat))
    
    x_center_km = (xh_km.min() + xh_km.max()) / 2.0
    y_center_km = (yh_km.min() + yh_km.max()) / 2.0
    
    lon_1d = center_lon + (xh_km - x_center_km) / km_per_deg_lon
    lat_1d = center_lat + (yh_km - y_center_km) / km_per_deg_lat
    
    lon_2d, lat_2d = np.meshgrid(lon_1d, lat_1d)
    
    return lon_2d, lat_2d


# ============================================================================
# DATA READING FUNCTIONS
# ============================================================================
def select_region(bounding_box, current_ds, lat_name='lat', lon_name='lon'):
    """Select a spatial region from the dataset."""
    min_lon, min_lat, max_lon, max_lat = bounding_box
    lats = current_ds[lat_name]
    lons = current_ds[lon_name]
    region_mask = (lats >= min_lat) & (lats <= max_lat) & (lons >= min_lon) & (lons <= max_lon)
    return current_ds.where(region_mask, drop=True)


def select_region_2d(bounding_box, lon_2d, lat_2d, data):
    """Select a spatial region using 2D arrays."""
    lon_min, lat_min, lon_max, lat_max = bounding_box
    mask = (lon_2d >= lon_min) & (lon_2d <= lon_max) & (lat_2d >= lat_min) & (lat_2d <= lat_max)
    
    if data.ndim == 2:
        masked_data = np.where(mask, data, np.nan)
    else:
        masked_data = np.where(mask[np.newaxis, :, :], data, np.nan)
    
    return masked_data


def read_extracted_mrms_data(case_name):
    """Read pre-extracted MRMS precipitation rate data (10-min NetCDF)."""
    config = CASE_CONFIG[case_name]
    
    extracted_file = os.path.join(MRMS_CONFIG['extracted_dir'], f"MRMS_{case_name}_10min.nc")
    
    if not os.path.exists(extracted_file):
        print(f"Warning: Extracted MRMS file not found: {extracted_file}")
        print(f"  Run extract_mrms_precipitation.py first!")
        return None
    
    print(f"Reading extracted MRMS data: {os.path.basename(extracted_file)}")
    ds = xr.open_dataset(extracted_file)
    
    # Filter by time range
    start_dt = np.datetime64(config['start_datetime'])
    end_dt = np.datetime64(config['end_datetime'])
    ds = ds.sel(time=slice(start_dt, end_dt))
    
    print(f"  Time steps: {len(ds['time'])}")
    
    return ds


def read_extracted_wrf_data(model_name, case_name, experiment='control'):
    """Read pre-extracted WRF precipitation data with 10-minute resampling."""
    config = CASE_CONFIG[case_name]
    model_config = MODEL_CONFIG[model_name]
    
    extracted_file = os.path.join(
        model_config['extracted_dir'],
        f"{model_name}_{case_name}_{experiment}.nc"
    )
    
    if not os.path.exists(extracted_file):
        print(f"Warning: Extracted file not found: {extracted_file}")
        return None
    
    print(f"Reading extracted data: {os.path.basename(extracted_file)}")
    ds = xr.open_dataset(extracted_file)
    
    # Filter by time range
    start_dt = np.datetime64(config['start_datetime'])
    end_dt = np.datetime64(config['end_datetime'])
    ds = ds.sel(Time=slice(start_dt, end_dt))
    
    # Resample RAINNC to 10-minute intervals if needed
    if 'RAINNC' in ds:
        ds_resampled = ds.resample(Time='10min').nearest()
        
        # Calculate rain rate from resampled data
        time_diff = ds_resampled['Time'].diff('Time').dt.total_seconds() / 3600.0
        rainnc_diff = ds_resampled['RAINNC'].diff('Time')
        rain_rate = rainnc_diff / time_diff
        
        first_rate = (rain_rate.isel(Time=0) * 0)
        first_rate = first_rate.assign_coords(Time=ds_resampled['Time'].isel(Time=0)).expand_dims('Time')
        ds_resampled['RAIN_RATE'] = xr.concat([first_rate, rain_rate], dim='Time')
        
        print(f"  Time steps: {len(ds_resampled['Time'])}")
        return ds_resampled
    
    return ds


def read_extracted_cm1_data(case_name, experiment='control'):
    """Read pre-extracted CM1 precipitation data."""
    config = CASE_CONFIG[case_name]
    
    extracted_file = os.path.join(CM1_CONFIG['extracted_dir'], f"CM1_{case_name}_{experiment}.nc")
    
    if not os.path.exists(extracted_file):
        print(f"Warning: Extracted file not found: {extracted_file}")
        return None
    
    print(f"Reading extracted data: {os.path.basename(extracted_file)}")
    ds = xr.open_dataset(extracted_file)
    
    # Filter by time range
    start_dt = np.datetime64(config['start_datetime'])
    end_dt = np.datetime64(config['end_datetime'])
    ds = ds.sel(time=slice(start_dt, end_dt))
    
    # Get coordinates
    if 'xh' in ds.coords and 'yh' in ds.coords:
        xh = ds['xh'].values
        yh = ds['yh'].values
        lon_2d, lat_2d = km_to_latlon(xh, yh)
        ds['XLONG'] = xr.DataArray(lon_2d, dims=['yh', 'xh'])
        ds['XLAT'] = xr.DataArray(lat_2d, dims=['yh', 'xh'])
    
    # Map CM1 variable names to standard names
    # CM1 uses 'prate' for precipitation rate and 'rain' for accumulated
    # Check both data_vars and all variables
    ds_vars = list(ds.data_vars) + list(ds.coords)
    print(f"  Available variables: {list(ds.data_vars)}")
    
    if 'prate' in ds_vars:
        ds['RAIN_RATE'] = ds['prate'] * 3600.0  # Convert from mm/s to mm/hr
        print(f"  Mapped prate -> RAIN_RATE (converted mm/s to mm/hr)")
    if 'rain' in ds_vars:
        ds['RAINNC'] = ds['rain']  # Accumulated precipitation
        print(f"  Mapped rain -> RAINNC")
    
    print(f"  Time steps: {len(ds['time'])}")
    
    return ds


def read_extracted_iconkit_data(case_name, experiment='control'):
    """Read pre-extracted ICON_KIT precipitation data."""
    config = CASE_CONFIG[case_name]
    
    extracted_file = os.path.join(ICONKIT_CONFIG['extracted_dir'], f"ICON_KIT_{case_name}_{experiment}.nc")
    
    if not os.path.exists(extracted_file):
        print(f"Warning: Extracted file not found: {extracted_file}")
        return None
    
    print(f"Reading extracted data: {os.path.basename(extracted_file)}")
    ds = xr.open_dataset(extracted_file)
    
    # Filter by time range
    start_dt = np.datetime64(config['start_datetime'])
    end_dt = np.datetime64(config['end_datetime'])
    ds = ds.sel(time=slice(start_dt, end_dt))
    
    # Create 2D lat/lon arrays
    lon_1d = ds['lon'].values
    lat_1d = ds['lat'].values
    lon_2d, lat_2d = np.meshgrid(lon_1d, lat_1d)
    
    ds['XLONG'] = xr.DataArray(lon_2d, dims=['lat', 'lon'])
    ds['XLAT'] = xr.DataArray(lat_2d, dims=['lat', 'lon'])
    
    # Convert tot_prec_rate from kg/m²/s to mm/hr
    if 'tot_prec_rate' in ds:
        ds['RAIN_RATE'] = ds['tot_prec_rate'] * 3600.0
    
    # Rename tot_prec to RAINNC for consistency
    if 'tot_prec' in ds:
        ds['RAINNC'] = ds['tot_prec']
    
    print(f"  Time steps: {len(ds['time'])}")
    
    return ds


def read_extracted_e3sm_data(case_name, experiment='control'):
    """Read pre-extracted E3SM_THREAD precipitation data.
    
    Uses case-specific directory:
    - June case: E3SM_THREAD (original)
    - August case: E3SM_THREAD_New (new rerun data)
    """
    config = CASE_CONFIG[case_name]
    
    # Get case-specific E3SM directory
    e3sm_dir = get_e3sm_dir(case_name)
    
    extracted_file = os.path.join(e3sm_dir, f"E3SM_THREAD_{case_name}_{experiment}.nc")
    
    if not os.path.exists(extracted_file):
        print(f"Warning: Extracted file not found: {extracted_file}")
        print(f"  Run Extract_E3SM_precip.py first!")
        return None
    
    print(f"Reading extracted data: {os.path.basename(extracted_file)}")
    print(f"  Using E3SM directory: {e3sm_dir}")
    ds = xr.open_dataset(extracted_file)
    
    # Filter by time range
    start_dt = np.datetime64(config['start_datetime'])
    end_dt = np.datetime64(config['end_datetime'])
    ds = ds.sel(time=slice(start_dt, end_dt))
    
    # Create 2D lat/lon arrays
    lon_1d = ds['lon'].values
    lat_1d = ds['lat'].values
    lon_2d, lat_2d = np.meshgrid(lon_1d, lat_1d)
    
    ds['XLONG'] = xr.DataArray(lon_2d, dims=['lat', 'lon'])
    ds['XLAT'] = xr.DataArray(lat_2d, dims=['lat', 'lon'])
    
    # Convert precip_rate from mm/s to mm/hr
    if 'precip_rate' in ds:
        ds['RAIN_RATE'] = ds['precip_rate'] * 3600.0
    
    # Rename precip_accum to RAINNC for consistency
    if 'precip_accum' in ds:
        ds['RAINNC'] = ds['precip_accum']
    
    print(f"  Time steps: {len(ds['time'])}")
    
    return ds


def read_extracted_rams_data(case_name, experiment='control'):
    """Read pre-extracted RAMS precipitation data."""
    config = CASE_CONFIG[case_name]
    
    extracted_file = os.path.join(RAMS_CONFIG['extracted_dir'], f"RAMS_{case_name}_{experiment}.nc")
    
    if not os.path.exists(extracted_file):
        print(f"Warning: Extracted file not found: {extracted_file}")
        print(f"  Run Extract_RAMS_precip.py first!")
        return None
    
    print(f"Reading extracted data: {os.path.basename(extracted_file)}")
    ds = xr.open_dataset(extracted_file)
    
    # Filter by time range
    start_dt = np.datetime64(config['start_datetime'])
    end_dt = np.datetime64(config['end_datetime'])
    ds = ds.sel(time=slice(start_dt, end_dt))
    
    # RAMS has 2D lat/lon arrays with dims ['y', 'x']
    # Assign to XLONG/XLAT for consistency with other models
    ds['XLONG'] = ds['lon']
    ds['XLAT'] = ds['lat']
    
    # Convert tot_prec_rate from kg/m²/s to mm/hr
    if 'tot_prec_rate' in ds:
        ds['RAIN_RATE'] = ds['tot_prec_rate'] * 3600.0
    
    # Rename tot_prec to RAINNC for consistency
    if 'tot_prec' in ds:
        ds['RAINNC'] = ds['tot_prec']
    
    print(f"  Time steps: {len(ds['time'])}")
    
    return ds


# ============================================================================
# PLOTTING FUNCTIONS
# ============================================================================
def plot_multimodel_timeseries(mrms_ds, wrf_datasets, cm1_ds, iconkit_ds, e3sm_ds, rams_ds, config, save_path=None):
    """Plot multi-model precipitation time series comparison using large domain."""
    # Use the large domain box for all cases
    box = LARGE_DOMAIN_BOX
    lon_min, lon_max = box['lon_min'], box['lon_max']
    lat_min, lat_max = box['lat_min'], box['lat_max']
    
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 9))
    
    # MRMS data
    if mrms_ds is not None:
        mrms_sub = select_region([lon_min, lat_min, lon_max, lat_max], mrms_ds)
        
        # MRMS precip_rate is in mm/hr
        mrms_rate = mrms_sub['precip_rate'].where(mrms_sub['precip_rate'] > 0, 0)
        mrms_hourly = mrms_rate.mean(dim=['lat', 'lon'], skipna=True).values
        
        # Accumulated precipitation: sum of (rate * time_interval)
        # 10 min = 1/6 hour
        mrms_accum = np.cumsum(mrms_hourly * (10.0 / 60.0))
        
        mrms_time = mrms_ds['time'].values
        
        ax1.plot(mrms_time, mrms_hourly, linestyle=MRMS_CONFIG['linestyle'], 
                 color=MRMS_CONFIG['color'], linewidth=MRMS_CONFIG['linewidth'], 
                 label=MRMS_CONFIG['label'])
        ax2.plot(mrms_time, mrms_accum, linestyle=MRMS_CONFIG['linestyle'], 
                 color=MRMS_CONFIG['color'], linewidth=MRMS_CONFIG['linewidth'], 
                 label=MRMS_CONFIG['label'])
    
    # WRF models
    for model_name, (wrf_ds, model_config) in wrf_datasets.items():
        if wrf_ds is None:
            continue
        
        lons = wrf_ds['XLONG'].values
        lats = wrf_ds['XLAT'].values
        
        if lons.ndim == 1:
            lon_2d, lat_2d = np.meshgrid(lons, lats)
        else:
            lon_2d, lat_2d = lons, lats
        
        rainnc = wrf_ds['RAINNC'].values
        rain_rate = wrf_ds['RAIN_RATE'].values if 'RAIN_RATE' in wrf_ds else None
        
        rainnc_masked = select_region_2d([lon_min, lat_min, lon_max, lat_max], lon_2d, lat_2d, rainnc)
        rainnc_from_start = rainnc_masked - rainnc_masked[0:1, :, :]
        wrf_accum_ts = np.nanmean(rainnc_from_start, axis=(1, 2))
        
        if rain_rate is not None:
            rain_rate_masked = select_region_2d([lon_min, lat_min, lon_max, lat_max], lon_2d, lat_2d, rain_rate)
            wrf_sub = xr.Dataset({
                'RAIN_RATE': xr.DataArray(rain_rate_masked, dims=['Time', 'south_north', 'west_east'])
            })
        else:
            wrf_sub = wrf_ds
        
        spatial_dims = [d for d in wrf_sub['RAIN_RATE'].dims if d != 'Time']
        
        wrf_hourly = wrf_sub['RAIN_RATE'].where(wrf_sub['RAIN_RATE'] > 0, 0).mean(dim=spatial_dims, skipna=True)
        wrf_time = wrf_ds['Time']
        
        ax1.plot(wrf_time, wrf_hourly, linestyle=model_config['linestyle'], 
                 color=model_config['color'], linewidth=2, label=model_config['label'])
        ax2.plot(wrf_time, wrf_accum_ts, linestyle=model_config['linestyle'], 
                 color=model_config['color'], linewidth=2, label=model_config['label'])
    
    # CM1 data
    if cm1_ds is not None:
        lons = cm1_ds['XLONG'].values
        lats = cm1_ds['XLAT'].values
        
        rain_rate = cm1_ds['RAIN_RATE'].values
        rainnc = cm1_ds['RAINNC'].values
        
        rain_rate_masked = select_region_2d([lon_min, lat_min, lon_max, lat_max], lons, lats, rain_rate)
        rainnc_masked = select_region_2d([lon_min, lat_min, lon_max, lat_max], lons, lats, rainnc)
        
        cm1_hourly = np.nanmean(rain_rate_masked, axis=(1, 2))
        cm1_hourly = np.where(cm1_hourly > 0, cm1_hourly, 0)
        
        rainnc_from_start = rainnc_masked - rainnc_masked[0:1, :, :]
        cm1_accum = np.nanmean(rainnc_from_start, axis=(1, 2))
        
        cm1_time = cm1_ds['time'].values
        
        ax1.plot(cm1_time, cm1_hourly, linestyle=CM1_CONFIG['linestyle'], 
                 color=CM1_CONFIG['color'], linewidth=2, label=CM1_CONFIG['label'])
        ax2.plot(cm1_time, cm1_accum, linestyle=CM1_CONFIG['linestyle'], 
                 color=CM1_CONFIG['color'], linewidth=2, label=CM1_CONFIG['label'])
    
    # ICON_KIT data
    if iconkit_ds is not None:
        lons = iconkit_ds['XLONG'].values
        lats = iconkit_ds['XLAT'].values
        
        rain_rate = iconkit_ds['RAIN_RATE'].values
        rainnc = iconkit_ds['RAINNC'].values
        
        rain_rate_masked = select_region_2d([lon_min, lat_min, lon_max, lat_max], lons, lats, rain_rate)
        rainnc_masked = select_region_2d([lon_min, lat_min, lon_max, lat_max], lons, lats, rainnc)
        
        iconkit_hourly = np.nanmean(rain_rate_masked, axis=(1, 2))
        iconkit_hourly = np.where(iconkit_hourly > 0, iconkit_hourly, 0)
        
        rainnc_from_start = rainnc_masked - rainnc_masked[0:1, :, :]
        iconkit_accum = np.nanmean(rainnc_from_start, axis=(1, 2))
        
        iconkit_time = iconkit_ds['time'].values
        
        ax1.plot(iconkit_time, iconkit_hourly, linestyle=ICONKIT_CONFIG['linestyle'], 
                 color=ICONKIT_CONFIG['color'], linewidth=2, label=ICONKIT_CONFIG['label'])
        ax2.plot(iconkit_time, iconkit_accum, linestyle=ICONKIT_CONFIG['linestyle'], 
                 color=ICONKIT_CONFIG['color'], linewidth=2, label=ICONKIT_CONFIG['label'])
    
    # E3SM_THREAD data
    if e3sm_ds is not None:
        lons = e3sm_ds['XLONG'].values
        lats = e3sm_ds['XLAT'].values
        
        rain_rate = e3sm_ds['RAIN_RATE'].values
        rainnc = e3sm_ds['RAINNC'].values
        
        rain_rate_masked = select_region_2d([lon_min, lat_min, lon_max, lat_max], lons, lats, rain_rate)
        rainnc_masked = select_region_2d([lon_min, lat_min, lon_max, lat_max], lons, lats, rainnc)
        
        e3sm_hourly = np.nanmean(rain_rate_masked, axis=(1, 2))
        e3sm_hourly = np.where(e3sm_hourly > 0, e3sm_hourly, 0)
        
        rainnc_from_start = rainnc_masked - rainnc_masked[0:1, :, :]
        e3sm_accum = np.nanmean(rainnc_from_start, axis=(1, 2))
        
        e3sm_time = e3sm_ds['time'].values
        
        ax1.plot(e3sm_time, e3sm_hourly, linestyle=E3SM_CONFIG['linestyle'], 
                 color=E3SM_CONFIG['color'], linewidth=2, label=E3SM_CONFIG['label'])
        ax2.plot(e3sm_time, e3sm_accum, linestyle=E3SM_CONFIG['linestyle'], 
                 color=E3SM_CONFIG['color'], linewidth=2, label=E3SM_CONFIG['label'])
    
    # RAMS data
    if rams_ds is not None:
        lons = rams_ds['XLONG'].values
        lats = rams_ds['XLAT'].values
        
        rain_rate = rams_ds['RAIN_RATE'].values
        rainnc = rams_ds['RAINNC'].values
        
        rain_rate_masked = select_region_2d([lon_min, lat_min, lon_max, lat_max], lons, lats, rain_rate)
        rainnc_masked = select_region_2d([lon_min, lat_min, lon_max, lat_max], lons, lats, rainnc)
        
        rams_hourly = np.nanmean(rain_rate_masked, axis=(1, 2))
        rams_hourly = np.where(rams_hourly > 0, rams_hourly, 0)
        
        rainnc_from_start = rainnc_masked - rainnc_masked[0:1, :, :]
        rams_accum = np.nanmean(rainnc_from_start, axis=(1, 2))
        
        rams_time = rams_ds['time'].values
        
        ax1.plot(rams_time, rams_hourly, linestyle=RAMS_CONFIG['linestyle'], 
                 color=RAMS_CONFIG['color'], linewidth=2, label=RAMS_CONFIG['label'])
        ax2.plot(rams_time, rams_accum, linestyle=RAMS_CONFIG['linestyle'], 
                 color=RAMS_CONFIG['color'], linewidth=2, label=RAMS_CONFIG['label'])
    
    # Format plots
    ax1.set_ylabel('Precip Rate (mm hr$^{-1}$)', fontsize=12)
    ax1.set_title(f'Precipitation Rate Time Series - {config["date_label"]} ({lat_min}°N-{lat_max}°N, {abs(lon_min)}°W-{abs(lon_max)}°W)', 
                  loc='left', fontsize=11)
    ax1.legend(loc='upper right', fontsize=9, ncol=2)
    ax1.set_ylim(config['rate_ylim'])  # Apply case-specific Y-axis limits
    ax1.grid(True, alpha=0.5, linestyle=':')
    ax1.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %HZ"))
    plt.setp(ax1.xaxis.get_majorticklabels(), rotation=45, ha='right')
    
    ax2.set_ylabel('Mean Precip Total (mm)', fontsize=12)
    ax2.set_xlabel('Time', fontsize=12)
    ax2.set_title(f'Accumulated Precipitation Spatial Mean - {config["date_label"]} ({lat_min}°N-{lat_max}°N, {abs(lon_min)}°W-{abs(lon_max)}°W)', 
                  loc='left', fontsize=11)
    ax2.legend(loc='upper left', fontsize=9, ncol=2)
    if 'accum_ylim' in config:
        ax2.set_ylim(config['accum_ylim'])  # Apply case-specific Y-axis limits
    ax2.grid(True, alpha=0.5, linestyle=':')
    ax2.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %HZ"))
    plt.setp(ax2.xaxis.get_majorticklabels(), rotation=45, ha='right')
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        print(f"  Saved: {save_path}")
    plt.close()


# ============================================================================
# MAIN PROCESSING
# ============================================================================
def process_case(case_name, experiment='control'):
    """Process a single case with all models."""
    config = CASE_CONFIG[case_name]
    print("\n" + "="*60)
    print(f"Processing: {config['date_label']} - {experiment}")
    print(f"Large Domain: {LARGE_DOMAIN_BOX['lat_min']}°N-{LARGE_DOMAIN_BOX['lat_max']}°N, "
          f"{abs(LARGE_DOMAIN_BOX['lon_min'])}°W-{abs(LARGE_DOMAIN_BOX['lon_max'])}°W")
    print(f"E3SM_THREAD directory: {config['e3sm_dir']}")
    print("="*60)
    
    case_output_dir = os.path.join(OUTPUT_DIR, case_name)
    os.makedirs(case_output_dir, exist_ok=True)
    
    # Load MRMS data
    print("\nLoading MRMS data...")
    mrms_ds = read_extracted_mrms_data(case_name)
    
    # Load WRF model data
    wrf_datasets = {}
    for model_name in MODEL_CONFIG.keys():
        print(f"\nLoading {model_name} data...")
        wrf_ds = read_extracted_wrf_data(model_name, case_name, experiment)
        wrf_datasets[model_name] = (wrf_ds, MODEL_CONFIG[model_name])
    
    # Load CM1 data
    print("\nLoading CM1 data...")
    cm1_ds = read_extracted_cm1_data(case_name, experiment)
    
    # Load ICON_KIT data
    print("\nLoading ICON_KIT data...")
    iconkit_ds = read_extracted_iconkit_data(case_name, experiment)
    
    # Load E3SM_THREAD data (uses case-specific directory)
    print("\nLoading E3SM_THREAD data...")
    e3sm_ds = read_extracted_e3sm_data(case_name, experiment)
    
    # Load RAMS data
    print("\nLoading RAMS data...")
    rams_ds = read_extracted_rams_data(case_name, experiment)
    
    # Generate time series
    print("\nGenerating multi-model time series (Large Domain)...")
    suffix = f"_{experiment}" if experiment != 'control' else ""
    plot_multimodel_timeseries(mrms_ds, wrf_datasets, cm1_ds, iconkit_ds, e3sm_ds, rams_ds, config,
                               save_path=os.path.join(case_output_dir, f'timeseries_multimodel_largedomain{suffix}.png'))
    
    print(f"\nCase {case_name} ({experiment}) complete!")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Generate multi-model precipitation time series comparison including RAMS (Large Domain: 28.5N-30.45N, 96.1W-94W)'
    )
    parser.add_argument('--case', choices=['june17', 'aug7', 'both'], default='both',
                        help='Case to process (default: both)')
    parser.add_argument('--experiment', choices=['control', 'high_ccn', 'low_ccn'], default='control',
                        help='Experiment to compare (default: control)')
    args = parser.parse_args()
    
    cases = ['june17', 'aug7'] if args.case == 'both' else [args.case]
    
    for case in cases:
        process_case(case, args.experiment)
    
    print("\n" + "="*60)
    print("All processing complete!")
    print("="*60)