# Reconstructing ORCHIDEE's forcing files as matrices for a gien area
# L. Fita, CIMA. December 2017
#
import numpy as np
from netCDF4 import Dataset as NetCDFFile
import os
import re
import numpy.ma as ma
# Importing generic tools file 'generic_tools.py'
import generic_tools as gen
import nc_var_tools as ncvar
import subprocess as sub
import module_ForSci as Sci
from optparse import OptionParser

parser = OptionParser()
parser.add_option("-y", "--year", dest="year", help="year to process",               \
  metavar="VALUE")

(opts, args) = parser.parse_args()

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

filen = 'cruncep_halfdeg_' + opts.year + '.nc'

# Variable whcih provides the indices of a 1D vector from the dimy, dimx space
indvar = 'land'

# 2D longitude, latitude matrices
lonvar = 'nav_lon'
latvar = 'nav_lat'

# Range to retrieve
Xmin=-90.25
#Xmin='all'
Xmax=-33.25
Ymin=-67.25
Ymax=15.25

# Variables to reconstruct
#variable = 'all'
variable = 'time,LWdown,Qair,PSurf,SWdown,Tair,Wind_E,Wind_N,Snowf,Rainf'

# Resolution
resX = 0.5
resY = 0.5

# Minimum difference from matrix to localized point
maxdiff = 0.05

# Projection of the matrix
matProj = 'latlon'

#######    #######
## MAIN
    #######
main = 'ORforcing_reconstruct.py'
fname = 'ORforcing_reconstruct.py'

onc = NetCDFFile(filen, 'r')

availProj = ['latlon']

ofilen = 'reconstruct_matrix_' + opts.year + '.nc'

ncvars = onc.variables.keys()

varstocheck = [indvar, lonvar, latvar]
for vn in varstocheck:
    if not gen.searchInlist(ncvars, vn):
        print gen.errormsg
        print '  ' + main + ": file '" + filen + "' does not have variable '" + vn + \
          "' !!"
        print '    available ones:', ncvars
        quit(-1)

oind = onc.variables[indvar]
olon = onc.variables[lonvar]
olat = onc.variables[latvar]

vecdimn = oind.dimensions[0]
Xdimn = olon.dimensions[1]
Ydimn = olon.dimensions[0]

# Fortran indices, first 1
indv = oind[:]
lonv = olon[:]
latv = olat[:]
indv = indv - 1

dimx = lonv.shape[1]
dimy = lonv.shape[0]

veclon = np.zeros((len(indv)), dtype=np.float)
veclat = np.zeros((len(indv)), dtype=np.float)

# Construct veclon, veclat
for iid in range(len(indv)):
    iy = int((indv[iid]-1)/dimx)
    ix = indv[iid] - iy*dimx - 1
    veclon[iid] = lonv[iy,ix]
    veclat[iid] = latv[iy,ix]

if not gen.searchInlist(availProj, matProj):
    print errormsg
    print '  ' + fname + ": projection '" + matProj + "' not available !!"
    print '    available ones:', availProj
    quit(-1)

ncdims = onc.dimensions

if variable == 'all':
    varns = ncvars
else:
    varns = variable.split(',')
    for vn in varns:
        if not gen.searchInlist(ncvars, vn):
            print gen.errormsg
            print '  ' + fname + ": file '" + ncfile + "' does not have " +          \
              " variable '" + vn + "' !!"
            print '  available ones:', ncvars
            quit(-1)

if gen.searchInlist(olon.ncattrs(), 'units'):
    xunits = olon.getncattr('units')
else:
    xunits = '-'
if gen.searchInlist(olat.ncattrs(), 'units'):
    yunits = olat.getncattr('units')
else:
   yunits = '-'

if type(Xmin) == type('2') and Xmin == 'all':
    Xmin = np.min(lonv)
    Xmax = np.max(lonv)
    Ymin = np.min(latv)
    Ymax = np.max(latv)

# Matrix values
if matProj == 'latlon':
    dimx = int((Xmax - Xmin+resX)/resX)
    dimy = int((Ymax - Ymin+resY)/resY)

matindt, matXt, matYt = Sci.module_scientific.reconstruct_matrix(vectorxpos=veclon,  \
  vectorypos=veclat, dvec=veclon.shape[0], xmin=Xmin, xmax=Xmax, ymin=Ymin,ymax=Ymax,\
  dmatx=dimx, dmaty=dimy, matproj=matProj, maxdiff=maxdiff)

matind = matindt.transpose()
# Fortran like, First 1
matind = np.where(matind != -1, matind - 1, matind)

# Creation of file
onewnc = NetCDFFile(ofilen, 'w')

# Dimensions
newdim = onewnc.createDimension('x', dimx)
newdim = onewnc.createDimension('y', dimy)

# Variable-dimension
newvar = onewnc.createVariable('lon', 'f8', ('y', 'x'))
newvar[:] = matXt.transpose()
ncvar.basicvardef(newvar, 'lon', 'Longitude', 'degrees_east')   

newvar = onewnc.createVariable('lat', 'f8', ('y', 'x'))
newvar[:] = matYt.transpose()
ncvar.basicvardef(newvar, 'lat', 'Latitude', 'degrees_north')   
onewnc.sync()

# Variable indices
newvar = onewnc.createVariable('vec1D_matind', 'i', ('y', 'x'), fill_value=-1)
newvar[:] = matind
ncvar.basicvardef(newvar, 'vec1D_matind', 'matrix with the equivalencies from 1D ' + \
  'vector indices', '-')
ncvar.set_attribute(newvar, 'coordinates', 'lon lat')

# Looking for equivalencies in the 1D vector
matlonlat = matind.copy()
for j in range(dimy):
    for i in range(dimx):
        if matind[j,i] != -1:
            matlonlat[j,i] = indv[matind[j,i]]

newvar = onewnc.createVariable('lonlat_matind', 'i', ('y', 'x'), fill_value=-1)
newvar[:] = matlonlat
ncvar.basicvardef(newvar, 'lonlat_matind', 'matrix with the equivalencies from ' +   \
  '2D lon, lat matrices', '-')
ncvar.set_attribute(newvar, 'coordinates', 'lon lat')
onewnc.sync()

# Getting variables
for vn in varns:
    if not onewnc.variables.has_key(vn):
        ovar = onc.variables[vn]
        indn = ovar.dimensions
        vardims = []
        shapevar = []
        for dn in indn:
            if not gen.searchInlist(onewnc.dimensions, dn) and dn != Xdimn and   \
              dn != Ydimn:
                if onc.dimensions[dn].isunlimited():
                    newdim = onewnc.createDimension(dn, None)
                else:
                    newdim = onewnc.createDimension(dn, len(onc.dimensions[dn]))
                
            if dn == vecdimn: 
                vardims.append('y')
                vardims.append('x')
                shapevar.append(dimy)
                shapevar.append(dimx)
            else: 
                vardims.append(dn)
                shapevar.append(len(onc.dimensions[dn]))

        if ovar.dtype == type(int(2)):
            newvar= onewnc.createVariable(vn,ncvar.nctype(ovar.dtype),tuple(vardims),\
              fill_value=gen.fillValueI)
            varvals = np.ones(tuple(shapevar), dtype=ovar.dtype)*gen.fillValueI
        elif ovar.dtype == type(np.int32(2)):
            newvar= onewnc.createVariable(vn,ncvar.nctype(ovar.dtype),tuple(vardims),\
              fill_value=gen.fillValueI)
            varvals = np.ones(tuple(shapevar), dtype=ovar.dtype)*gen.fillValueI
        elif ovar.dtype == type(np.int64(2)):
            newvar= onewnc.createVariable(vn,ncvar.nctype(ovar.dtype),tuple(vardims),\
              fill_value=gen.fillValueI)
            varvals = np.ones(tuple(shapevar), dtype=ovar.dtype)*gen.fillValueI
        elif ovar.dtype == type(np.float(2.)):
            newvar= onewnc.createVariable(vn,ncvar.nctype(ovar.dtype),tuple(vardims),\
              fill_value=gen.fillValueF)
            varvals = np.ones(tuple(shapevar), dtype=ovar.dtype)*gen.fillValueF
        elif ovar.dtype == type(np.float32(2.)):
            newvar= onewnc.createVariable(vn,ncvar.nctype(ovar.dtype),tuple(vardims),\
              fill_value=gen.fillValueF)
            varvals = np.ones(tuple(shapevar), dtype=ovar.dtype)*gen.fillValueF
        else:
            print gen.errormsg
            print '  ' + fname + ': variable type:', ovar.dtype, ' not ready !!'
            quit(-1)

        print '  reconstructing:', vn, '...'
        # Filling variable. It would be faster if we can avoid this loop... I'm feeling lazy!
        if not gen.searchInlist(vardims,'x') and not gen.searchInlist(vardims,'y'):
            newvar[:] = ovar[:] 
        else:
            ovart = ovar[:].transpose()
            if newvar.dtype == type(float(2.)) or newvar.dtype == type(np.float(2.)) \
              or newvar.dtype == type(np.float32(2)) or                              \
              newvar.dtype == type(np.float64(2)):
                newvals = Sci.module_scientific.fill3dr_2dvec(matind=matindt,        \
                  inmat=ovart, id1=ovart.shape[0], id2=ovart.shape[1],               \
                  od1=newvar.shape[2], od2=newvar.shape[1], od3=newvar.shape[0])
            else:
                newvals = Sci.module_scientific.fill3di_2dvec(matind=matindt,        \
                  inmat=ovart, id1=ovart.shape[0], id2=ovart.shape[1],               \
                  od1=newvar.shape[2], od2=newvar.shape[1], od3=newvar.shape[0])
            newvar[:] = newvals.transpose()

        # Attributes
        for atn in ovar.ncattrs():
            if atn != '_FillValue' and atn != 'units':
                atv = ovar.getncattr(atn)
                ncvar.set_attribute(newvar, atn, atv)
        ncvar.set_attribute(newvar, 'coordinates', 'lon lat')
        onewnc.sync()
   
# Global attributes
for atn in onc.ncattrs():
    atv = ovar.getncattr(atn)
    ncvar.set_attribute(onewnc, atn, atv)
onewnc.sync()
ncvar.add_global_PyNCplot(onewnc, main, fname, '0.1')
onc.close()

# Reconstructing times
otime = onewnc.variables['time']
ncvar.set_attribute(otime, 'units', 'seconds since ' + opts.year + '-01-01 00:00:00')

onewnc.sync()
onewnc.close()

print fname + ": Successful writing of file '" + ofilen + ".nc' !!"