### ===========================================================================
###
### Checks integrals
###
### ===========================================================================
##
##  Warning, to install, configure, run, use any of Olivier Marti's
##  software or to read the associated documentation you'll need at least
##  one (1) brain in a reasonably working order. Lack of this implement
##  will void any warranties (either express or implied).
##  O. Marti assumes no responsability for errors, omissions,
##  data loss, or any other consequences caused directly or indirectly by
##  the usage of his software by incorrectly or partially configured
##  personal.
##
## SVN information
__Author__   = "$Author$"
__Date__     = "$Date$"
__Revision__ = "$Revision$"
__Id__       = "$Id:  $"
__HeadURL    = "$HeadURL$"

import cdms2
import numpy as np
import glob
import sys, argparse, textwrap
import nemo

# Creating a parser
parser = argparse.ArgumentParser (
    description = """Compute integral of matching coupling fields on atmosphere and ocean side.""",
    epilog='-------- This is the end of the help message --------')

# Adding arguments
parser.add_argument ( '--file'      , help='Input file, e.g. dia_tICO40_to_teORCA1.2_1storder_false.nc', type=str, default='dia_tico_to_torc_1storder_false.nc' )
parser.add_argument ( '--grids'     , help='grids file', type=str, default='grids.nc' )
parser.add_argument ( '--areas'     , help='areas file', type=str, default='areas.nc' )
parser.add_argument ( '--masks'     , help='masks file', type=str, default='masks.nc' )
parser.add_argument ( '--fracs'     , help='fracs file', type=str, default=None )
parser.add_argument ( '--guess'     , help='If set, build grids, areas and masks file name from the dia file name', action="store_true", default=False )
parser.add_argument ( '--src_int'   , help='Source field is integrated'      , action="store_true" )
parser.add_argument ( '--dst_int'   , help='Destination field is integrated' , action="store_true" )
parser.add_argument ( '--src_nofrac', help='Do not apply frac on source'     , action="store_true" )
parser.add_argument ( '--dst_nofrac', help='Do not apply frac on destination', action="store_true" )

# Parse command line
myargs = parser.parse_args()

#
name_input = myargs.file

# Determination des modeles
Elements = myargs.file.split('_')

#
Name_src = Elements[1][1:] ; Name_dst = Elements[3][1:]
grid_src = Elements[1][0]  ; grid_dst = Elements[3][0]

Grid_src = grid_src.upper() ; Grid_dst = grid_dst.upper()
if Grid_src in ( 'A', 'C') : Grid_src = 'T'
if Grid_dst in ( 'A', 'C') : Grid_dst = 'T'

# src side
if Name_src.count('ico') != 0 or Name_src.count('lmd') != 0 : name_src = Name_src ; ShortNames = True
else :
    if Name_src.count('ICO') : name_src = 'ico' ; ShortNames = False
    if Name_src.count('LMD') : name_src = 'lmd' ; ShortNames = False

if Name_dst.count('orc') != 0    : name_src = Name_src ; ShortNames = True
else :
    if Name_src.count('ORC')!= 0 : name_src = 'orc'    ; ShortNames = False
        
# dst side
if Name_dst.count('ico')!= 0 or Name_dst.count('lmd') != 0 : name_dst = Name_dst ; ShortNames = True
else :
    if Name_dst.count('ICO') : name_dst = 'ico' ; ShortNames = False
    if Name_dst.count('LMD') : name_dst = 'lmd' ; ShortNames = False

if Name_dst.count('orc') != 0 : name_dst = Name_dst ; ShortNames = True
else :
    if Name_dst.count('ORC') != 0 : name_dst = 'orc' ; ShortNames = False
    
if Name_src.count('ORC') : CplModel = Name_src + 'x' + Name_dst
if Name_dst.count('ORC') : CplModel = Name_dst + 'x' + Name_src

print ('CplModel : ' + CplModel )
print ('Input file: ' + name_input )
     
# Coordonnees et masques
if myargs.fracs == None :
    if Name_src.count('ORC') : name_frc = Name_dst + '_grid_maskFrom_' + Name_src + '.nc'
    if Name_dst.count('ORC') : name_frc = Name_src + '_grid_maskFrom_' + Name_dst + '.nc'
else:
    name_frc = myargs.fracs

if myargs.guess :
    n_grids =  'grids_' + CplModel + '.nc'
    n_masks =  'masks_' + CplModel + '.nc'
    n_areas =  'areas_' + CplModel + '.nc'
else :
    n_grids = myargs.grids
    n_masks = myargs.masks
    n_areas = myargs.areas

print ('Opening grids file: ' + n_grids )
f_grids = cdms2.open ( n_grids )
print ('Opening masks file: ' + n_masks )
f_masks = cdms2.open ( n_masks )
print ('Opening areas file: ' + n_areas )
f_areas = cdms2.open ( n_areas )
print ('Opening fracs file: ' + name_frc )   
f_frac  = cdms2.open ( name_frc     )

#
msk_src = np.float64(1.0) - f_masks ( grid_src + name_src + '.msk', squeeze=True )
lon_src = f_grids ( grid_src + name_src + '.lon', squeeze=True )
lat_src = f_grids ( grid_src + name_src + '.lat', squeeze=True )
if myargs.src_int : srf_src = np.ones ( (msk_src.shape), dtype=np.float64 )
else              : srf_src = f_areas ( grid_src + name_src + '.srf', squeeze=True )

#
msk_dst = np.float64(1.0) - f_masks ( grid_dst + name_dst + '.msk', squeeze=True )
lon_dst = f_grids ( grid_dst + name_dst + '.lon', squeeze=True )
lat_dst = f_grids ( grid_dst + name_dst + '.lat', squeeze=True )
if myargs.dst_int : srf_dst = np.ones ( (msk_dst.shape), dtype=np.float64)
else              : srf_dst = f_areas ( grid_dst + name_dst + '.srf' )

#
if name_src in ( 'ico', 'lmd' ) :
    if myargs.src_nofrac : frc_src = np.ones ( (msk_src.shape), dtype=np.float64)
    else                 : frc_src = f_frac ( 'OceFrac', squeeze=True )
if name_src == 'orc'     : frc_src = np.ones ( (msk_src.shape), dtype=np.float64)
if name_dst in ( 'ico', 'lmd' ) :
    if myargs.dst_nofrac : frc_dst = np.ones ( (msk_dst.shape), dtype=np.float64)
    else                 : frc_dst = f_frac ( 'OceFrac', squeeze=True )
if name_dst == 'orc'     : frc_dst = np.ones ( (msk_dst.shape), dtype=np.float64)

# Periodicity
nperio_src = 0 ; nperio_dst = 0

if Name_src == 'ORCA2.3'  : nperio_src = 4
if Name_src == 'eORCA1.2' : nperio_src = 6
if Name_src == 'eORCA025' : nperio_src = 6

if Name_dst == 'ORCA2.3'  : nperio_dst = 4
if Name_dst == 'eORCA1.2' : nperio_dst = 6
if Name_dst == 'eORCA025' : nperio_dst = 6

# NEMO periodicity.
if Name_src in ('ORCA2.3', 'eORCA1.2', 'eORCA025') : msk_src = nemo.lbc_mask ( msk_src, nperio=nperio_src, cd_type=Grid_src )
if Name_dst in ('ORCA2.3', 'eORCA1.2', 'eORCA025') : msk_dst = nemo.lbc_mask ( msk_dst, nperio=nperio_dst, cd_type=Grid_dst )
    
# Surfaces
area_src = np.sum ( srf_src * msk_src * frc_src )
area_dst = np.sum ( srf_dst * msk_dst * frc_dst )

print ( Name_src )
print 'mask: {:12.3} {:12.3} {:12.3}'.format( np.min(msk_src), np.max(msk_src), np.sum(msk_src) )
print 'frac: {:12.3} {:12.3} {:12.3}'.format( np.min(frc_src), np.max(frc_src), np.sum(frc_src) )
print 'area: {:12.3} {:12.3} {:12.3}'.format( np.min(srf_src), np.max(srf_src), np.sum(srf_src) )

print ( Name_dst )
print 'mask: {:12.3} {:12.3} {:12.3}'.format( np.min(msk_dst), np.max(msk_dst), np.sum(msk_dst) )
print 'frac: {:12.3} {:12.3} {:12.3}'.format( np.min(frc_dst), np.max(frc_dst), np.sum(frc_dst) )
print 'area: {:12.3} {:12.3} {:12.3}'.format( np.min(srf_dst), np.max(srf_dst), np.sum(srf_dst) )
print (' ')
print ( "Surfaces   : %14.6e %14.6e %11.3e"%( area_src, area_dst, (area_src-area_dst)/(area_src+area_dst)*0.5 ) )
print (' ')

##
f_input = cdms2.open ( name_input )
# Loop over fields
for num in np.arange (1, 7) :
    name_src = 'field{:02}_src'.format(num)
    name_dst = 'field{:02}_dst'.format(num)
    print ( '{:02}'.format(num) + ':' + name_src + ':' + name_dst)

    # Reading field
    v_src = f_input ( name_src )
    v_dst = f_input ( name_dst )

    # Compute integrals
    sum_src = np.sum (v_src * srf_src * msk_src * frc_src)
    sum_dst = np.sum (v_dst * srf_dst * msk_dst * frc_dst)
    
    # Compute average
    mean_src = sum_src / area_src
    mean_dst = sum_dst / area_dst

    #
    print ( "Integrals : %14.6e %14.6e %11.3e"%( sum_src , sum_dst , (sum_src - sum_dst)/(sum_src + sum_dst)*0.5 ) )
    print ( "Averages  : %14.6e %14.6e %11.3e"%( mean_src, mean_dst, (mean_src-mean_dst)/(mean_src+mean_dst)*0.5 ) )
    print ( " " )
#=======================================