#-------------------------------------------------------------------------------
# Name:        waterloop-gewasverdamping.py
# Purpose:     Rise of groundwater table after surface water level rise
#              Ernst(1958) gives three formula's:
#               1. River without entry resitance
#               2. Compensation for evaporation increase (no entry resistance)
#               3. River with entry resitance
#              This script contains formulas 1 and 2 for description of
#              formula 2 at Grondwaterformules.nl
# Source:      L.F. Ernst, 1959. Verhoging van grondwaterstanden en vermindering
#              van afvoer door opstuwing van beken. Instituut voor Cultuurtechniek
#              en Watrerhuishouding, Wageningen. Mededeling 8.
# Created:     10-03-2012
# Copyright:   (c) Grondwaterformules
# Licence:     Python 2.6 with SciPy
#-------------------------------------------------------------------------------
#!/usr/bin/env python
from scipy import *
from scipy import sqrt as sqrt
from scipy import exp as exp
from scipy.special import erfc as erfc
from scipy.special import erf as erf
from scipy import cosh as cosh
from pylab import *

def Ernst1958_1(x,t,kD,S):
    beta = sqrt(kD/S)
    return  erfc(x/(2*beta*sqrt(t)))

def Ernst1958_2(x,t,kD,S,h0,C0):
    beta = sqrt(kD/S)
    g = sqrt(C0/(kD*h0))

    p1 = cosh(g*x)
    p2 = 0.5*exp(-g*x)
    p3 = erf(beta*g*sqrt(t)-x/(2*beta*sqrt(t)))
    p4 = 0.5*exp(g*x)
    p5 = erf(beta*g*sqrt(t)+x/(2*beta*sqrt(t)))
    return p1+p2*p3-p4*p5,[p1,p2,p3,p4,p5]

def plotdh(x,t,dh1,dh2):

    # Define colors
    orange = '#FF9900'
    blue   = '#2F64B2'

    # Create plots
    graph = subplot(111)

    # Set ticklines
    xticklines = graph.get_xticklines()
    yticklines = graph.get_yticklines()

    for line in xticklines+yticklines:
        line.set_linewidth(3)

    for line in xticklines:
        line.set_visible(False)

    # Set gridlines
    gridlines = graph.get_xgridlines()
    gridlines.extend( graph.get_ygridlines() )
    for line in gridlines:
        line.set_linestyle('-')
    grid(True)

    # Set ticklabels
    xticklabels = graph.get_xticklabels()
    yticklabels = graph.get_yticklabels()

    for label in xticklabels+yticklabels:
        label.set_color('k')
        label.set_fontsize('medium')
        label.set_visible(True)

    # Create annotation
    ylabel('verandering grondwaterstand (%h0)')
    xlabel('afstand tot de watergang (m)')

    # Plot lowring in time
    for ti in range(len(t)):
        dh1i = dh1[ti][:]*100
        dh2i = dh2[ti][:]*100
        graph.plot(x, dh1i, color=blue, lw=2.0)
        graph.plot(x, dh2i, color=orange, lw=2.0)

    # Filename
    filename = "waterloop-gewasverdamping.png"

    # Save figure to file
    fig = figure(1)
    mydpi = 100.0
    fig.set_dpi(mydpi)
    fig.set_size_inches(6.0,4.0)
    fig.savefig(filename,dpi=mydpi, facecolor='w', edgecolor='w')

    # Show the plot on screen
    show()

if __name__ == '__main__':

    kD = 700.0
    S = 0.25
    h0 = 1.0
    C0 = 0.004

    # Print small table of both formulas
    x = array([100,700])
    t = array([1,50,100])
    [X,T] = meshgrid(x,t)
    hfrac1=Ernst1958_1(X,T,kD,S)
    hfrac2,P = Ernst1958_2(X,T,kD,S,h0,C0)

    print "    t       x    Ernst1    Ernst2       p1       p2       p3       p4       p5"
    for xi in range(len(x)):
        for ti in range(len(t)):
            print '%5d'    %t[ti],
            print '%7.2f'  %x[xi],
            print '%9.5f'  %hfrac1[ti][xi],
            print '%9.5f'  %hfrac2[ti][xi],
            print '% 7.5f' %P[0][ti][xi],
            print '% 7.5f' %P[1][ti][xi],
            print '% 7.5f' %P[2][ti][xi],
            print '% 7.5f' %P[3][ti][xi],
            print '% 7.5f' %P[4][ti][xi]

    # Make a plot
    x = linspace(0,700,100)
    t = array([10,31,121])
    [X,T]  = meshgrid(x,t)
    hfrac1 = Ernst1958_1(X,T,kD,S)
    hfrac2,P = Ernst1958_2(X,T,kD,S,h0,C0)
    plotdh(x,t,hfrac1,hfrac2)