import os, sys
import numpy as np
from netCDF4 import Dataset

#
# read data (1)
#
nc_f   = Dataset('Clim_IDX_2020.nc', 'r')
year1  = nc_f.variables['Year'][:]
month1 = nc_f.variables['Month'][:]
nino34 = nc_f.variables['nino34'][:]
ao     = nc_f.variables['ao'][:]
nc_f.close()

#
# read data (2)
#
nc_f   = Dataset('R_Monthly.nc', 'r')
year2  = nc_f.variables['Year'][:]
month2 = nc_f.variables['Month'][:]
stname = nc_f.variables['Stname'][:]
station= nc_f.variables['Station'][:]

for k in range(len(station)):
    str_out = stname[k,0]+stname[k,1]+stname[k,2]+stname[k,3]+stname[k,4]+stname[k,5]+stname[k,6]+stname[k,7]+stname[k,8]+stname[k,9]+stname[k,10]+stname[k,11]+stname[k,12]+stname[k,13]+stname[k,14]+stname[k,15]+stname[k,16]+stname[k,17]+stname[k,18]+stname[k,19]
    str_out = " ".join(str_out.decode("utf-8").split())
    if str_out == 'HANOI':
        index = k
        print('index='+str(k),str_out)

rain   = nc_f.variables['R'][:][:,:,index]
nc_f.close()

#
# print_out lead=0->3
#
for l in [0,1,2,3]:
    for j in range(len(month2)):

        # print( 'month = '+str(j+1)+', lead = '+str(l)+' \n' )

        sys.stdout=open('data/HANOI_month='+str(int(month2[j]))+'_lead='+str(l)+'.txt','w')

        print( 'month year nino34 ao rainfall' )
        # y
        # x
        k = -1
        for i in range(len(year1)):
            if year1[i] >= 1961 and year1[i] <= 2014:
                k += 1
                li = i
                if l == 0:
                    lj = j
                else:
                    check = int(month2[j]) - l
                    if check > 0:
                        lj = check - 1
                    else:
                        lj = check + 12 - 1
                        li = i - 1
                
                print( int(month2[j]),int(year2[k]),nino34[lj,li],ao[lj,li],round(rain[j,k],2) )

        sys.stdout.close()