import xarray as xr
import numpy as np
import matplotlib.pyplot as plt
import dask

# Directory for input and output files
input_dir = '/nird/projects/6hr/'
output_dir = '/nird/projects/tphwl_6hr/'

# Define the range of years you want to process
years = range(1985, 2015)  # This will include 1985 to 2014
#years = range(2015, 2100)  

for year in years:
    # Load the NetCDF files for the current year
    uas_ds = xr.open_dataset(f'{input_dir}uas_6hr_{year}.nc') #, chunks={'time': 100})
    vas_ds = xr.open_dataset(f'{input_dir}vas_6hr_{year}.nc') 
    tas_ds = xr.open_dataset(f'{input_dir}tas_6hr_{year}.nc')
    ps_ds = xr.open_dataset(f'{input_dir}ps_6hr_{year}.nc')
    huss_ds = xr.open_dataset(f'{input_dir}huss_6hr_{year}.nc')
    rlds_ds = xr.open_dataset(f'{input_dir}rlds_6hr_{year}.nc')
    
    # Ensure both datasets have the same dimensions and coordinates
    assert uas_ds.dims == vas_ds.dims, "Datasets do not have the same dimensions!"

    # Drop the 'bnds' dimension in rlds if necessary
    rlds_ds = rlds_ds.drop_dims('bnds')  # Ensures dimensions match

    # Align datasets
    uas_ds_aligned, vas_ds_aligned, ps_ds_aligned, huss_ds_aligned, rlds_ds_aligned = xr.align(
        uas_ds, vas_ds, ps_ds, huss_ds, rlds_ds, join='override'
    )

    # Calculate wind speed using Pythagorean sum of uas and vas
    wind_data = np.sqrt(uas_ds_aligned['uas'].values**2 + vas_ds_aligned['vas'].values**2)

    # Add wind speed to the base dataset (tas_ds)
    tas_ds['wind'] = (uas_ds_aligned['uas'].dims, wind_data)

    # Optionally, update the attributes for the new wind variable
    tas_ds['wind'].attrs['long_name'] = 'wind speed'
    tas_ds['wind'].attrs['units'] = 'm/s'
    tas_ds['wind'].attrs['_FillValue'] = -1.e+20
    tas_ds['wind'].attrs['standard_name'] = 'wind'
    tas_ds['wind'].attrs['grid_mapping'] = 'rotated_pole'
    tas_ds['wind'].attrs['cell_methods'] = 'time: point time: mean'
    tas_ds['wind'].attrs['coordinates'] = 'height_10m lat lon'
    tas_ds['wind'].attrs['missing_value'] = -1.e+20

    # Add other variables from aligned datasets to tas_ds
    tas_ds['ps'] = ps_ds_aligned['ps']
    tas_ds['huss'] = huss_ds_aligned['huss']
    tas_ds['rlds'] = rlds_ds_aligned['rlds']

    # Save the updated dataset for the current year
    tas_ds.to_netcdf(f'{output_dir}tphwl_6hr_{year}.nc')

print("Processing complete.")