import numpy as np
import xarray as xr

def MOCpsi(vh, vmsk=None):
  """Sums 'vh' zonally and cumulatively in the vertical to yield an overturning stream function, psi(y,z)."""
  shape = list(vh.shape); shape[-3] += 1
  psi = np.zeros(shape[:-1])
  if len(shape)==3:
    for k in range(shape[-3]-1,0,-1):
      if vmsk is None: psi[k-1,:] = psi[k,:] - vh[k-1].sum(axis=-1)
      else: psi[k-1,:] = psi[k,:] - (vmsk*vh[k-1]).sum(axis=-1)
  else:
    for n in range(shape[0]):
      for k in range(shape[-3]-1,0,-1):
        if vmsk is None: psi[n,k-1,:] = psi[n,k,:] - vh[n,k-1].sum(axis=-1)
        else: psi[n,k-1,:] = psi[n,k,:] - (vmsk*vh[n,k-1]).sum(axis=-1)
  return psi

def get_z(rg, depth, var_name):
  """Returns 3d interface positions from netcdf group rg, based on dimension data for variable var_name"""
  if 'e' in rg.variables: # First try native approach
    if len(rg.variables['e'])==3: return rg.variables['e'][:]
    elif len(rg.variables['e'])==4: return rg.variables['e'][0]
  if var_name not in rg.variables: raise Exception('Variable "'+var_name+'" not found in netcdf file')
  if len(rg.variables[var_name].shape)<3: raise Exception('Variable "'+var_name+'" must have 3 or more dimensions')
  try: vdim = rg.variables[var_name].dimensions[-3]
  # handle xarray dataset, dimensions = dims
  except: vdim = rg.variables[var_name].dims[-3]
  if vdim not in rg.variables: raise Exception('Variable "'+vdim+'" should be a [CF] dimension variable but is missing')
  #if 'edges' in rg.variables[vdim].ncattrs():
  try: zvar = getattr(rg.variables[vdim],'edges')
  except:
    if 'zw' in rg.variables: zvar = 'zw'
    elif 'zl' in rg.variables: zvar = 'zl'
    elif 'z_l' in rg.variables: zvar = 'z_l'
    else: raise Exception('Cannot figure out vertical coordinate from variable "'+var_name+'"')

  print(zvar)
  if not len(rg.variables[zvar].shape)==1: raise Exception('Variable "'+zvar+'" was expected to be 1d')
  if type(rg) == xr.core.dataset.Dataset:
    zw = rg[zvar][:].data
  else:
    zw = rg.variables[zvar][:]
  Zmod = np.zeros((zw.shape[0], depth.shape[0], depth.shape[1] ))
  print(zw[:])
  for k in range(zw.shape[0]):
    Zmod[k] = np.minimum( depth, abs(zw[k]) )
  return Zmod


def findExtrema(y, z, psi, min_lat=-90., max_lat=90., min_depth=0., mult=1., plot = True):
  psiMax = mult*np.amax( mult * np.ma.array(psi)[(y>=min_lat) & (y<=max_lat) & (z>min_depth)] )
  idx = np.argmin(np.abs(psi-psiMax))
  (j,i) = np.unravel_index(idx, psi.shape)
  return j,i,psi[j,i]