#########################################################################################
#                                Lamb's WTs with ERA-Interim                            #
#########################################################################################
# run it from the bash with:
#
# Rscript WT_v2.R 1984
#
# being 1984 the year you want to classify the WTs;
# in this way you can run up to 8 jobs at the same time, each one producing its output files!
# You can also run it for a sequence of years with the syntax:
# 
# Rscript WT_vs.R 1980 2014
#
# and it will compute each year from 1980 to 2014, each one after finishing the previous one.
# If you want to run many years in parallel with just only 1 command, run from the bash:
#
# for y in {1980..2014}; do Rscript WT_v2.R &; done
#
# but it'd need 24 processors! In practice, better to run only 4-8 years at time


library(s2dverification)  # for the function Load()
source('/home/Earth/ncortesi/Downloads/scripts/Rfunctions.R') # for the calendar functions
Load.path <- '/home/Earth/ncortesi/Downloads/scripts/IC3.conf'      # Load() file path

workdir="/scratch/Earth/ncortesi/RESILIENCE/WT"    # working dir where to put the output maps and files
subdatadir="/scratch/Earth/ncortesi/RESILIENCE/WT_subdata"    # dir where to put the intermediate daily psl data (1 file for each subarray)

mslp.rean='ERAintDailyHighRes'  #'ERAintDailyLowRes'  # daily reanalysis dataset used for MSLP data

year.start=1980     # starting year of the MSLP daily data (from the 1st of january)
year.end=1980       # ending year of the MSLP daily data (up to the 31 of December)

low.res.size=5      # even numb. of MSLP nearby grid points to compute the psl at each of the 16 points (both for lon and lat: actual n. of pts used is its square)

##########################################################################################

args <- commandArgs(TRUE)

if(length(args) == 1) year.start <- year.end <- as.integer(args[1])
if(length(args) == 2) {year.start <- as.integer(args[1]); year.end <- as.integer(args[2])}

WTs.type<-c("NE","E","SE","S","SW","W","NW","N","C","C.NE","C.E","C.SE","C.S","C.SW","C.W","C.NW","C.N","A","A.NE","A.E","A.SE","A.S","A.SW","A.W","A.NW","A.N")
radius <- (low.res.size-1)/2  # number of grid points to use as search radius for averaging the psl at each Lamb grid point

# Load 1 day of MSLP data from the reanalysis chosen, just to detect the number of latitude and longitude points:
MSLP <- Load('psl', NULL, mslp.rean, paste0(year.start,'0101'), storefreq = 'daily', leadtimemax = 1, output = 'lonlat', configfile = Load.path)

n.lat <- length(MSLP$lat)  # number of latitude values
n.lon <- length(MSLP$lon)  # number of longitude values

lon.pos <- ifelse(min(MSLP$lon>=0), TRUE, FALSE)  # set lon.pos to TRUE if the longitude values of the reanalysis grid has no negative values, FALSE if not.

# for each grid point of the reanalysis, compute the Lamb WT classification with Lamb grid centered on that point.
# we exclude only central points > +77 degrees lat and < -77 deg. because the Lamb grid needs 12.5 deg. north/south of the central point.
n.lat.unused.poles <- 20 # number of unused latitude values near each of the two poles (must exclude last 10 degrees N/S, so it depends on the spatial res. of the reanalysis)
#n.lat.unused.equat <- 20 # number of unused latitude values near each side of the equator (must exclude 10 degrees N/S for side, so it depends on the sp.res. of the reanalysis)

#pos.lat.eq.north <- tail(which(MSLP$lat >= 0),1) # note than the eventual point at lat=0 is always included
#pos.lat.eq.south <- head(which(MSLP$lat < 0),1) # note than the eventual point at lat=0 is always excluded
#pos.lat.unused.eq.north <- (pos.lat.eq.north-n.lat.unused.poles+1):pos.lat.eq.north
#pos.lat.unused.eq.south <- pos.lat.eq.south:(pos.lat.eq.south+n.lat.unused.equat-1)

#lat.used <- MSLP$lat[-c(1:n.lat.unused.poles,(n.lat-n.lat.unused.poles):n.lat, pos.lat.unused.eq.north, pos.lat.unused.eq.south)]  # latitude values used as central points
lat.used <- round(MSLP$lat[-c(1:n.lat.unused.poles,(n.lat-n.lat.unused.poles):n.lat)],3) # latitude values used as central points rounded to the third decimal
n.lat.used <- length(lat.used)
    
lon.used <- round(MSLP$lon,3) # longitude values used as central points rounded to the third decimal (precision: ~100 m)
n.lon.used <- length(lon.used)

n.grid.points <- length(lat.used)*length(lon.used)

#n.days.tot <- n.days.in.a.yearly.period(year.start,year.end)
#MSLP.year <- Load('psl', NULL, mslp.rean, paste0(year.start,'0101'), storefreq = 'daily', leadtimemax = n.days.in.a.year(y), output = 'lonlat',  configfile = Load.path)
#gc()

## sub <- split.array.big(dimensions=c(n.days, n.lat.used, n.lon.used), along=3)
## for(y in year.start:year.end){
##    cat(paste0("Splitting year: ",y))
   
##    MSLP.year <- Load('psl', NULL, mslp.rean, paste0(year.start,'0101'), storefreq = 'daily', leadtimemax = n.days.in.a.year(y), output = 'lonlat',  configfile = Load.path)

##    for(s in 1:sub$n.sub){
##       MSLP.year.sub <- MSLP.year$obs[1,1,1,,,sub$int[[s]]]
##       save(MSLP.year.sub, file=paste0(subdatadir,"/","MSLP_",mslp.rean,"_year_",y,"_sub_",s,".RData"), compress=FALSE)
##    }
## }
## ### lonmin=lon.used[head(sub$int[[s]],1)], lonmax=lon.used[tail(sub$int[[s]],1)])

# load all MSLP data one year at time:
for(y in year.start:year.end){
    print(paste0("Year: ",y))
    n.days <- n.days.in.a.year(y)
    MSLP.year <- Load('psl', NULL, mslp.rean, paste0(y,'0101'), storefreq = 'daily', leadtimemax = n.days, output = 'lonlat',  configfile = Load.path)
    gc()

    # create a subdir where to put the yearly output (there are too many files to put them in only 1 dir)
    if(!dir.exists(file.path(workdir,paste0(y,"_Rdata")))) dir.create(file.path(workdir,paste0(y,"_Rdata"))) # create a subdir where to put the yerly output (there are too many files to put them in only 1 dir)

    if(!dir.exists(file.path(workdir,paste0(y,"_txt")))) dir.create(file.path(workdir,paste0(y,"_txt"))) # create a subdir where to put the yerly output (there are too many files to put them in only 1 dir)

    # loop over the lat/lon of the central point of the Lamb classification:
    for(latc in lat.used){
        #latc <- lat.used[1]; lonc <- lon.used[1];  # for the debug

        # first, computes the latitude coefficients:
        SF1<-1/cos(latc*pi/180) 
        ZS1<-1/(2*cos(latc*pi/180)^2)
        ZW1<-sin(latc*pi/180)/sin((latc-5)*pi/180)
        ZW2<-sin(latc*pi/180)/sin((latc+5)*pi/180)
        
        for(lonc in lon.used){
            pos.latc <- which(lat.used==latc)
            pos.lonc <- which(lon.used==lonc)
            cat(paste0("Grid point saved: ",(pos.latc-1)*n.lon.used+pos.lonc ,"/", n.grid.points, " "), "\r")
            
            # assign the lat and lon values of the 16 grid points of the Lamb grid centered over the central point:
            lat <- lon <- c()  # lat and lon of the 16 points of Lamb's grid
            lat[1] <- lat[2] <- latc + 10
            lat[3] <- lat[4] <- lat[5] <- lat[6] <- latc + 5
            lat[7] <- lat[8] <- lat[9] <- lat[10] <- latc
            lat[11] <- lat[12] <- lat[13] <- lat[14] <- latc - 5
            lat[15] <- lat[16] <- latc - 10

            lon[1] <- lon[4] <- lon[8] <- lon[12] <- lon[15] <- lonc - 5
            lon[3] <- lon[7] <- lon[11] <- lonc - 15
            lon[2] <- lon[5] <- lon[9] <- lon[13] <- lon[16] <- lonc + 5
            lon[6] <- lon[10] <- lon[14] <- lonc + 15

            # longitude correction for reanalysis with positive-only longitude (to stay always positive):
            if(lon.pos){
                if(lonc - 5 < 0){
                    lon[1] <- lon[4] <- lon[8] <- lon[12] <- lon[15] <- lonc - 5 + 360
                    lon[3] <- lon[7] <- lon[11] <- lonc - 15 + 360
                }
                if(lonc - 15 < 0 && lon - 5 >= 0){
                    lon[3] <- lon[7] <- lon[11] <- lonc - 15 + 360
                }
                if(lonc + 5 >= 360){
                    lon[2] <- lon[5] <- lon[9] <- lon[13] <- lon[16] <- lonc + 5 - 360
                    lon[6] <- lon[10] <- lon[14] <- lonc + 15 - 360
                }
                if(lonc + 15 >= 360 && lonc + 5 < 360){
                    lon[6] <- lon[10] <- lon[14] <- lonc + 15 - 360
                }
            }

            # longitude correction for reanalysis with [-180, +180] longitude (to stay always in the same interval):
            # MUST STILL UPDATE FOR CASE lonc + 15 >= 360 !!!
            if(!lon.pos){
                if(lonc - 5 < -180){
                    lon[1] <- lon[4] <- lon[8] <- lon[12] <- lon[15] <- lonc - 5 + 360
                    lon[3] <- lon[7] <- lon[11] <- lonc - 15 + 360
                }
                if(lonc + 5 >= 180){
                    lon[2] <- lon[5] <- lon[9] <- lon[13] <- lon[16] <- lonc + 5 - 360
                    lon[6] <- lon[10] <- lon[14] <- lonc + 15 - 360
                }
            }                
            
            # lat and lon position of the 16 Lamb grid points inside the MSLP reanalysis:
            pos.lat <- pos.lon <- c()
            for(p in 1:16){
                pos <- nearest(lat[p], lon[p], MSLP$lat, MSLP$lon)
                pos.lat[p] <- pos[1]
                pos.lon[p] <- pos[2]
            }

            #WTs <- data.frame()
            
            # mean psl at each of the 16 grid points:
            psl <- array(NA,c(n.days,16)) # for each day of the year, MSLP of the 16 points of Lamb's grid
            #for(p in 1:16) psl[,p] <- MSLP.year$obs[1,1,1,,pos.lat[p],pos.lon[p]]
            
            for(p in 1:16) {
                 
                 pos.lon.low.res <- (pos.lon[p]-radius):(pos.lon[p]+radius)
                
                 if(head(pos.lon.low.res,1) <= 0){
                     ss <- which(pos.lon.low.res <= 0)
                     pos.lon.low.res <- c(pos.lon.low.res[ss] + n.lon.used, pos.lon.low.res[-ss])
                 }
                 if(tail(pos.lon.low.res,1) > n.lon.used){
                     ss <- which(pos.lon.low.res > n.lon.used)
                     pos.lon.low.res <- c(pos.lon.low.res[-ss], pos.lon.low.res[ss]-n.lon.used)
                 }
                
                 pos.lat.low.res <- (pos.lat[p]-radius):(pos.lat[p]+radius)  # it never has to be corrected because we excluded the poles

                 MSLP.low.res <- MSLP.year$obs[1,1,1,,pos.lat.low.res,pos.lon.low.res]
                 psl[,p] <- apply(MSLP.low.res, 1, mean)
            }

            # compute geostrophical indexes (don't round them to the 3rd decimal number as done by Trigo and Martin-Vide, not to introduce an error that change 3-4% of WTs!!!):
            SF<-WF<-F<-D<-D2<-c()

            WF <- 0.5*(psl[,12]+psl[,13]) - 0.5*(psl[,4]+psl[,5])
            SF <- SF1*(0.25 * (psl[,5]+2*psl[,9]+psl[,13]) - 0.25 * (psl[,4]+2*psl[,8]+psl[,12]))
            F <- (SF*SF + WF*WF)^0.5

            if(latc < 0){
                WF <- -WF  # because in the southern hemisphere, cyclones and anticyclones spins in the opposite sense!
                SF <- -SF
            }

            # D e' la direzione del flusso; dato che l'angolo risultato di atan e' espresso in radianti, bisogna convertirlo in gradi moltiplicandolo per (180/pi)
            # D e' l'angolo tra la direzione del vento e la sua proiezione in direzione nord-sud (SF)
            D <- (180/pi) * atan(WF/SF)

            # trasforma l'angolo D iniziale in un angolo D2 sempre positivo che si calcola a partire del Nord andando in senso orario:
            Q1 <- which(WF>0 & SF>0)
            Q2 <- which(WF>0 & SF<=0)
            Q3 <- which(WF<0 & SF<=0)   # calcola quali angoli D appartengono al quadrante 1 (Q1), quali al quadrante 2, quali al 3 e quali al 4.
            Q4 <- which(WF<0 & SF>0)    # il primo quadrante e' in alto a destra, il secondo in basso a destra, il terzo in basso a sinistra, il quarto in alto a sinistra.

            # the wind direction is usually different from the angle D where the wind is blowing to, so a correction has to be made:
            D2 <- D
            D2[Q1] <- D[Q1] + 180  
            D2[Q2] <- D[Q2] + 360
            D2[Q4] <- D[Q4] + 180
            
            # individua i WTs direzionali:
            NE <- which(D2>=22.5 & D2<67.5)
            E <- which(D2>=67.5 & D2<112.5)
            SE <- which(D2>=112.5 & D2<157.5)
            S <- which(D2>=157.5 & D2<202.5)
            SW <- which(D2>=202.5 & D2<247.5)
            W <- which(D2>=247.5 & D2<292.5)
            NW <- which(D2>=292.5 & D2<337.5)
            N <- which(D2>=337.5 | D2<22.5)

            # aggiunge una colonna a destra con il WT direzionale associato all'angolo:
            D2.dir<-rep(NA,length(D2))
            D2.dir[NE]<-"NE";D2.dir[E]<-"E";D2.dir[SE]<-"SE";D2.dir[S]<-"S";D2.dir[SW]<-"SW";D2.dir[W]<-"W";D2.dir[NW]<-"NW";D2.dir[N]<-"N"

            # calcolo componenti vorticita':
            ZS<-ZW<-Z<-c()
            
            ZS <- ZS1*(0.25 * (psl[,6]+2*psl[,10]+psl[,14]) - 0.25 * (psl[,5]+2*psl[,9]+psl[,13]) - 0.25 * (psl[,4]+2*psl[,8]+psl[,12]) + 0.25 * (psl[,3]+2*psl[,7]+psl[,11]))
            ZW <- ZW1 * (0.5*(psl[,15]+psl[,16]) - 0.5*(psl[,8]+psl[,9])) - ZW2 * (0.5*(psl[,8]+psl[,9]) - 0.5*(psl[,1]+psl[,2]))
            Z <- ZS + ZW

            # Pure and hybrid WTs
            cyc.pure <- which(Z>2*F)
            anticyc.pure <- which(Z<(-2*F))
            hybrid.cyc <- which(Z>0 & abs(Z)<2*F & abs(Z)>F)
            hybrid.anticyc <- which(Z<0 & abs(Z)<2*F & abs(Z)>F)
            indeter <- which(F<6 & abs(Z)<6) # tipo di tempo indeterminato (U); the choice of 6 depend on grid size and should be changed if grid res.is higher!

            # create a data frame with all the info about the WT classification:    

            #years.period <- months.period <- days.period <- c()
            #for(y in year.start:year.end) years.period <- c(years.period, rep(y, n.days.in.a.year(y)))
            #for(y in year.start:year.end) months.period <- c(months.period, seq.months.in.a.year(y))
            #for(y in year.start:year.end) days.period <- c(days.period, seq.days.in.a.year(y))
            years.period <- rep(y, n.days)
            months.period <- seq.months.in.a.year(y)
            days.period <- seq.days.in.a.year(y)
            WT<-data.frame(Year=years.period, Month=months.period, Day=days.period, Direction=round(D2,3), WT.dir=D2.dir, stringsAsFactors=FALSE)

            WT$Year<-as.numeric(WT$Year);WT$Month<-as.numeric(WT$Month);WT$Day<-as.numeric(WT$Day);WT$Direction<-as.numeric(WT$Direction);WT$WT.dir<-as.character(WT$WT.dir)

            # colonna che identifica quali giorni sono di tipo ibrido e quali di tipo direzionale:
            WT$Hyb <- rep("Directional",length(D2))
            WT$Hyb[cyc.pure] <- "Pure_C"
            WT$Hyb[anticyc.pure] <- "Pure_A"
            WT$Hyb[hybrid.cyc] <- "Hybrid_C"
            WT$Hyb[hybrid.anticyc] <- "Hybrid_A"

             # colonna con la lista dei 26 WTs officiali:
            WT$WT <- WT$D2.dir
            WT$WT[cyc.pure] <- "C"
            WT$WT[anticyc.pure] <- "A"
            WT$WT[hybrid.cyc] <- paste("C.",WT$WT.dir[hybrid.cyc],sep="")
            WT$WT[hybrid.anticyc] <- paste("A.", WT$WT.dir[hybrid.anticyc],sep="")
            #WT$WT[indeter] <- "U"   # comment this line not to include the Unclassified U WT

            WT$WT.num <- match(WT$WT,WTs.type)# aggiunge una colonna con la conversione dei 26 tipi di tempo a numeri da 1 a 26

            # colonna con la lista dei 10 WTs ottenuti unendo a C e ad A gli otto WT direzionali puri;
            # i WTs ibridi si aggiungono ognuno al suo tipo direzionale corrispondente (es: C.SW si aggiunge a SW, ecc.)
            WT$WT10 <- WT$WT.dir
            WT$WT10[cyc.pure] <- "C"
            WT$WT10[anticyc.pure] <- "A"
            WT$WT10.num <- match(WT$WT10,WTs.type) # conversione dei 10 tipi di tempo a numeri da 1 a 9 + il tipo A che si converte nel numero 18

            # add the other geostrophical indexes:
            WT$SF=round(SF,5)
            WT$WF=round(WF,5)
            WT$F=round(F,5)
            WT$ZS=round(ZS,5)
            WT$ZW=round(ZW,5)
            WT$Z=round(Z,5)

            #WTs<-rbind(WTs,WT)
            
            # save the full data table as a .txt to exchange it with other people: 
            write.table(WT,file=paste0(workdir,"/",y,"_txt/WTs_",mslp.rean,"_year_",y,"_lat_",latc,"_lon_",lonc,".txt"),row.names=FALSE, col.names=TRUE, quote=FALSE, sep=" ")
            # save only the classification with 10 WTs, to have a smaller (binary) file:
            #WTs<-cbind(WT$Year, WT$Month, WT$Day, WT$WT10.num)
            WTs <- WT$WT10.num
            save(WTs, file=paste0(workdir,"/",y,"_Rdata/10WTs_",mslp.rean,"_year_",y,"_lat_",latc,"_lon_",lonc,".RData"), compress=FALSE)
            
            rm(WT,SF,WF,F,D,ZS,ZW,W)
            
        } # close for on lonc
    } # close for on latc
    cat("\n")
    
} # close for on y 
