#-------------------------------------------------------------------------------
# Name:        wellfun-DeGlee.py
# Purpose:     Grondwaterverlaging volgens De Glee (1930)
# Created:     7 januari 2012
# Copyright:   (c) Grondwaterformules.nl
#-------------------------------------------------------------------------------
#!/usr/bin/env python

from pylab import *
from scipy import *
from scipy.special import kn as BesselK

def DeGlee(r,Q0,kD,c):
    labda = math.sqrt(kD*c)
    phi = Q0/(2*pi*kD)*BesselK(0,r/labda)
    q = Q0*r/labda*BesselK(1,r/labda)
    return phi,q,labda

def plot(r,VAR,plotvar='PHI'):

    # 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 the plot
    if plotvar == 'PHI':
        graph.plot(r, -VAR, color=orange, lw=2.0)

        # Set axes limits
        # If you see a plot without lines, delete or comment next line:
        graph.set_ylim(-3, 0)

        # Create annotation
        ylabel('verlaging (m)')
        xlabel('afstand tot de put (m)')

        # Filename
        filename = "deglee_phi.png"

    elif plotvar =='Q':
        graph.plot(r, VAR, color=orange, lw=2.0)

        #Create annotation
        ylabel('debiet (m2/dag)')
        xlabel('afstand tot de put (m)')

        # Filename
        filename = "deglee_qx.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__':

    r = arange(0,5000.0,0.1) # r is een numpy array
    Q0 = 1000000.0/365.0
    kD = 800.0
    c  = 1500.0

    phi,q,labda = DeGlee(r,Q0,kD,c)
    plot(r,phi,'PHI')
    plot(r,q,'Q')