# L. Fita, LMD-Jussieu. February 2015
## e.g. # validation_sim.py -d X@west_east@None,Y@south_north@None,T@Time@time -D X@XLONG@longitude,Y@XLAT@latitude,T@time@time -k single-station -l 4.878773,43.915876,12. -o /home/lluis/DATA/obs/HyMeX/IOP15/sfcEnergyBalance_Avignon/OBSnetcdf.nc -s /home/lluis/PY/wrfout_d01_2012-10-18_00:00:00.tests -v HFX@H 
import numpy as np
import os
import re
from optparse import OptionParser
from netCDF4 import Dataset as NetCDFFile

main = 'validarion_sim.py'
errormsg = 'ERROR -- errror -- ERROR -- error'
warnmsg = 'WARNING -- warning -- WARNING -- warning'

# version
version=1.0

# Filling values for floats, integer and string
fillValueF = 1.e20
fillValueI = -99999
fillValueS = '---'

def index_3mat(matA,matB,matC,val):
    """ Function to provide the coordinates of a given value inside three matrix simultaneously
    index_mat(matA,matB,matC,val)
      matA= matrix with one set of values
      matB= matrix with the other set of values
      matB= matrix with the third set of values
      val= triplet of values to search
    >>> index_mat(np.arange(27).reshape(3,3,3),np.arange(100,127).reshape(3,3,3),np.arange(200,227).reshape(3,3,3),[22,122,222])
    [2 1 1]
    """
    fname = 'index_3mat'

    matAshape = matA.shape
    matBshape = matB.shape
    matCshape = matC.shape

    for idv in range(len(matAshape)):
        if matAshape[idv] != matBshape[idv]:
            print errormsg
            print '  ' + fname + ': Dimension',idv,'of matrices A:',matAshape[idv],  \
              'and B:',matBshape[idv],'does not coincide!!'
            quit(-1)
        if matAshape[idv] != matCshape[idv]:
            print errormsg
            print '  ' + fname + ': Dimension',idv,'of matrices A:',matAshape[idv],  \
              'and C:',matCshape[idv],'does not coincide!!'
            quit(-1)

    minA = np.min(matA)
    maxA = np.max(matA)
    minB = np.min(matB)
    maxB = np.max(matB)
    minC = np.min(matC)
    maxC = np.max(matC)

    if val[0] < minA or val[0] > maxA:
        print warnmsg
        print '  ' + fname + ': first value:',val[0],'outside matA range',minA,',',  \
          maxA,'!!'
    if val[1] < minB or val[1] > maxB:
        print warnmsg
        print '  ' + fname + ': second value:',val[1],'outside matB range',minB,',',  \
          maxB,'!!'
    if val[2] < minC or val[2] > maxC:
        print warnmsg
        print '  ' + fname + ': second value:',val[2],'outside matC range',minC,',',  \
          maxC,'!!'

    dist = np.zeros(tuple(matAshape), dtype=np.float)
    dist = np.sqrt((matA - np.float(val[0]))**2 + (matB - np.float(val[1]))**2 +     \
      (matC - np.float(val[2]))**2)

    mindist = np.min(dist)
   
    matlist = list(dist.flatten())
    ifound = matlist.index(mindist)

    Ndims = len(matAshape)
    valpos = np.zeros((Ndims), dtype=int)
    baseprevdims = np.zeros((Ndims), dtype=int)

    for dimid in range(Ndims):
        baseprevdims[dimid] = np.product(matAshape[dimid+1:Ndims])
        if dimid == 0:
            alreadyplaced = 0
        else:
            alreadyplaced = np.sum(baseprevdims[0:dimid]*valpos[0:dimid])
        valpos[dimid] = int((ifound - alreadyplaced )/ baseprevdims[dimid])

    return valpos

def index_2mat(matA,matB,val):
    """ Function to provide the coordinates of a given value inside two matrix simultaneously
    index_mat(matA,matB,val)
      matA= matrix with one set of values
      matB= matrix with the pother set of values
      val= couple of values to search
    >>> index_mat(np.arange(27).reshape(3,3,3),np.arange(100,127).reshape(3,3,3),[22,111])
    [2 1 1]
    """
    fname = 'index_2mat'

    matAshape = matA.shape
    matBshape = matB.shape

    for idv in range(len(matAshape)):
        if matAshape[idv] != matBshape[idv]:
            print errormsg
            print '  ' + fname + ': Dimension',idv,'of matrices A:',matAshape[idv],  \
              'and B:',matBshape[idv],'does not coincide!!'
            quit(-1)

    minA = np.min(matA)
    maxA = np.max(matA)
    minB = np.min(matB)
    maxB = np.max(matB)

    if val[0] < minA or val[0] > maxA:
        print warnmsg
        print '  ' + fname + ': first value:',val[0],'outside matA range',minA,',',  \
          maxA,'!!'
    if val[1] < minB or val[1] > maxB:
        print warnmsg
        print '  ' + fname + ': second value:',val[1],'outside matB range',minB,',',  \
          maxB,'!!'

    dist = np.zeros(tuple(matAshape), dtype=np.float)
    dist = np.sqrt((matA - np.float(val[0]))**2 + (matB - np.float(val[1]))**2)

    mindist = np.min(dist)
   
    matlist = list(dist.flatten())
    ifound = matlist.index(mindist)

    Ndims = len(matAshape)
    valpos = np.zeros((Ndims), dtype=int)
    baseprevdims = np.zeros((Ndims), dtype=int)

    for dimid in range(Ndims):
        baseprevdims[dimid] = np.product(matAshape[dimid+1:Ndims])
        if dimid == 0:
            alreadyplaced = 0
        else:
            alreadyplaced = np.sum(baseprevdims[0:dimid]*valpos[0:dimid])
        valpos[dimid] = int((ifound - alreadyplaced )/ baseprevdims[dimid])

    return valpos

def index_mat(mat,val):
    """ Function to provide the coordinates of a given value inside a matrix
    index_mat(mat,val)
      mat= matrix with values
      val= value to search
    >>> index_mat(np.arange(27).reshape(3,3,3),22)
    [2 1 1]
    """

    fname = 'index_mat'

    matshape = mat.shape

    matlist = list(mat.flatten())
    ifound = matlist.index(val)

    Ndims = len(matshape)
    valpos = np.zeros((Ndims), dtype=int)
    baseprevdims = np.zeros((Ndims), dtype=int)

    for dimid in range(Ndims):
        baseprevdims[dimid] = np.product(matshape[dimid+1:Ndims])
        if dimid == 0:
            alreadyplaced = 0
        else:
            alreadyplaced = np.sum(baseprevdims[0:dimid]*valpos[0:dimid])
        valpos[dimid] = int((ifound - alreadyplaced )/ baseprevdims[dimid])

    return valpos

def coincident_CFtimes(tvalB, tunitA, tunitB):
    """ Function to make coincident times for two different sets of CFtimes
    tvalB= time values B
    tunitA= time units times A to which we want to make coincidence
    tunitB= time units times B
    >>> coincident_CFtimes(np.arange(10),'seconds since 1949-12-01 00:00:00',
      'hours since 1949-12-01 00:00:00')
    [     0.   3600.   7200.  10800.  14400.  18000.  21600.  25200.  28800.  32400.]
    >>> coincident_CFtimes(np.arange(10),'seconds since 1949-12-01 00:00:00',
      'hours since 1979-12-01 00:00:00')
    [  9.46684800e+08   9.46688400e+08   9.46692000e+08   9.46695600e+08
       9.46699200e+08   9.46702800e+08   9.46706400e+08   9.46710000e+08
       9.46713600e+08   9.46717200e+08]
    """
    import datetime as dt
    fname = 'coincident_CFtimes'

    trefA = tunitA.split(' ')[2] + ' ' + tunitA.split(' ')[3]
    trefB = tunitB.split(' ')[2] + ' ' + tunitB.split(' ')[3]
    tuA = tunitA.split(' ')[0]
    tuB = tunitB.split(' ')[0]

    if tuA != tuB:
        if tuA == 'microseconds':
            if tuB == 'microseconds':
                tB = tvalB*1.
            elif tuB == 'seconds':
                tB = tvalB*10.e6
            elif tuB == 'minutes':
                tB = tvalB*60.*10.e6
            elif tuB == 'hours':
                tB = tvalB*3600.*10.e6
            elif tuB == 'days':
                tB = tvalB*3600.*24.*10.e6
            else:
                print errormsg
                print '  ' + fname + ": combination of time untis: '" + tuA +        \
                  "' & '" + tuB + "' not ready !!"
                quit(-1)
        elif tuA == 'seconds':
            if tuB == 'microseconds':
                tB = tvalB/10.e6
            elif tuB == 'seconds':
                tB = tvalB*1.
            elif tuB == 'minutes':
                tB = tvalB*60.
            elif tuB == 'hours':
                tB = tvalB*3600.
            elif tuB == 'days':
                tB = tvalB*3600.*24.
            else:
                print errormsg
                print '  ' + fname + ": combination of time untis: '" + tuA +        \
                  "' & '" + tuB + "' not ready !!"
                quit(-1)
        elif tuA == 'minutes':
            if tuB == 'microseconds':
                tB = tvalB/(60.*10.e6)
            elif tuB == 'seconds':
                tB = tvalB/60.
            elif tuB == 'minutes':
                tB = tvalB*1.
            elif tuB == 'hours':
                tB = tvalB*60.
            elif tuB == 'days':
                tB = tvalB*60.*24.
            else:
                print errormsg
                print '  ' + fname + ": combination of time untis: '" + tuA +        \
                  "' & '" + tuB + "' not ready !!"
                quit(-1)
        elif tuA == 'hours':
            if tuB == 'microseconds':
                tB = tvalB/(3600.*10.e6)
            elif tuB == 'seconds':
                tB = tvalB/3600.
            elif tuB == 'minutes':
                tB = tvalB/60.
            elif tuB == 'hours':
                tB = tvalB*1.
            elif tuB == 'days':
                tB = tvalB*24.
            else:
                print errormsg
                print '  ' + fname + ": combination of time untis: '" + tuA +        \
                  "' & '" + tuB + "' not ready !!"
                quit(-1)
        elif tuA == 'days':
            if tuB == 'microseconds':
                tB = tvalB/(24.*3600.*10.e6)
            elif tuB == 'seconds':
                tB = tvalB/(24.*3600.)
            elif tuB == 'minutes':
                tB = tvalB/(24.*60.)
            elif tuB == 'hours':
                tB = tvalB/24.
            elif tuB == 'days':
                tB = tvalB*1.
            else:
                print errormsg
                print '  ' + fname + ": combination of time untis: '" + tuA +        \
                  "' & '" + tuB + "' not ready !!"
                quit(-1)
        else:
            print errormsg
            print '  ' + fname + ": time untis: '" + tuA + "' not ready !!"
            quit(-1)
    else:
        tB = tvalB*1.

    if trefA != trefB:
        trefTA = dt.datetime.strptime(trefA, '%Y-%m-%d %H:%M:%S')
        trefTB = dt.datetime.strptime(trefB, '%Y-%m-%d %H:%M:%S')

        difft = trefTB - trefTA
        diffv = difft.days*24.*3600.*10.e6 + difft.seconds*10.e6 + difft.microseconds
        print '  ' + fname + ': different reference refA:',trefTA,'refB',trefTB
        print '    difference:',difft,':',diffv,'microseconds'

        if tuA == 'microseconds':
            tB = tB + diffv
        elif tuA == 'seconds':
            tB = tB + diffv/10.e6
        elif tuA == 'minutes':
            tB = tB + diffv/(60.*10.e6)
        elif tuA == 'hours':
            tB = tB + diffv/(3600.*10.e6)
        elif tuA == 'dayss':
            tB = tB + diffv/(24.*3600.*10.e6)
        else:
            print errormsg
            print '  ' + fname + ": time untis: '" + tuA + "' not ready !!"
            quit(-1)

    return tB

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

strCFt="Refdate,tunits (CF reference date [YYYY][MM][DD][HH][MI][SS] format and " +  \
  " and time units: 'weeks', 'days', 'hours', 'miuntes', 'seconds')"

kindobs=['multi-points', 'single-station', 'trajectory']
strkObs="kind of observations; 'multi-points': multiple individual punctual obs " +  \
  "(e.g., lightning strikes), 'single-station': single station on a fixed position,"+\
  "'trajectory': following a trajectory"
#sumc,[constant]: add [constant] to variables values
#subc,[constant]: substract [constant] to variables values
#mulc,[constant]: multipy by [constant] to variables values
#divc,[constant]: divide by [constant] to variables values

parser = OptionParser()
parser.add_option("-d", "--dimensions", dest="dims",
  help="[DIM]@[simdim]@[obsdim] ',' list of couples of dimensions names from each source ([DIM]='X','Y','Z','T'; None, no value)",
  metavar="VALUES")
parser.add_option("-D", "--vardimensions", dest="vardims",
  help="[DIM]@[simvardim]@[obsvardim] ',' list of couples of variables names with dimensions values from each source ([DIM]='X','Y','Z','T'; None, no value)", metavar="VALUES")
parser.add_option("-k", "--kindObs", dest="obskind", type='choice', choices=kindobs, 
  help=strkObs, metavar="FILE")
parser.add_option("-l", "--stationLocation", dest="stloc",  
  help="longitude, latitude and height of the station (only for 'single-station')", 
  metavar="FILE")
parser.add_option("-o", "--observation", dest="fobs",
  help="observations file to validate", metavar="FILE")
parser.add_option("-s", "--simulation", dest="fsim",
  help="simulation file to validate", metavar="FILE")
parser.add_option("-v", "--variables", dest="vars",
  help="[simvar]@[obsvar]@[[oper]@[val]] ',' list of couples of variables to validate and if necessary operation and value", metavar="VALUES")

(opts, args) = parser.parse_args()

#######    #######
## MAIN
    #######

ofile='validation_sim.nc'

if opts.dims is None:
    print errormsg
    print '  ' + main + ': No list of dimensions are provided!!'
    print '    a ',' list of values X@[dimxsim]@[dimxobs],...,T@[dimtsim]@[dimtobs]'+\
      ' is needed'
    quit(-1)
else:
    print main +': couple of dimensions _______'
    dims = {}
    ds = opts.dims.split(',')
    for d in ds:
        dsecs = d.split('@')
        if len(dsecs) != 3:
            print errormsg
            print '  ' + main + ': wrong number of values in:',dsecs,' 3 are needed !!'
            print '    [DIM]@[dimnsim]@[dimnobs]'
            quit(-1)
        dims[dsecs[0]] = [dsecs[1], dsecs[2]]
        print dsecs[0],':',dsecs[1],',',dsecs[2]
        

if opts.vardims is None:
    print errormsg
    print '  ' + main + ': No list of variables with dimension values are provided!!'
    print '    a ',' list of values X@[vardimxsim]@[vardimxobs],...,T@' +  \
      '[vardimtsim]@[vardimtobs] is needed'
    quit(-1)
else:
    print main +': couple of variable dimensions _______'
    vardims = {}
    ds = opts.vardims.split(',')
    for d in ds:
        dsecs = d.split('@')
        if len(dsecs) != 3:
            print errormsg
            print '  ' + main + ': wrong number of values in:',dsecs,' 3 are needed !!'
            print '    [DIM]@[vardimnsim]@[vardimnobs]'
            quit(-1)
        vardims[dsecs[0]] = [dsecs[1], dsecs[2]]
        print dsecs[0],':',dsecs[1],',',dsecs[2]

if opts.obskind is None:
    print errormsg
    print '  ' + main + ': No kind of observations provided !!'
    quit(-1)
else:
    obskind = opts.obskind
    if obskind == 'single-station':
        if opts.stloc is None:
            print errormsg
            print '  ' + main + ': No station location provided !!'
            quit(-1)
        else:
            stationdesc = [np.float(opts.stloc.split(',')[0]),                       \
              np.float(opts.stloc.split(',')[1]), np.float(opts.stloc.split(',')[2])]

if opts.fobs is None:
    print errormsg
    print '  ' + main + ': No observations file is provided!!'
    quit(-1)
else:
    if not os.path.isfile(opts.fobs):
        print errormsg
        print '   ' + main + ": observations file '" + opts.fobs + "' does not exist !!"
        quit(-1)

if opts.fsim is None:
    print errormsg
    print '  ' + main + ': No simulation file is provided!!'
    quit(-1)
else:
    if not os.path.isfile(opts.fsim):
        print errormsg
        print '   ' + main + ": simulation file '" + opts.fsim + "' does not exist !!"
        quit(-1)

if opts.vars is None:
    print errormsg
    print '  ' + main + ': No list of couples of variables is provided!!'
    print '    a ',' list of values [varsim]@[varobs],... is needed'
    quit(-1)
else:
    valvars = []
    vs = opts.dims.split(',')
    for v in vs:
        vsecs = v.split('@')
        if len(dsecs) < 2:
            print errormsg
            print '  ' + main + ': wrong number of values in:',vsecs,                \
              ' at least 2 are needed !!'
            print '    [varsim]@[varobs]@[[oper][val]]'
            quit(-1)
        valvars.append(vsecs)

# Openning observations trajectory
##
oobs = NetCDFFile(opts.fobs, 'r')

valdimobs = {}
for dn in dims:
    print dn,':',dims[dn]
    if dims[dn][1] != 'None':
        if not oobs.dimensions.has_key(dims[dn][1]):
            print errormsg
            print '  ' + main + ": observations file does not have dimension '" +    \
              dims[dn][1] + "' !!"
            quit(-1)
        if vardims[dn][1] != 'None':
            if not oobs.variables.has_key(vardims[dn][1]):
                print errormsg
                print '  ' + main + ": observations file does not have varibale " +  \
                  "dimension '" + vardims[dn][1] + "' !!"
                quit(-1)
            valdimobs[dn] = oobs.variables[vardims[dn][1]][:]
    else:
        if dn == 'X':
            valdimobs[dn] = stationdesc[0]
        elif dn == 'Y':
            valdimobs[dn] = stationdesc[1]
        elif dn == 'Z':
            valdimobs[dn] = stationdesc[2]

osim = NetCDFFile(opts.fsim, 'r')

valdimsim = {}
for dn in dims:
    if not osim.dimensions.has_key(dims[dn][0]):
        print errormsg
        print '  ' + main + ": simulation file does not have dimension '" +          \
          dims[dn][0] + "' !!"
        quit(-1)
    if not osim.variables.has_key(vardims[dn][0]):
        print errormsg
        print '  ' + main + ": simulation file does not have varibale dimension '" + \
          vardims[dn][0] + "' !!"
        quit(-1)
    valdimsim[dn] = osim.variables[vardims[dn][0]][:]

# General characteristics
dimtobs = len(valdimobs['T'])
dimtsim = len(valdimsim['T'])

print main +': observational time-steps:',dimtobs,'simulation:',dimtsim

if obskind == 'multi-points':
    trajpos = np.zeros((2,dimt),dtype=int)
    for it in dimtobs:
        trajpos[:,it] = index_2mat(valdimsim['X'],valdimsim['Y'],                    \
          [valdimobs['X'][it],valdimobss['Y'][it]])

elif obskind == 'single-station':
    stsimpos = index_2mat(valdimsim['Y'],valdimsim['X'],[valdimobs['Y'],             \
      valdimobs['X']])
    print main + ': station point in simulation:', stsimpos
    print '    station position:',valdimobs['X'],',',valdimobs['Y']
    print '    simulation coord.:',valdimsim['X'][tuple(stsimpos)],',',              \
      valdimsim['Y'][tuple(stsimpos)]
elif obskind == 'trajectory':
    if dims.has_key('Z'):
        trajpos = np.zeros((3,dimt),dtype=int)
        for it in dimtobs:
            trajpos[0:1,it] = index_2mat(valdimsim['X'],valdimsim['Y'],              \
              [valdimobs['X'][it],valdimobss['Y'][it]])
            trajpos[2,it] = index_mat(valdimsim['Z'],valdimobs['Z'][it])
    else:
        trajpos = np.zeros((2,dimt),dtype=int)
        for it in dimtobs:
            trajpos[:,it] = index_2mat(valdimsim['X'],valdimsim['Y'],                \
              [valdimobs['X'][it],valdimobss['Y'][it]])

# Getting times
tobj = oobs.variables[vardims['T'][1]]
obstunits = tobj.getncattr('units')
tobj = osim.variables[vardims['T'][0]]
simtunits = tobj.getncattr('units')

simobstimes = coincident_CFtimes(valdimsim['T'], obstunits, simtunits)

print 'Lluis:',simobstimes