import numpy as np
import utm

import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
params = {
    'image.interpolation': 'nearest',
    'image.cmap': 'gray',
    'savefig.dpi': 300,  # to adjust notebook inline plot size
    'axes.labelsize': 18, # fontsize for x and y labels (was 10)
    'axes.titlesize': 18,
    'font.size': 18,
    'legend.fontsize': 18,
    'xtick.labelsize': 18,
    'ytick.labelsize': 18,
    'text.usetex':False
}
matplotlib.rcParams.update(params)
from mpl_toolkits.axes_grid1.inset_locator import inset_axes

velModFile = "./FinalVpVsLongValleyDAS2023.npz"
with np.load(velModFile) as velFile:
    Vp = velFile["vp"]
    Vs = velFile["vs"]
    dy, dx, dz = velFile["ds"]
    ny, nx, nz = Vp.shape
    minLon, maxLon, minLat, maxLat = velFile['boundbox']
    # Space axes and UTM parameters for lat-lon conversion
    oy, ox, oz = 0.0, 0.0, velFile['oz']
    zAxis = np.linspace(oz, oz+(nz-1)*dz, nz)
    xAxis = np.linspace(ox, ox+(nx-1)*dx, nx)
    yAxis = np.linspace(oy, oy+(ny-1)*dy, ny)
    latAxis = np.linspace(minLat, maxLat, ny)
    lonAxis = np.linspace(minLon, maxLon, nx)
    _,_,zone_n,zone_lt = utm.from_latlon(0.5*(minLat+maxLat),0.5*(minLon+maxLon))
    NLL_origin_utm = utm.from_latlon(minLat,minLon,zone_n,zone_lt)

z_slice = -1.0 #[km]
idx_z = int(np.floor(((z_slice - oz) / dz) + 0.5))
textstr="Depth: %s [km]" % z_slice
# Index of the depth slice to plot
depth_slice = Vs[:,:,idx_z]
minval=Vs[:,:,idx_z].min()
maxval=Vs[:,:,idx_z].max()
# Plotting depth slice initial model
fig, ax = plt.subplots(figsize=(14,14))
im = ax.imshow(depth_slice, cmap=plt.get_cmap("jet_r"), 
               extent=[minLon,maxLon,minLat,maxLat], vmin=minval, 
               vmax=maxval, origin="lower", aspect=1.0)
ax.set_xlabel("Longitude [deg]")
ax.set_ylabel("Latitude [deg]")
ax.grid()
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="5%", pad=0.2)
cbar = plt.colorbar(im, orientation="vertical", cax=cax, format='%.1f')
cbar.set_label('Vs [km/s]')
# Depth slice box
props = dict(boxstyle='square', facecolor='white', alpha=1.0)
ax.text(0.007, 0.992, textstr, transform=ax.transAxes, fontsize=18,
        verticalalignment='top', bbox=props)