#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Oct  6 08:42:56 2022

@author: idra
"""

# This python script can be used to open a netcdf file, and rotate certain variables by a desired
# angle in a portion of the domain


from netCDF4 import Dataset
import matplotlib.pyplot as plt
from scipy.ndimage.interpolation import rotate
from scipy.ndimage.measurements import center_of_mass
import numpy as np

# Path and name of the file

path_in = "./"

domain = 3
name_in = "geo_em.d0"+str(domain)+".nc"

# Rotation angle (in degree °) 

angle = 28.7

# I want to rotate only part of the domains, so I have to indicate the borderds of the part 

if domain == 1:
    
    # d01 
       
    north = 115 
    south = 100
    west = 50 
    east = 65
    
    Dx = 5
    
elif domain == 2:
    
    # d02
   
    north = 310
    south = 290
    west = 135
    east = 155
    
    Dx = 10

elif domain == 3:
    
    #d03

    north = 500   
    south = 430
    west = 155    
    east = 215

    Dx = 30

# I write the name of the variables to be rotated (For different version of WPS the name and the number of
# the variables of interest can change)

variables = ['LANDMASK','LU_INDEX', 'HGT_M', 'SOILTEMP', 'SOILCTOP', 'SCT_DOM', 'SOILCBOT',
             'SCB_DOM', 'ALBEDO12M', 'GREENFRAC', 'LAI12M', 'SNOALB', 'VAR_SSO', 'LAKE_DEPTH']
   
# Open the file

ncfile = Dataset(path_in+name_in) 

# Read the topography

topo = ncfile.variables['HGT_M']

# Extract only the part of the domain of interest

topo_slice = topo[0, south:north, west:east]

# And calculate the center of mass 

topo_bin = np.where(topo_slice > 1, 1, 0)   
CoM = center_of_mass(topo_bin)

x_point = CoM[1]
y_point = CoM[0]

ncfile.close()

# The rotate function do the rotation with respect to the center of the portion, so I change the borders
# to ensure that the center is coincident with the center of mass

south_com = int(y_point)+south-Dx
north_com = int(y_point)+south+Dx

west_com = int(x_point)+west-Dx
east_com = int(x_point)+west+Dx

south = south_com
north = north_com

west = west_com
east = east_com

for var in variables:   # For every variable
    
    print(var)
    
    # Open the file and extract the variable

    ncfile = Dataset(path_in+name_in)
    
    h = ncfile.variables[var]
    
    if len(h.shape) == 2:
           
        # I am interested only in a part of the domain
    
        h_zoom = h[south:north+1, west:east+1]
           
        # Initialize two different array
        
        h_scipy=h
        
        h_plot = h
        
        
    elif len(h.shape) == 3:
        
        h_zoom = h[0, south:north+1, west:east+1]
        
        h_scipy=h[0]
        
        h_plot = h[0]

    elif len(h.shape) == 4:
        
        h_zoom = h[0,0, south:north+1, west:east+1]
        
        h_scipy=h[0,0]
        
        h_plot = h[0,0]

    # Now the topography will be rotated
    
    if (var == 'HGT_M' or var == 'VAR_SSO'):
        
        rotated = rotate(h_zoom, angle, reshape=False, order=1, prefilter=False, mode='nearest')
    
    else:
        
        rotated = rotate(h_zoom, angle, reshape=False, order=0, prefilter=False, mode='nearest')
    
    # Now I replace the original topography with the rotated one

    i = west        
    
    while i <= east:
        
        j = south   
        
        while j <= north:
            
            h_scipy[j,i] = rotated[j-south,i-west]
                     
            j+=1
        i+=1

    # To check the results the original and rotated variables will be plotted in a folder called "Rotation_plot"
    
    if var == 'HGT_M':
        color = 'terrain'
    else:
        color = 'viridis'       # Valore di default
    
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15,9))
    fig.suptitle(var)
    ax1.contourf(h_plot[south:north,west:east], cmap=color)
    ax1.set_title("Normal")
    ax1.axis('equal')
    ax2.contourf(h_scipy[south:north,west:east], cmap=color)
    ax2.set_title("Rotated of "+str(angle)+"°")
    ax2.axis('equal')
    fig.tight_layout()
    plt.savefig(path_in+"/Rotation_plot/Domain_d0"+str(domain)+"_"+var+".png")
    plt.close(fig)
    
    ncfile.close()
    
    # Now the rotated variables will be written in the original file
    
    ncfile = Dataset(path_in+name_in, 'r+')

    ncfile.variables[var][:]=[h_scipy[:]]

    ncfile.close()