#!/usr/bin/env python                                                                                       \
                                                           
#We would like to rotate the wind by using a skript in python. First it is important to double check the unit of angle. cos and sin in Python consider angle in radians

                                                  
# -*- coding: utf-8 -*-                                                                                    
import netCDF4
import numpy as np
import os
import sys
import datetime
import calendar
import xarray as xr

#import netCDF4, sys, datetime, calendar, os


def main():
    """                                                                                                      
    testing era5 vector rotation                                                                            
                                                                                                             
    by Hiroshi Sumata / 2023.06.09                                                                          
    """

    #    path = "/cluster/work/users/andriani/metroms_run/arctic-4km/era5_A4_intrp/"
    fname_angle = '/cluster/work/users/andriani/metroms_run/arctic-4km/era5_A4_intrp/arctic4km_grd.nc'
    nc2 = netCDF4.Dataset(fname_angle, 'r')
    #nc2 = xr.open_dataset(fname_angle) #The open_dataset method will open the fname_angle but not read the entire file into memory #XARRAY
    angle = nc2.variables['angle'][:]
    #angle = nc2.angle.values #XARRAY
    nc2.close()


    fname = '/cluster/work/users/andriani/metroms_run/arctic-4km/era5_A4_intrp/era5_2008_11.nc'
    #    fname_angle = '/cluster/work/users/andriani/metroms_run/arctic-4km/era5_A4_intrp/arctic4km_grd.nc'

    nc = netCDF4.Dataset(fname, mode = 'r+', clobber = True)
    #nc = xr.open_dataset(fname) #XARRAY
    Uwind = nc.variables['Uwind'][:]
    Vwind = nc.variables['Vwind'][:]
    #Uwind = nc.Uwind.values #XARRAY
    #Vwind = nc.Vwind.values #XARRAY


    #    nc2 = netCDF4.Dataset(fname_angle, 'r')
    #    angle = nc2.variables['angle'][:]
    #    angle_units = nc2.variables['angle'].units
    #    nc2.close()

    #    if angle_units == 'radians':
    #        print('It is in radians.It does not need to be converted')
    #    else:
    #        print('It needs to be converted from degrees to radians')
    #        angle = angle[:,:]*(np.pi/180)


    u =np.copy(Uwind)
    v =np.copy(Vwind)

    for n in range(Uwind.shape[0]):
        u[n, :, :] = Uwind[n, :, :] * np.cos(angle[:, :]) + Vwind[n, :, :] * np.sin(angle[:, :])
        v[n, :, :] = Vwind[n, :, :] * np.cos(angle[:, :]) - Uwind[n, :, :] * np.sin(angle[:, :])


    nc.variables['Uwind'][:] = u[:]
    nc.variables['Vwind'][:] = v[:]

    ################
    #Try to use a loop since the datasets are in 1TB. Save them again in one time slot each time
    #for n in range(Uwind.shape[0]):
       # Uwind[n,:,:] = u[n,:,:]
       # Vwind[n,:,:] = v[n,:,:]

    #nc.to_netcdf('cluster/work/users/andriani/metroms_run/arctic-4km/era5_A4_intrp/modified_era5_2009_1.nc')

    nc.close()

    print ('finished saving')

if __name__ == '__main__':
    main()