import gempy as gp
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

def extract_domain(sol, unit):
    '''
    Extract domain coordinates from gempy model by unit name
    arguments:
        sol: Gempy solution object.
        unit: id of gempy unit to be extracted
    returns:
        dom_x, dom_y, dom_z: coordinates of domain
    '''

    # new version with rounding, definitely necessary
    rounded_lithblock = sol.lith_block.round(0)
    rounded_lithblock = rounded_lithblock.astype(int)

    # mask by array of input surfaces (by id, can be from different series)
    mask = np.isin(rounded_lithblock, unit)

    # get coordinates by mask
    # krig_lith = sol.lith_block[mask]
    dom_grid = sol.grid.values[mask]
    #    dom_grid = sol.grid.values

    return dom_grid

pd.set_option('precision', 2)
data_path="C:\\Users\\17138\\geoLOGIC\\gridding\\test\\"
path_to_data = data_path

geo_data = gp.create_data('combination',
                          extent=[755589.109, 780838.701, 7034865.1, 7053841.76, -1735.9, -1271.1],
                          resolution=[125, 50, 50],
                          path_o=path_to_data + "input_gempy1_orientations.csv",
                          path_i=path_to_data + "input_gempy1.csv")

geo_data.get_data()

gp.map_stack_to_surfaces(geo_data, {"Strat_Series1": ('surface1'),
                         "Basement_Series":('basement')})

gp.plot_2d(geo_data, direction='y')

interp_data = gp.set_interpolator(geo_data, theano_optimizer='fast_compile')

# %%
sol = gp.compute_model(geo_data)

# %%
# Displaying the result in x and y direction:
#

# %%
gp.plot_2d(geo_data, cell_number=20,
           direction='y', show_data=False, show_boundaries=True)

grid = extract_domain(sol, 1) # extract single unit by id

# Following part is to find lowest z value for each unique xy coordinate

# very much hacked and not very nice, but hopefully does the job
# make zeros array as start
res = np.zeros((1,3))

# loop over unique x and y coordiantes and save minimum
jcntr=0
for i in np.unique(grid[:,0]):
    for j in np.unique(grid[:,1]):
        mask = (grid[:,0]==i)&(grid[:,1]==j)
        
        temp1 = grid[mask]
        trc, trc1=temp1.shape
        if trc!= 0 and trc1 != 0 :
            temp2 = np.where(temp1[:,2] == np.amin(temp1[:,2]))
            temp3 = temp1[temp2]
        else:
            print("no value found for number: ",jcntr)
            jcntr=jcntr+ 1
        # save result
        res = np.vstack((res, temp3))
# remove first zero entry
res = res[1:,:]
fig = plt.figure(figsize=(15, 10))
ax = fig.add_subplot(111)
c = ax.scatter(res[:, 0], res[:, 1], c=res[:, 2], s=3, cmap="viridis")
ax.set_aspect('equal')
plt.colorbar(c, shrink=0.5)
plt.show()