from fipy import Variable, FaceVariable, CellVariable, Grid1D, ExplicitDiffusionTerm, TransientTerm, DiffusionTerm, Viewer
from fipy.tools import numerix
import numpy as np
import pandas as pd


nx = 6000        #amount of steps
dx = 1.          #stepsize
L = int(nx * dx) #depth of well

mesh = Grid1D(nx=nx, dx=dx)
m = Grid1D(nx=1., dx=0.1)




T = CellVariable(name="temp vs depth DH4", mesh=mesh, value= (mesh.cellCenters[0]) / 75)
X = mesh.faceCenters[0]



toyr = 60 * 60 * 24 * 365 #to year from s

#variables into FaceVariables
k= FaceVariable(mesh=mesh)
dens= FaceVariable(mesh=mesh)
cp = FaceVariable(mesh=mesh)



k.setValue(0)
dens.setValue(0)
cp.setValue(0)

#df.D.iloc[ii]

for ii in range(len(df.d)):
    if df.lithology.iloc[ii] == 'sandstone':
        k.setValue(3 * toyr, where=(X == ii))
        dens.setValue(2200, where=(X == ii))
        cp.setValue(850, where=(X == ii))
    elif df.lithology.iloc[ii] == 'shale/siltstone':
        k.setValue(2 * toyr, where=(X == ii))
        dens.setValue(2400, where=(X == ii))
        cp.setValue(850, where=(X == ii))
    elif df.lithology.iloc[ii] == 'ign_intrusion':
        k.setValue(2 * toyr, where=(X == ii))
        dens.setValue(2700, where=(X == ii))
        cp.setValue(850, where=(X == ii))
    elif df.lithology.iloc[ii] == 'shale':
        k.setValue(2 * toyr, where=(X == ii))
        dens.setValue(2600, where=(X == ii))
        cp.setValue(850, where=(X == ii))


D = k / (dens * cp)
#find a way to apply different densities to D


time = Variable()


fluxRight = 0.1 #UNIT' heatflow influx
T.faceGrad.constrain([fluxRight], mesh.facesRight)


eqI = TransientTerm() == DiffusionTerm(coeff=D)


if __name__ == '__main__':
    viewer = Viewer(vars=(T), limits={'xmin': 0, 'xmax': 6000, 'ymin':-15, 'ymax':250} )
    viewer.plot()


D = FaceVariable(mesh=mesh, value=D)

dt = 1

while time() < 1500:
    if surface_temp.value.iloc[time] == 1:
        valueLeft = 0
        T.constrain(valueLeft, mesh.facesLeft)
    elif surface_temp.value.iloc[time] == 2:
        valueLeft = -5
        T.constrain(valueLeft, mesh.facesLeft)
    elif surface_temp.value.iloc[time] == 3:
        valueLeft = -15
        T.constrain(valueLeft, mesh.facesLeft)
    #here add the temperature of -6 which is the CMR approximation for the
    #surface temp 150000 - 800000 yrs ago

    time.setValue(time() + dt)
    eqI.solve(var=T, dt=dt)
    if __name__ == '__main__':
        viewer.plot()