import numpy as np
import pandas as pd
%pylab inline

import scipy as sp
import matplotlib.pyplot as plt
import os, sys

#connect echoRD Tools
pathdir='../' #path to echoRD
lib_path = os.path.abspath(pathdir)
sys.path.append(lib_path)
import vG_conv as vG
from hydro_tools import plotparticles_t,hydroprofile,plotparticles_gen

try:
   import cPickle as pickle
except:
   import pickle


"""
Explanation: This is the 2D testcase of the echoRD model
New version for generic cases based on pre-defined macropore definitions
Make sure to have the dependencies numpy, pandas, scipy, matplotlib and echoRD installed.
This is the Python 3 version.
Preamble
End of explanation
"""


#connect to echoRD
import run_echoRD as rE
#connect and load project
[dr,mc,mcp,pdyn,cinf,vG]=rE.loadconnect(pathdir='../',mcinif='mcini_g63_nomac',experimental=True)
mc=rE.preproc_echoRD(mc, dr, mcp,'column_no.pick')
#mc=mcp.mcpick_out(mc,'weiherbach.pick')
runname='column_no'
mc.advectref='Shipitalo'

import macropore_ini as mpo
mpo.mac_plot(mc.macP,mc,'macs_'+runname+'.pdf')

mc.md_macdepth

mc.precf

imshow(mc.soilgrid,interpolation=None,cmap='jet') 
print(np.unique(mc.soilgrid))

precTS=pd.read_csv(mc.precf, sep=',',skiprows=3)

precTS.tstart=60
precTS.tend=60+1800
precTS.total=0.01
precTS.intense=precTS.total/(precTS.tend-precTS.tstart)
precTS

#use modified routines for binned retention definitions
mc.part_sizefac=500
mc.gridcellA=mc.mgrid.vertfac*mc.mgrid.latfac
mc.particleA=abs(mc.gridcellA.values)/(2*mc.part_sizefac) #assume average ks at about 0.5 as reference of particle size
mc.particleD=2.*np.sqrt(mc.particleA/np.pi)
mc.particleV=3./4.*np.pi*(mc.particleD/2.)**3.
mc.particleV/=np.sqrt(abs(mc.gridcellA.values)) #assume grid size as 3rd dimension
mc.particleD/=np.sqrt(abs(mc.gridcellA.values))
mc.particlemass=dr.waterdensity(np.array(20),np.array(-9999))*mc.particleV #assume 20°C as reference for particle mass
                                                                        #DEBUG: a) we assume 2D=3D; b) change 20°C to annual mean T?
mc=dr.ini_bins(mc)
mc=dr.mc_diffs(mc,np.max(np.max(mc.mxbin)))

#mc.inimf = 'ini_moist6.dat'

#mc.inimf='./moisture/07_4_1_moist.dat'
[mc,particles,npart]=dr.particle_setup(mc,True)

#non-particle setting dependent infiltration breakpoint
dummy=np.zeros((100,len(np.unique(mc.soilgrid))))
j=0
for soil in np.unique(mc.soilgrid)-1:
    #plot(i,(mc.theta100[100-i,4]/(mc.particleD*0.5))*mc.D100[i,4]+mc.ku100[i,4],'.') #q
    for i in np.arange(100):
        dummy[i,j]=pdyn.filmthick(mc.psi100[i,soil])/((mc.theta100[100-i,soil]/(mc.particleD*0.5))*mc.D100[i,soil]+mc.ku100[i,soil])
        #plot(i,((0.001*mc.particleD[0])/((mc.theta100[100-i,soil]/(mc.particleD*0.5))*mc.D100[i,soil]+mc.ku100[i,soil]))/60,'.') #time to infilt
        #plot(i,(0.5*mc.particleD[0])/mc.ku100[i,soil]/60,'.') #time to infilt
    j+=1

dummyD=pd.DataFrame(dummy.T)
#dummyD.index=mc.soilmatrix.type
#dummyD.T.iloc[:,8].plot()
dummyD.T.plot()
plot([0,100],[60,60],'-',label='1min')
plot([0,100],[300,300],'-',label='5min')
plot([0,100],[1800,1800],'-',label='30min')
yscale('log')
#legend(ncol=3)
title('exfiltration time [s] at pore wall at different theta*')

mc.particlemass

mc.soilmatrix

#new reference
mc.maccon=np.where(mc.macconnect.ravel()>0)[0] #index of all connected cells
#mc.md_macdepth=np.abs(mc.md_macdepth)

#define bin assignment mode for infiltration particles
mc.LTEdef='instant'#'ks' #'instant' #'random'
mc.LTEmemory=mc.soilgrid.ravel()*0.

mc.prects=False
#theta=mc.zgrid[:,1]*0.+0.273
#[mc,particles,npart]=rE.particle_setup_obs(theta,mc,vG,dr,pdyn)
[thS,npart]=pdyn.gridupdate_thS(particles.lat,particles.z,mc)
[A,B]=plotparticles_t(particles,thS/100.,mc,vG,store=True)

t=0
i=0
plotparticles_gen(particles,mc,pdyn,vG,runname,t,i,saving=False,relative=False,fsize=(1.55,4.2),w_rat=[3.2,0.6],h_rat=[0.8,9])

"""
Explanation: Prepare echoRD
End of explanation
"""


mc.LTEpercentile=70 #new parameter
particles.LTEbin=np.fmin(particles.LTEbin.values,len(mc.D)-1)

t_end=6.*3600.
saveDT=True

#1: MDA
#2: MED
#3: rand
infiltmeth='MDA'
#3: RWdiff
#4: Ediss
#exfiltmeth='RWdiff'
exfiltmeth='Ediss'
#5: film_uconst
#6: dynamic u
film=True
#7: maccoat1
#8: maccoat10
#9: maccoat100
macscale=1. #scale the macropore coating 
clogswitch=False
infiltscale=False#0.8


#mc.dt=0.11
#mc.splitfac=5
#pdyn.part_diffusion_binned_pd(particles,npart,thS,mc)


drained=pd.DataFrame(np.array([]))
leftover=0
output=60. #mind to set also in TXstore.index definition

dummy=np.floor(t_end/output)
#prepare store arrays
TSstore=np.zeros((int(dummy),np.shape(thS)[0],np.shape(thS)[1]))
NPstore=np.zeros((int(dummy),len(mc.zgrid[:,1])+1))
colnames=['part1','part2','part3','part4','part5']
TXstore=pd.DataFrame(np.zeros((int(dummy),len(colnames))),columns=colnames)
t=0.
#loop through plot cycles
for i in np.arange(dummy.astype(int)):
    #plot and store states
    plotparticles_gen(particles,mc,pdyn,vG,runname,t,i,saving=True,relative=False)
    [TXstore.iloc[i],NPstore[i,:]]=plotparticles_t(particles,thS/100.,mc,vG,runname,t,i,saving=True,store=True)
    #store thS
    TSstore[i,:,:]=thS
    [particles,npart,thS,leftover,drained,t]=rE.CAOSpy_rundx1(i*output,(i+1)*output,mc,pdyn,cinf,precTS,particles,leftover,drained,6.,splitfac=1,prec_2D=False,maccoat=macscale,saveDT=saveDT,clogswitch=clogswitch,infilt_method=infiltmeth,exfilt_method=exfiltmeth,film=film,infiltscale=infiltscale)
    
    #if i/5.==np.round(i/5.):
    #    with open(''.join(['./results/Z',runname,'_Mstat.pick']),'wb') as handle:
    #       pickle.dump(pickle.dumps([pickle.dumps(particles),pickle.dumps([leftover,drained,t])]), handle, protocol=2)

    
    

"""
Explanation: Run Model
End of explanation
"""


