########################################################################
# Stratigraphic columns
# ~~~~~~~~~~~~~~~~~~~~~~~
# We define the stratigraphic column using a nested dictionary

stratigraphic_column = {}
stratigraphic_column['strati2'] = {}
stratigraphic_column['strati2']['unit1'] = {'min': 1, 'max': 10, 'id': 0}
stratigraphic_column['strati'] = {}
stratigraphic_column['strati']['unit2'] = {'min': -60, 'max': 0, 'id': 1}
stratigraphic_column['strati']['unit3'] = {'min': -250, 'max': -60, 'id': 2}
stratigraphic_column['strati']['unit4'] = {'min': -330, 'max': -250, 'id': 3}
stratigraphic_column['strati']['unit5'] = {
    'min': -np.inf,
    'max': -330,
    'id': 4
}

########################################################
# Adding stratigraphic column to the model
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# The stratigraphic column can be added to the geological model. Allowing
# for the `model.evaluate_model(xyz)` function to be called.

model.set_stratigraphic_column(stratigraphic_column)

viewer = LavaVuModelViewer(model)
viewer.add_model(cmap='tab20')
viewer.rotate([-85.18760681152344, 42.93233871459961, 0.8641873002052307])
viewer.display()
view.nsteps = np.array([50, 50, 50])

for sg in model.feature_name_index:
    if ('super' in sg):
        view.add_data(model.features[model.feature_name_index[sg]])
view.nelements = 1e5
view.add_model_surfaces(function=function, filename=filename, faults=False)
view.nelements = 1e6
view.add_model_surfaces(function=function,
                        filename=filename,
                        strati=False,
                        displacement_cmap='rainbow')
view.lv.webgl(vtk_path + model_name)
view.nsteps = np.array([200, 200, 200])

view.add_model()

view.lv.control.Range('alpha', label="Global Opacity")
view.lv.control.DualRange(['xmin', 'xmax'],
                          label="x clip",
                          step=0.01,
                          values=[0.0, 1.0])
view.lv.control.DualRange(['ymin', 'ymax'],
                          label="y clip",
                          step=0.01,
                          values=[0.0, 1.0])
view.lv.control.DualRange(['zmin', 'zmax'],
                          label="z clip",
                          step=0.01,
                          values=[0.0, 1.0])
view.lv.control.Range(command='background', range=(0, 1), step=0.1, value=0.8)