#!/usr/bin/python

import numpy
from pyOM import pyOM

class pyOM_cdf_scipy(pyOM):
   """
   pyOM with snapshot output in netcdf format using scipy
   """
   def __init__(self):
     pyOM.__init__(self)
     # try to load module with netcdf bindings
     try:
       from scipy.io.netcdf import NetCDFFile as NF
     except ImportError:
       from scipy.io.netcdf import netcdf_file as NF
         
     self.NF=NF
     self._spval = -1e33*numpy.ones( (1,),'f')  
     return
  
   def diagnose(self):
     """ diagnose the model variables
     """
     pyOM.diagnose(self)
     self.write_cdf()
     return

   def define_grid_cdf(self,fid):
       M=self.fortran.pyom_module         # fortran module with model variables
       
       fid.createDimension('xt',int(M.nx))
       fid.createDimension('xu',int(M.nx))
       fid.createDimension('yt',int(M.ny))
       fid.createDimension('yu',int(M.ny))
       fid.createDimension('zt',int(M.nz))
       fid.createDimension('zw',int(M.nz))
       fid.createDimension('Time',None)
       
       xt=fid.createVariable('xt','f',('xt',) )
       xt.long_name='Zonal coordinate on T grid'; xt.units='m'
       xt[:]=M.xt.astype('f')
  
       xu=fid.createVariable('xu','f',('xu',) )
       xu.long_name='Zonal coordinate on U grid'; xu.units='m'
       xu[:]=M.xu.astype('f')
  
       yt=fid.createVariable('yt','f',('yt',) )
       yt.long_name='Meridional coordinate on T grid'; yt.units='m'
       yt[:]=M.yt.astype('f')
  
       yu=fid.createVariable('yu','f',('yu',) )
       yu.long_name='Meridional coordinate on V grid'; yu.units='m'
       yu[:]=M.yu.astype('f')
  
       zt=fid.createVariable('zt','f',('zt',) )
       zt.long_name='Vertical coordinate on T grid'; zt.units='m'
       zt[:]=M.zt.astype('f')

       zw=fid.createVariable('zw','f',('zw',) )
       zw.long_name='Vertical coordinate on W grid'; zw.units='m'
       zw[:]=M.zw.astype('f')
       return
    
   
   def init_cdf(self,fid):
     """ intitialize netcdf diagnostics
     """
     M=self.fortran.pyom_module         # fortran module with model variables
     
       
     self.define_grid_cdf(fid)
     Time=fid.createVariable('Time','f',('Time',) ) 
     Time.long_name = 'Time since start'; Time.units='Seconds'
   
     id=fid.createVariable('b','f',('Time','zt','yt','xt') )
     id.long_name = 'Buoyancy'; id.units='m/s^2'
     id.missing_value = self._spval
  
     id=fid.createVariable('p','f',('Time','zt','yt','xt') )
     id.long_name = 'Pressure'; id.units='m^2/s^2'
  
     id=fid.createVariable('u','f',('Time','zt','yt','xu') )
     id.long_name = 'Zonal velocity'; id.units='m/s'
  
     id=fid.createVariable('v','f',('Time','zt','yu','xt') )
     id.long_name = 'Meridional velocity'; id.units='m/s'
  
     id=fid.createVariable('w','f',('Time','zw','yt','xt') )
     id.long_name = 'Vertical velocity'; id.units='m/s'
       
     id=fid.createVariable('K_b','f',('Time','zw','yt','xt') )
     id.long_name = 'Vertical diffusivity'; id.units='m^2/s'
       
     id=fid.createVariable('taux','f',('Time','yt','xu') )
     id.long_name = 'Zonal wind stress'; id.units='m^2/s^2'
  
     id=fid.createVariable('tauy','f',('Time','yt','xu') )
     id.long_name = 'Meridional wind stress'; id.units='m^2/s^2'
       
     id=fid.createVariable('surface_flux','f',('Time','yt','xt') )
     id.long_name = 'Surface buoyancy flux'; id.units='m^2/s^4'
  
     if M.enable_vert_friction_trm:
        id=fid.createVariable('A_trm','f',('Time','zw','yt','xt') )
        id.long_name = 'Vertical viscosity'; id.units='m^2/s'
       
     if M.enable_back_state:
         id=fid.createVariable('back','f',('zt','yt','xt') )
         id.long_name = 'Background Buoyancy'; id.units='m/s^2'
         id=fid.createVariable('u0','f',('zt','yt','xu') )
         id.long_name = 'Background zonal flow'; id.units='m/s^2'
     return


   def write_cdf(self):
     M=self.fortran.pyom_module         # fortran module with model variables
     if hasattr(self,'NF'):
       self.fortran.pe0_recv_3d(M.b[:,:,:,M.tau-1])
       self.fortran.pe0_recv_3d(M.u[:,:,:,M.tau-1])
       self.fortran.pe0_recv_3d(M.v[:,:,:,M.tau-1])
       self.fortran.pe0_recv_3d(M.w[:,:,:,M.tau-1])
       self.fortran.pe0_recv_3d(M.p_full[:,:,:,M.tau-1])
       self.fortran.pe0_recv_2d(M.surface_taux)
       self.fortran.pe0_recv_2d(M.surface_tauy)
       self.fortran.pe0_recv_2d(M.surface_flux)
       self.fortran.pe0_recv_3d(M.k_b)
       if M.enable_vert_friction_trm:
           self.fortran.pe0_recv_3d(M.a_trm)

           
       if M.my_pe == 0:   
          file=('snap_%10d.cdf'% self.itt).replace(' ','0')
          print ' writing to file ',file
          fid= self.NF(file,'w')
          self.init_cdf(fid)
          
          tid=fid.variables['Time'];
          i=0
          tid[i]=self.time
          a=M.b[:,:,:,M.tau-1].astype('f')
          b=self._spval*numpy.ones(a.shape,'f')
          fid.variables['b'][i,:,:,:]  = numpy.where(M.maskt[:,:,:]==0.,b,a).transpose().astype('f')
          fid.variables['u'][i,:,:,:]  = M.u[:,:,:,M.tau-1].transpose().astype('f')
          fid.variables['v'][i,:,:,:]  = M.v[:,:,:,M.tau-1].transpose().astype('f')
          fid.variables['w'][i,:,:,:]  = M.w[:,:,:,M.tau-1].transpose().astype('f')
          fid.variables['p'][i,:,:,:]  = M.p_full[:,:,:,M.tau-1].transpose().astype('f')
          fid.variables['taux'][i,:,:] = M.surface_taux[:,:].transpose().astype('f')
          fid.variables['tauy'][i,:,:] = M.surface_tauy[:,:].transpose().astype('f')
          fid.variables['surface_flux'][i,:,:] = M.surface_flux[:,:].transpose().astype('f')
          fid.variables['K_b'][i,:,:,:]  = M.k_b[:,:,:].transpose().astype('f')
          if M.enable_vert_friction_trm:
               fid.variables['A_trm'][i,:,:,:] = M.a_trm[:,:,:].transpose().astype('f')
          if M.enable_back_state:
            fid.variables['back'][:,:,:] = M.back[:,:,:,M.tau-1].transpose().astype('f')
            fid.variables['u0'][:,:,:]   = M.u0[:,:,:].transpose().astype('f')
          fid.close()
     return   





if __name__ == "__main__":
   print 'I will do nothing'
