import numpy as np
from simpeg import utils
from discretize.utils import mesh_builder_xyz, refine_tree_xyz

# First we need to define the direction of the inducing field
# As a simple case, we pick a vertical inducing field of magnitude
# 50,000nT.
# From old convention, field orientation is given as an
# azimuth from North (positive clockwise)
# and dip from the horizontal (positive downward).
h0_amplitude, h0_inclination, h0_declination = (50000.0, 90.0, 0.0)

# Create a mesh
h = [5, 5, 5]
padDist = np.ones((3, 2)) * 100
nCpad = [2, 4, 2]

# Create grid of points for topography
# Lets create a simple Gaussian topo and set the active cells
[xx, yy] = np.meshgrid(np.linspace(-200.0, 200.0, 50), np.linspace(-200.0, 200.0, 50))

b = 100
A = 50
zz = A * np.exp(-0.5 * ((xx / b) ** 2.0 + (yy / b) ** 2.0))

# We would usually load a topofile
topo = np.c_[utils.mkvc(xx), utils.mkvc(yy), utils.mkvc(zz)]

# Create and array of observation points
xr = np.linspace(-100.0, 100.0, 20)
yr = np.linspace(-100.0, 100.0, 20)
X, Y = np.meshgrid(xr, yr)
Z = A * np.exp(-0.5 * ((X / b) ** 2.0 + (Y / b) ** 2.0)) + 5

xyzLoc = np.c_[utils.mkvc(X.T), utils.mkvc(Y.T), utils.mkvc(Z.T)]

# Build mesh with deprecated functions
old_mesh = mesh_builder_xyz(
    xyzLoc,
    h,
    padding_distance=padDist,
    mesh_type="TREE",
)
old_mesh = refine_tree_xyz(
    old_mesh,
    topo,
    method="surface",
    octree_levels=nCpad,
    octree_levels_padding=nCpad,
    finalize=True,
)

# Build mesh with new methods
new_mesh = mesh_builder_xyz(
    xyzLoc,
    h,
    padding_distance=padDist,
    mesh_type="TREE",
)

new_mesh.refine_surface(topo, padding_cells_by_level=nCpad, finalize=True)

print(f"old_mesh.n_cells: {old_mesh.n_cells}")
print(f"new_mesh.n_cells: {new_mesh.n_cells}")