#!/usr/bin/env python
# coding: utf-8

# In[162]:


import underworld.function as fn
from underworld import UWGeodynamics as GEO
from underworld import visualisation as vis
import numpy as np
import gflex
u = GEO.UnitRegistry


# In[163]:


half_rate = 10 * u.millimeter / u.year
model_length = 44.8e3 * u.meter
surfaceTemp = 273.15 * u.degK
baseModelTemp = 1603.15 * u.degK
bodyforce = 2700 * u.kilogram / u.metre**3 * 9.81 * u.meter / u.second**2
rigidbasedensity = 4000. * u.kilogram / u.metre**3
velocity = 0.4 * u.centimeter / u.year


# In[164]:


Model = GEO.Model(elementRes=(90, 30),
                  minCoord=(0. * u.kilometer, -12. * u.kilometer),
                  maxCoord=(48. * u.kilometer, 4. * u.kilometer),
                  gravity=(0.0, -9.81 * u.meter / u.second**2))


# In[165]:


Model.diffusivity = 9e-7 * u.metre**2 / u.second
Model.capacity    = 1000. * u.joule / (u.kelvin * u.kilogram)


# In[166]:


air = Model.add_material(name="Air", shape=GEO.shapes.Layer(top=Model.top, bottom=0 * u.kilometer))
air.density = 1. * u.kilogram / u.metre**3
air.diffusivity = 1e-6 * u.metre**2 / u.second
air.capacity = 1000. * u.joule / (u.kelvin * u.kilogram)
Loose_Sediment = Model.add_material(name="Loose_Sediment", shape=GEO.shapes.Layer2D(top=0. * u.kilometer, bottom=-0.2 * u.kilometer))
Loose_Sediment.radiogenicHeatProd = 7.67e-7 * u.watt / u.meter**3
Loose_Sediment.density  = 2000. * u.kilogram / u.metre**3


# In[167]:


Left_Top = Model.add_material(name="Left_Top_Layer", shape=GEO.shapes.Polygon([(0,-0.20*u.kilometer), (32*u.kilometer,-0.20*u.kilometer), (32*u.kilometer,-2.60*u.kilometer), (0,-2.60*u.kilometer)]))
Left_Middle = Model.add_material(name="Left_Middle_Layer", shape=GEO.shapes.Polygon([(0,-2.60*u.kilometer), (32*u.kilometer,-2.60*u.kilometer), (32*u.kilometer,-7.40*u.kilometer), (0,-7.40*u.kilometer)]))
Left_Bottom = Model.add_material(name="Left_Bottom_Layer", shape=GEO.shapes.Polygon([(0,-7.40*u.kilometer), (32*u.kilometer,-7.40*u.kilometer), (32*u.kilometer,-12*u.kilometer), (0,-12*u.kilometer)]))

Right_Top = Model.add_material(name="Left_Top_Layer", shape=GEO.shapes.Polygon([(32*u.kilometer,-0.20*u.kilometer), (48*u.kilometer,-0.20*u.kilometer), (48*u.kilometer,-2.60*u.kilometer), (32*u.kilometer,-2.60*u.kilometer)]))
Right_Middle = Model.add_material(name="Left_Middle_Layer", shape=GEO.shapes.Polygon([(32*u.kilometer,-2.60*u.kilometer), (48*u.kilometer,-2.60*u.kilometer), (48*u.kilometer,-5.0*u.kilometer), (32*u.kilometer,-5.0*u.kilometer)]))
Right_Bottom = Model.add_material(name="Left_Bottom_Layer", shape=GEO.shapes.Polygon([(32*u.kilometer,-5.0*u.kilometer), (48*u.kilometer,-5.0*u.kilometer), (48*u.kilometer,-12*u.kilometer), (32*u.kilometer,-12*u.kilometer)]))


# In[168]:



#Left_Top = Model.add_material(name="Left_Top_Layer", shape=GEO.shapes.Layer2D(top=-0.2* u.kilometer, bottom=-2.4 * u.kilometer,maxCoord=(32. * u.kilometer, -0.2* u.kilometer))
#Left_Top = Model.add_material(name="Left_Top_Layer", shape=GEO.shapes.Layer2D(top=-0.2* u.kilometer, bottom=-2.4 * u.kilometer,maxX=32.* u.kilometer))
#Left_Middle = Model.add_material(name="Left_Middle_Layer", shape=GEO.shapes.Layer2D(top=Left_Top.bottom, bottom=-7.2 * u.kilometer, maxX=32.* u.kilometer))
#Left_Bottom = Model.add_material(name="Left_Bottom_Layer", shape=GEO.shapes.Layer2D(top=Left_Middle.bottom, bottom=Model.bottom, maxX=32.* u.kilometer))

#right_top_layer_shape = fn.shape.Box(minX=32. * u.kilometer, maxX=48. * u.kilometer, minY=-12 * u.kilometer, maxY=-0.2 * u.kilometer)
#Right_Top = Model.add_material(name="Right_Top_Layer", shape=GEO.shapes.Layer2D(top=Model.top, bottom=-2.4  * u.kilometer, minX=32.* u.kilometer))
#Right_Middle = Model.add_material(name="Right_Middle_Layer", shape=GEO.shapes.Layer2D(top=Right_Top.bottom, bottom=-7.2 * u.kilometer, minX=32.* u.kilometer))
#Right_Bottom = Model.add_material(name="Right_Bottom_Layer", shape=GEO.shapes.Layer2D(top=Right_Middle.bottom, bottom=Model.bottom, minX=32.* u.kilometer))


# In[169]:


# Create a numeric mapping with float values
#value_map = dict(zip([Left_Top.index, Left_Middle.index, Left_Bottom.index, Right_Top.index, Right_Middle.index, Right_Bottom.index, air.index,Loose_Sediment.index], 
 #                   [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]))

# Make sure a value for the key 0 (or the default material) is included
#if 0 not in value_map:
  #  value_map[0] = 8.0  # Or any other unused float value

# Create a color function with a default value
#colour_fn = fn.branching.map(fn_key=Model.materialField, mapping=value_map, fn_default=8.0)

# Create a color list based on the order in value_map
#colour_list = ["blue", "green", "red", "cyan", "magenta", "yellow", "white", "grey", "black"]  
# Black is for the default material
# Visualize using the color function and color list
#Fig = vis.Figure(figsize=(900,300))
#Fig.Points(Model.swarm,Model.materialField, colour_fn, fn_size=2.0, colours=colour_list)
#Fig.show()


# In[185]:


from underworld import visualisation as vis
Fig = vis.Figure(figsize=(900,300))
Fig.Points(Model.swarm, Model.materialField, fn_size=2.0, colours=
          'white black yellow green blue green yellow blue green blue ')
Fig.show()


# In[171]:


air.density = 1. * u.kilogram / u.metre**3
Loose_Sediment.density  = 2000. * u.kilogram / u.metre**3
Left_Top.density = 2800. * u.kilogram / u.metre**3
Left_Middle.density= 3200. * u.kilogram / u.metre**3
Left_Bottom.density= 3200. * u.kilogram / u.metre**3

Right_Top.density = 2800. * u.kilogram / u.metre**3
Right_Middle.density= 3200. * u.kilogram / u.metre**3
Right_Bottom.density= 3200. * u.kilogram / u.metre**3


# In[172]:


rh = GEO.ViscousCreepRegistry()

Model.minViscosity = 5e18 * u.pascal * u.second
Model.maxViscosity = 5e23 * u.pascal * u.second

air.viscosity = 5e18 * u.pascal * u.second
Loose_Sediment.viscosity = 1e19 * u.pascal * u.second
Left_Top.viscosity = 1e21 * u.pascal * u.second
Left_Middle.viscosity = 1e22 * u.pascal * u.second
Left_Bottom.viscosity = 1e22 * u.pascal * u.second

Right_Top.viscosity = 1e22 * u.pascal * u.second
Right_Middle.viscosity = 1e23 * u.pascal * u.second
Right_Bottom.viscosity = 1e23 * u.pascal * u.second


# In[173]:


Left_Top.elasticity = GEO.Elasticity(shear_modulus=8e10 * u.pascal,
                        observation_time=20000 * u.year)
Left_Middle.elasticity = GEO.Elasticity(shear_modulus=1e11 * u.pascal,
                        observation_time=20000 * u.year)
Left_Bottom.elasticity = GEO.Elasticity(shear_modulus=1e11 * u.pascal,
                        observation_time=20000 * u.year)

Right_Top.elasticity = GEO.Elasticity(shear_modulus=8e10 * u.pascal,
                        observation_time=20000 * u.year)
Right_Middle.elasticity = GEO.Elasticity(shear_modulus=1e11 * u.pascal,
                        observation_time=20000 * u.year)
Right_Bottom.elasticity = GEO.Elasticity(shear_modulus=1e11 * u.pascal,
                        observation_time=20000 * u.year)


# In[174]:


#Model.init_model()

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

#Model.set_heatFlowBCs(bottom=(-0.044 * u.watt / u.metre**2, Beam))


# In[192]:


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


# In[193]:


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
   
    DissipationField.data[:] = Dissipation.evaluate(Model.swarm)
    
### for pre solve update of swarm variable
Model.pre_solve_functions["A-pre"] = updateVD


# In[194]:


#Model.init_model()
npoints = 100
x = np.linspace(GEO.nd(Model.minCoord[0]), GEO.nd(Model.maxCoord[0]), 100)
y = 0.

#surface_tracers = Model.add_passive_tracers(name="Surface", vertices=[x,y])
#moho_tracers = Model.add_passive_tracers(name="Moho", vertices=[x,y-GEO.nd(24.*u.kilometer)])


npoints = int(Model.maxCoord[0].to(u.kilometer).magnitude)
x_surface = np.linspace(GEO.nd(Model.minCoord[0]), GEO.nd(Model.maxCoord[0]), npoints)
y_surface = -0.01 * u.kilometer


# In[195]:


def Hillslope_diffusion_basic():
    from scipy.interpolate import interp1d
    from scipy.interpolate import InterpolatedUnivariateSpline

    x = GEO.dimensionalise(surface_tracers_erosion.data[:,0], u.meter).magnitude
    z = GEO.dimensionalise(surface_tracers_erosion.data[:,1], u.meter).magnitude

    dx = (Model.maxCoord[0].to(u.meter).magnitude)/npoints

    total_time = (GEO.dimensionalise(Model._dt, u.year)).magnitude

    D = ((1.0e3 * u.meter**2 / u.year).to(u.meter**2 / u.year)).magnitude

    dt = min((0.2 * dx * dx / D), total_time)

    nts = int(round(total_time/dt))

    print('total time:', total_time, 'timestep:', dt, 'No. of its:', nts)

    z_orig = z.copy()

    for i in range(nts):
        qs = -D * np.diff(z)/dx
        dzdt = -np.diff(qs)/dx
        z[1:-1] += dzdt*dt


    x_nd = GEO.nd(x*u.meter)

    z_nd = GEO.nd(z*u.meter)


    if x_nd.shape[0] > 0.:
        f1 = interp1d(x_nd, z_nd, fill_value='extrapolate', kind='nearest')

        y_eroded_surface = f1(x_nd)

        y_eroded_surface[x_nd < GEO.nd(Update_material_LHS_Length*u.kilometer )] = 0.

        surface_tracers_erosion.data[:,1] = y_eroded_surface


# In[211]:


GEO.rcParams["default.outputs"].append("projStrainTensor")
GEO.rcParams["default.outputs"].append("projStressTensor")
from underworld.scaling import units as u
#colour_fn = fn.branching.map(fn_key=Model.materialField, mapping=value_map )
#Model.solver.set_inner_method("mumps")
#Model.solver.set_penalty(1e6)
Model.run_for(1000.* u.year, checkpoint_interval=5000.* u.year)


# In[216]:


from underworld import visualisation as vis
Fig = vis.Figure(figsize=(900,300))
Fig.Points(Model.swarm, Model.materialField, fn_size=2.0, colours=
          'white black yellow green blue green yellow blue green blue ')
Fig.show()


# In[213]:


Fig = vis.Figure(figsize=(1200,400))
Fig.Points(Model.swarm, DissipationField, fn_size=2.0)
Fig.show()


# In[214]:


Fig = vis.Figure(figsize=(900,300), title="strainRate Field", quality=3)
Fig.Surface(Model.mesh, Model.strainRate_2ndInvariant, logScale=True)
Fig.show()


# In[215]:


Fig = vis.Figure(figsize=(900,300), title="projStressField Field", quality=3)
Fig.Surface(Model.mesh, Model.projStressField, logScale=True,colours="coolwarm")
Fig.show()


# In[ ]: