from underworld import UWGeodynamics as GEO
from underworld import visualisation as vis
import numpy as np

u = GEO.u

velocity = 1e-9 * u.kilometre / u.hour
model_length = 20. * u.kilometre
model_height = 10. * u.kilometre
refViscosity = 1e7 * u.pascal * u.second
surfaceTemp = 298.15 * u.degK
baseModelTemp = 623.15 * u.degK
bodyforce = 2500 * u.kilogram / u.metre**3 * 9.81 * u.meter / u.second**2

KL = model_length
Kt = KL / velocity
KM = bodyforce * KL**2 * Kt**2
KT = (baseModelTemp - surfaceTemp)

GEO.scaling_coefficients["[length]"] = KL
GEO.scaling_coefficients["[time]"] = Kt
GEO.scaling_coefficients["[mass]"]= KM
GEO.scaling_coefficients["[temperature]"] = KT

Model = GEO.Model(elementRes=(100,50), 
                  minCoord=(0. * u.kilometre, -3.5 * u.kilometre), 
                  maxCoord=(20. * u.kilometre, 3.5 * u.kilometre))

#Model.outputDir="/home/jovyan/workspace/output-sandbox5/"

Model.minViscosity = 1e20 * u.pascal * u.second
Model.maxViscosity = 1e25 * u.pascal * u.second
Model.diffusivity = 1e-6 * u.metre**2 / u.second 
Model.capacity    = 1000. * u.joule / (u.kelvin * u.kilogram)

Model.mesh_advector(axis=0)

air        = Model.add_material(name="Air", shape=GEO.shapes.Layer2D(top=Model.top, bottom=0.0))
sand1      = Model.add_material(name="Sand1", shape=GEO.shapes.Layer2D(top=air.bottom, bottom=Model.bottom))
sand2      = Model.add_material(name="Sand2", shape=GEO.shapes.Layer2D(top=-1.0 * u.kilometre, bottom=-1.5*u.kilometre))
microbeads = Model.add_material(name="Microbeads", shape=GEO.shapes.Layer2D(top=-2.5 * u.kilometre, bottom=-3.0*u.kilometre))

import numpy as np

wedge = [( 15.* u.kilometre, 0. * u.kilometre),
        (20.* u.kilometre, 5.*u.kilometre * np.tan(np.radians(10))),
        (5, 0. * u.kilometre)]

#wedge = [( 15.* u.centimeter, 0. * u.centimeter),
 #        (Model.maxCoord[0], 10.*u.centimeter * np.tan(np.radians(10))),
 #        (Model.maxCoord[0], 0. * u.centimeter)]

sand3     = Model.add_material(name="Sand3", shape=GEO.shapes.Polygon(wedge))

import numpy as np

npoints = 100

coords = np.ndarray((npoints, 2))
coords[:, 0] = np.linspace(GEO.nd(Model.minCoord[0]), GEO.nd(Model.maxCoord[0]), npoints)
coords[:, 1] = GEO.nd(sand1.top)

Model.add_passive_tracers(name="Interface1", vertices=coords)

coords[:, 1] = GEO.nd(sand2.top)
Model.add_passive_tracers(name="Interface2", vertices=coords)

coords[:, 1] = GEO.nd(sand2.bottom)
Model.add_passive_tracers(name="Interface3", vertices=coords)

coords[:, 1] = GEO.nd(microbeads.top)
Model.add_passive_tracers(name="Interface4", vertices=coords)

coords[:, 1] = GEO.nd(microbeads.bottom)
Model.add_passive_tracers(name="Interface5", vertices=coords)

from underworld import visualisation as vis
Fig = vis.Figure(figsize=(900,300))
Fig.Points(Model.Interface1_tracers, pointSize=2.0)
Fig.Points(Model.Interface2_tracers, pointSize=2.0)
Fig.Points(Model.Interface3_tracers, pointSize=2.0)
Fig.Points(Model.Interface4_tracers, pointSize=2.0)
Fig.Points(Model.Interface5_tracers, pointSize=2.0)
Fig.Points(Model.swarm, Model.materialField, fn_size=2.0)
Fig.show()

air.density        = 10. * u.kilogram / u.metre**3
sand1.density      = 2550. * u.kilogram / u.metre**3
sand2.density      = 2550. * u.kilogram / u.metre**3
sand3.density      = 2550. * u.kilogram / u.metre**3
microbeads.density = 2200. * u.kilogram / u.metre**3

air.viscosity        = 1.0e5 * u.pascal * u.second
sand1.viscosity      = 1.0e21 * u.pascal * u.second
sand2.viscosity      = 1.0e21 * u.pascal * u.second
sand3.viscosity      = 1.0e21 * u.pascal * u.second
microbeads.viscosity = 1.0e21 * u.pascal * u.second

sandFriction = np.tan(np.radians(36.0))
sandFrictionW = np.tan(np.radians(31.0))

microbeadsFriction = np.tan(np.radians(22.0))
microbeadsFrictionW = np.tan(np.radians(21.0))

sand1.plasticity       = GEO.DruckerPrager(cohesion=10.*u.pascal, frictionCoefficient=sandFriction, frictionAfterSoftening=sandFrictionW)
sand2.plasticity       = GEO.DruckerPrager(cohesion=10.*u.pascal, frictionCoefficient=sandFriction, frictionAfterSoftening=sandFrictionW)
sand3.plasticity       = GEO.DruckerPrager(cohesion=10.*u.pascal, frictionCoefficient=sandFriction, frictionAfterSoftening=sandFrictionW)
microbeads.plasticity  = GEO.DruckerPrager(cohesion=10.*u.pascal, frictionCoefficient=microbeadsFriction, frictionAfterSoftening=microbeadsFrictionW)


import underworld.function as fn

conditions = [(Model.y < GEO.nd(Model.bottom +  0.2 * u.centimeter), 0.0),
              (Model.y > GEO.nd(Model.bottom +  0.4 * u.kilometre), GEO.nd(-2.5 * u.kilometre / u.hour)),
              (True, (GEO.nd(-2.5 * u.kilometre / u.hour) / GEO.nd(0.2 * u.kilometre) * (Model.y - GEO.nd(Model.bottom) - GEO.nd(0.2 * u.kilometre))))]

fn_condition = fn.branching.conditional(conditions)

Model.set_velocityBCs(left=[0.,0.], right=[fn_condition, 0.], bottom=[0.,0.])

Model.set_temperatureBCs(top=298.15 * u.degK, 
                        bottom=623.15 * u.degK, 
                        materials=[ (air, 293.15 * u.degK)])

A = Model.set_frictional_boundary(left=np.tan(np.radians(19.)), right=np.tan(np.radians(19.)), bottom=np.tan(np.radians(19.)), top=None, thickness=3)

from underworld import visualisation as vis
Fig = vis.Figure(figsize=(900,300))
Fig.Points(Model.swarm, A.friction, fn_size=1.0)
Fig.Mesh(Model.mesh)
Fig.show()

Model.init_model(temperature="steady-state", pressure="lithostatic")
Model.solver.set_inner_method("mumps")
Model.solver.set_penalty(1e6)

## Velocity Boundary Conditions
Model.set_velocityBCs(bottom=[0.,0.], top=[None, 0.])

## Viscous Dissipation Tool

DissipationField = Model.add_swarm_variable(name='Dissipation', restart_variable=True, count=1)

def updateVD():
     #Dissipation = Model.strainRate_2ndInvariant[0]*Model.projStressTensor[0] + Model.strainRate_2ndInvariant[1]*Model.projStressTensor[1] +Model.strainRate_2ndInvariant[2]*Model.projStressTensor[2]
    Dissipation = Model._viscosityFn * Model.strainRate_2ndInvariant**2
    ### in ND form
#    DissipationField.data[:] = Dissipation.evaluate(Model.swarm) 
    
    ### Dim form 
    DissipationField.data[:] = GEO.dimensionalise(Dissipation.evaluate(Model.swarm), 
                                u.pascal * u.second * u.second**2).to_base_units()

#### scaled   
### for pre solve update of swarm variable
Model.pre_solve_functions["A-pre"] = updateVD

### for post solve update of swarm variable
 #Model.post_solve_functions["A-post"] = updateVD

# ## run
#Model.run_for(nstep=2, checkpoint_interval=1*u.year)
from underworld import UWGeodynamics as GEO
from underworld import visualisation as vis
from underworld.scaling import units as u

fn_map = fn.branching.map(fn_key=2, mapping={1: 1.}, fn_default=0.)
Model.run_for(0.01 * u.hour, checkpoint_interval=1*u.year)

Fig = vis.Figure(figsize=(1200,400))
Fig.Points(Model.swarm, DissipationField, fn_size=2.0, colours='white green red purple blue')
Fig.show()

Fig = vis.Figure(figsize=(900,300),title="Material Field")
Fig.Points(Model.Interface1_tracers, pointSize=2.0)
Fig.Points(Model.Interface1_tracers, pointSize=2.0)
Fig.Points(Model.Interface1_tracers, pointSize=2.0)
Fig.Points(Model.Interface1_tracers, pointSize=2.0)
Fig.Points(Model.Interface1_tracers, pointSize=2.0)
Fig.Points(Model.swarm, Model.materialField, fn_size=3.0)
Fig.show()