#!/usr/bin/env python3
"""
WRF-DRI (WRF-MorrisonAero) Precipitation Extraction Script
Extracts RAINNC, XLAT, XLONG from WRF-DRI output files and saves to consolidated NetCDF files.

Model: WRF V4.6.1 with Morrison aerosol-aware microphysics (MP_PHYSICS=40)
Source: Desert Research Institute (DRI)
Output frequency: 10 min, 145 files per case, 24-hour simulations starting at 06 UTC.
"""
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"

# WRF-DRI source uses cntl / 3xhigh / 3xlow naming on disk; we map to the
# project-standard control / high_ccn / low_ccn keys for output consistency.
MODEL_CONFIG = {
    'WRF_DRI': {
        'june17': {
            'control':  '/global/cfs/projectdirs/m4486/Haochen/TRACER_Model/WRF-MorrisonAero/Jun17/cntl',
            'high_ccn': '/global/cfs/projectdirs/m4486/Haochen/TRACER_Model/WRF-MorrisonAero/Jun17/3xhigh',
            'low_ccn':  '/global/cfs/projectdirs/m4486/Haochen/TRACER_Model/WRF-MorrisonAero/Jun17/3xlow',
        },
        'aug7': {
            'control':  '/global/cfs/projectdirs/m4486/Haochen/TRACER_Model/WRF-MorrisonAero/Aug7/cntl',
            'high_ccn': '/global/cfs/projectdirs/m4486/Haochen/TRACER_Model/WRF-MorrisonAero/Aug7/3xhigh',
            'low_ccn':  '/global/cfs/projectdirs/m4486/Haochen/TRACER_Model/WRF-MorrisonAero/Aug7/3xlow',
        },
        'output_dir': os.path.join(OUTPUT_BASE, 'WRF_DRI'),
    },
}

# 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."""
    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 []

    all_files = sorted(list(set(all_files)))

    start_dt = case_config['start_datetime']
    end_dt = case_config['end_datetime']

    filtered_files = []
    for fpath in all_files:
        try:
            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:
            continue

    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.
    """
    if not file_list:
        print("  No files to process")
        return False

    print("  Processing {0} files...".format(len(file_list)))

    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)

            if 'RAINNC' in ds:
                rainnc = ds['RAINNC']
                if 'Time' in rainnc.dims:
                    rainnc = rainnc.isel(Time=0)
                rainnc_values = np.array(rainnc.values)
                if rainnc_values.ndim > 2:
                    rainnc_values = rainnc_values.squeeze()
                rainnc_list.append(rainnc_values)
                time_list.append(timestamp)

            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()

            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

    print("  Creating output dataset with {0} timesteps...".format(len(rainnc_list)))

    rainnc_array = np.stack(rainnc_list, axis=0)
    time_array = np.array(time_list, dtype='datetime64[ns]')

    print("  RAINNC shape: {0}".format(rainnc_array.shape))
    print("  XLAT shape: {0}".format(xlat.shape))
    print("  XLONG shape: {0}".format(xlong.shape))

    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-DRI (WRF-MorrisonAero) precipitation data',
            'source': 'WRF-MorrisonAero output files (DRI)',
            'microphysics': 'Morrison aerosol-aware (mp_physics=40)',
            'model': 'WRF_DRI',
            'created': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
        }
    )

    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("  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]

    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]

    output_dir = model_config['output_dir']
    os.makedirs(output_dir, exist_ok=True)

    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))

    if os.path.exists(output_file):
        print("  Output file already exists. Skipping...")
        print("  (Delete the file to re-extract)")
        return True

    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

    success = extract_precipitation(file_list, output_file)
    return success


# ============================================================================
# MAIN
# ============================================================================
def main():
    parser = argparse.ArgumentParser(
        description='Extract WRF-DRI precipitation data and save to NetCDF'
    )
    parser.add_argument('--model', choices=['WRF_DRI', '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()

    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-DRI 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))

    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

    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()