import numpy as np
from scipy.interpolate import griddata

def convert_pcm(field,int):
    """
    Convert a field into PCM field
    field: 1d or 3d array - variable to convert
    int: 1d or 3d array - interface depth of the variable
    """
    fdim=np.ndim(field)
    idim=np.ndim(int)
    nk=field.shape[0]
    print(fdim,idim)

    if (fdim == 1):
        field_pcm=np.zeros([nk*3])
        k1=0
        for k in np.arange(nk):
           field_pcm[k1  ] = field[k]
           field_pcm[k1+1] = field[k]
           field_pcm[k1+2] = field[k]
           k1=k1+3

    if (fdim == 3):
        jdm=field.shape[1];idm=field.shape[2]
        field_pcm=np.zeros([nk*3,jdm,idm])
        k1=0
        for k in np.arange(nk):
           field_pcm[k1  ,:,:] = field[k,:,:]
           field_pcm[k1+1,:,:] = field[k,:,:]
           field_pcm[k1+2,:,:] = field[k,:,:]
           k1=k1+3

    if (idim == 1):
        int_pcm=np.zeros([nk*3])
        k1=0
        for k in np.arange(nk):
           int_pcm[k1  ] = int[k]
           int_pcm[k1+1] = 0.5*(int[k] + int[k+1])
           int_pcm[k1+2] = int[k+1]
           k1=k1+3

    if (idim == 3):
        int_pcm=np.zeros([nk*3,jdm,idm])
        k1=0
        for k in np.arange(nk):
           int_pcm[k1  ,:,:] = int[k,:,:]
           int_pcm[k1+1,:,:] = 0.5*(int[k,:,:] + int[k+1,:,:])
           int_pcm[k1+2,:,:] = int[k+1,:,:]
           k1=k1+3

    return field_pcm,int_pcm



def adjust_minmax(field,levels):
   """
   Adjust min and max of an array to get closed contours at edge values
   Equivent to saturate colors
   field: list of floats 2D - field to be changed
   levels: list of floats 1D - levels to be drawn
   """
   # find min and max of field and levels
   field_min=np.nanmin(field);field_max=np.nanmax(field)
   print(field_min, field_max)
   vmin=np.nanmin(levels) ; vmax=np.nanmax(levels)
   # get the difference between min/max field and min/max levels
   diffn=vmin-field_min
   diffx=vmax-field_max
   # find at waht level is the field min and max
   indic_min=np.where(levels<=field_min)
   indic_max=np.where(levels>=field_max)
   # replace min/max field by the lower/higher level value respectively
   if (diffn < 0.):
       field[field <= field_min]=np.nanmax(levels[indic_min])

   if (diffx > 0.):
       field[np.logical_and(field>=field_max,field<1e5)]=np.nanmin(levels[indic_max])

   return field


def interp2plot(field,lonin,latin,reso,method='linear',ocean=True):
   """
   Interpolate curvilinear grid (typical HYCOM global grid) to regular grid
   field: list of floats 2D - field to be interpolated
   lonin: list of floats 2D - longitude of input grid
   latin: list of floats 2D - latitude of input grid
   reso: float - resolution of the output regular grid
   method: string - 'linear','nearest','cubic' from scipy griddata
   ocean: logical - True for ocean only or False for all array
   """
   # get lon and lat
   lonin[lonin >= 360]=lonin[lonin >= 360]-360.
   jdm=lonin.shape[0] ; idm=lonin.shape[1]
   lon1=lonin[0:jdm-1,:] ## last row always bogus
   lat1=latin[0:jdm-1,:] ## last row always bogus
   field1=field[0:jdm-1,:] ## last row always bogus
   ## extract only ocean points
   if (ocean):
      datain=field1[field1 < 1e10]
      lonin=lon1[field1 < 1e10]
      latin=lat1[field1 < 1e10]
   else:
      datain=np.reshape(field1,idm*(jdm-1))
      lonin=np.reshape(lon1,idm*(jdm-1))
      latin=np.reshape(lat1,idm*(jdm-1))

## create regular grid with resolution reso
   grid_x,grid_y=np.mgrid[0:360:reso,-80:90:reso]
   ## Interpolate on regular grid using chosen method !
   dataout=griddata((lonin,latin), datain, (grid_x, grid_y), method=method,fill_value=np.nan)
   return dataout,grid_x,grid_y

def interp2plot2(field,lonin,latin,lonout,latout,method='linear'):
   """
   Interpolate curvilinear grid (typical HYCOM global grid) to regular grid
   field: list of floats 2D - field to be interpolated
   lonin: list of floats 2D - longitude of input grid
   latin: list of floats 2D - latitude of input grid
   reso: float - resolution of the output regular grid
   method: string - 'linear','nearest','cubic' from scipy griddata
   """
   # get lon and lat
   lonin[lonin >= 360.]=lonin[lonin >= 360.]-360.
   lonin[lonin >= 180.]=lonin[lonin >= 180.]-360.
   # make lonin and latin 2-d if necessary
   lonin1=lonin
   latin1=latin
   if (lonin.ndim == 1):
      idm=lonin.shape[0] ; jdm=latin.shape[0]
      lonin1=np.zeros([jdm,idm])
      latin1=np.zeros([jdm,idm])
      for j in np.arange(jdm):
          lonin1[j,:]=lonin
      for i in np.arange(idm):
          latin1[:,i]=latin

    ## extract only ocean points
   datain=field[np.where(field < 1e10)]
   lonin1=lonin1[np.where(field < 1e10)]
   latin1=latin1[np.where(field < 1e10)]
   ## Interpolate on regular grid using chosen method !
   dataout=griddata((lonin1,latin1), datain, (lonout, latout), method=method)
   return dataout

def bytescl(array, max=None , min=None , nan=0, top=255 ):

   # see http://star.pst.qub.ac.uk/idl/BYTSCL.html
   # note that IDL uses slightly different formulae for bytscaling floats and ints.
   # here we apply only the FLOAT formula...

   if max is None: max = np.nanmax(array)
   if min is None: min = np.nanmin(array)

   #return (top+0.9999)*(array-min)/(max-min)
   #return np.max(np.min( ((top+0.9999)*(array-min)/(max-min)).astype(np.int16), top),0)
   idm=array.shape[1] ; jdm=array.shape[0]
   results=np.zeros([jdm,idm])
   for j in np.arange(jdm):
      for i in np.arange(idm):
         if (np.isnan(array[j,i] == False)):
            results[j,i]=np.nanmax(np.nanmin( int((top+0.9999)*(array[j,i]-min)/(max-min)), top), 0)


   return results