#!/usr/bin/env python3
'''
  Script to check water conservation in the IPSL coupled model
 
  Here, some common utilities to all scripts
   
  Warning, to install, configure, run, use any of included 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).  Authors 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$
 $Date$
 $Revision$
 $Id: ATM_waterbudget.py 6277 2022-12-08 09:24:05Z omamce $
 $HeadURL$
'''

import numpy as np
import configparser, re

VarInt      = 10
Rho         = 1000.0
Ra          = 6366197.7236758135 #-- Earth Radius (m)
Grav        = 9.81               #-- Gravity (m^2/s
ICE_rho_ice = 917.0              #-- Ice volumic mass (kg/m3) in LIM3
ICE_rho_sno = 330.0              #-- Snow volumic mass (kg/m3) in LIM3
OCE_rho_liq = 1026.0             #-- Ocean water volumic mass (kg/m3) in NEMO
ATM_rho     = 1000.0             #-- Water volumic mass in atmosphere (kg/m^3)
SRF_rho     = 1000.0             #-- Water volumic mass in surface reservoir (kg/m^3)
RUN_rho     = 1000.0             #-- Water volumic mass of rivers (kg/m^3)
ICE_rho_pnd = 1000.              #-- Water volumic mass in ice ponds (kg/m^3)
YearLength  = 365.25 * 24. * 60. * 60. #-- Year length (s)

def setBool (chars) :
    '''Convert specific char string in boolean if possible'''
    setBool = chars
    if type (chars) == str :
        for key in configparser.ConfigParser.BOOLEAN_STATES.keys () :
            if chars.lower() == key : setBool = configparser.ConfigParser.BOOLEAN_STATES[key]
    return setBool

def setNum (chars) :
    '''Convert specific char string in integer or real if possible'''
    if type (chars) == str :
        realnum  = re.compile ("^[-+]?[0-9]*\.?[0-9]+(e[-+]?[0-9]+)?$")
        isNumber = realnum.match(chars.strip()) != None
        isInt    = chars.strip().isdigit()
        #print (chars , isNumber , isInt)
        if isNumber :
            if isInt : setNum = int   (chars)
            else     : setNum = float (chars)
        else : setNum = chars
    else : setNum = chars
    return setNum

def setNone (chars) :
    '''Convert specific char string to None if possible'''
    if type (chars) == str :
        if chars in ['None', 'NONE', 'none'] : setNone = None
        else : setNone = chars
    else : setNone = chars
    return setNone

def unDefined (char) :
    '''True if a variable is not defined, or set to None'''
    if char in locals () :
        if locals()[char] == None : unDefined = True
        else            : unDefined = False
    else                : unDefined = True
    return unDefined

def Defined (char) :
    '''True if a variable is defined and not equal to None'''
    if char in globals () :
        if globals()[char] == None : Defined = False
        else            : Defined = True
    else                : Defined = False
    return Defined

def Ksum (tab) :
    '''
    Kahan summation algorithm, also known as compensated summation
    https://en.wikipedia.org/wiki/Kahan_summation_algorithm
    '''
    Ksum = 0.0                   # Prepare the accumulator.
    comp = 0.0                   # A running compensation for lost low-order bits.
   
    for xx in np.where ( np.isnan(tab), 0., tab ) : 
        yy = xx - comp           # comp is zero the first time around.
        tt = Ksum + yy           # Alas, sum is big, y small, so low-order digits of y are lost.
        comp = (tt - Ksum) - yy  # (tt - Ksum) cancels the high-order part of y; subtracting y recovers negative (low part of yy)
        Ksum = tt                # Algebraically, comp should always be zero. Beware overly-aggressive optimizing compilers !
                                 # Next time around, the lost low part will be added to y in a fresh attempt.
    return Ksum

def Ssum (tab) :
    '''
    Precision summation by sorting by absolute values
    '''
    Ssum = np.sum ( tab[np.argsort(np.abs(tab))] )
    return Ssum

def KSsum (tab) :
    '''
    Precision summation by sorting by absolute value, and applying Kahan summation
    '''
    KSsum = Ksum ( tab[np.argsort(np.abs(tab))] )
    return KSsum

# Choosing algorithm for precise summation
Psum = KSsum

def IsLeapYear ( Year, CalendarType="Gregorian" ) :
    ''' True is Year is a leap year ''' 

    # What is the calendar :
    if CalendarType in [ '360d', '360_day', 'noleap', '365_day'] : 
        IsLeapYear = False
        return IsLeapYear

    if CalendarType in [ 'all_leap', '366_day' ] : 
        IsLeapYear = True
        return IsLeapYear
    
    Year = int ( Year )
  
    # a year is a leap year if it is even divisible by 4
    # but not evenly divisible by 100
    # unless it is evenly divisible by 400

    # if it is evenly divisible by 400 it must be a leap year
    if np.mod ( Year, 400 ) == 0 : 
        IsLeapYear = True
        return IsLeapYear
        
    # if it is evenly divisible by 100 it must not be a leap year
    if np.mod ( Year, 100 ) == 0 :
        IsLeapYear = False
        return IsLeapYear

    # if it is evenly divisible by 4 it must be a leap year
    if np.mod ( Year, 4 ) == 0 :
        IsLeapYear = True
        return IsLeapYear

    IsLeapYear = False
    return IsLeapYear

def DateFormat ( Date ) :
    '''
    Get date format :
      [yy]yymmdd   is Gregorian
      [yy]yy-mm-dd is Human
    '''
    if type (Date) == str :
        if '-' in Date : 
            DateFormat = 'Human'
        else : 
            DateFormat = 'Gregorian'
    if type (Date) == int : DateFormat = 'Gregorian'
    return DateFormat

def PrintDate ( YE, MO, DA, Format ) :
    '''Return a date in the requested format'''
    if Format == 'Human'     : PrintDate = f'{YE:04d}-{MO:02d}-{DA:02d}'
    if Format == 'Gregorian' : PrintDate = f'{YE:04d}{MO:02d}{DA:02d}'
    return PrintDate
        
def FormatToGregorian ( Date ) :
    ''' from a yyyy-mm-dd or yyymmdd date format returns
        a yyymmdd date format
    '''
    YE, MO, DA = SplitDate ( Date )
    FormatToGregorian = f'{YE:04d}{MO:02d}{DA:02d}'
    return FormatToGregorian
  
def FormatToHuman ( Date ) :
    ''' From a yyyymmdd or yyymmdd date format returns
        a yyy-mm-dd date format
    '''
    YE, MO, DA = SplitDate ( Date )
    FormatToHuman =  f'{YE_new:04d}-{MO:02d}-{DA:02d}'
    return FormatToHuman

def SplitDate  ( Date, Out='int' ) :
    ''' Split Date in format [yy]yymmdd or [yy]yy-mm-dd
        to yy, mm, dd '''
    if type (Date) == str :
        if '-' in Date : 
            Dash = True
            YE, MO, DA = Date.split ('-')
        else : 
            Dash = False
            YE = Date[:-4] ; MO = Date[-4:-2] ; DA = Date[-2:]
            YearDigits = len (YE)
    if type (Date) == int :
        DA = np.mod ( Date, 100) 
        MO = np.mod ( Date//100, 100)
        YE = Date // 10000
    
    YE = int(YE) ; MO = int(MO) ; DA=int(DA)
    return YE, MO, DA

def DateAddYear ( Date, YearInc=1 ) :
    ''' Add on year to date in format [yy]yymmdd or [yy]yy-mm-dd'''
    Format = DateFormat ( Date )
    YE, MO, DA = SplitDate ( Date )
    YE_new = YE + YearInc
    DateAddYear = PrintDate ( YE_new, MO, DA, Format)
    return DateAddYear 

def DateMinusOneDay ( Date ) :
    ''' Substracts one day to date in format [yy]yymmdd or [yy]yy-mm-dd'''
    Format = DateFormat ( Date )
    mth_length = np.array ( [31, 28, 31,  30,  31,  30,  31,  31,  30,  31,  30,  31] )
    YE, MO, DA = SplitDate ( Date )
    if IsLeapYear ( YE) : mth_length[1] += 1
    
    YE = int(YE) ; MO = int(MO) ; DA=int(DA)
    if DA ==  1 : 
        if MO == 1 : DA_new = mth_length[-1]   ; MO_new = 12     ; YE_new = YE - 1
        else       : DA_new = mth_length[MO-2] ; MO_new = MO - 1 ; YE_new = YE
    else        :    DA_new = DA - 1           ; MO_new = MO     ; YE_new = YE

    DateMinusOneDay = PrintDate ( YE_new, MO_new, DA_new, Format)
    return DateMinusOneDay

def DatePlusOneDay ( Date ) :
    ''' Add one day to date in format [yy]yymmdd or [yy]yy-mm-dd'''
    Format = DateFormat ( Date )
    mth_length = np.array ( [31, 28, 31,  30,  31,  30,  31,  31,  30,  31,  30,  31] )
    YE, MO, DA = SplitDate ( Date )
    if IsLeapYear ( YE ) : mth_length[1] += 1
    
    YE_new = YE ; MO_new=MO ; DA_new=DA+1
    if DA_new > mth_length [MO_new-1] :
        DA_new = 1 ; MO_new = MO_new + 1
        if MO_new == 13 : 
            MO_new = 1 ; YE_new = YE_new+1

    DatePlusOneDay = PrintDate ( YE_new, MO_new, DA_new, Format )
    return DatePlusOneDay