#!/usr/bin/python

import numpy
from pyOM import pyOM

class pyOM_mat(pyOM):
   """
   pyOM with output in matlab format
   """
   def __init__(self):
     pyOM.__init__(self)
     import scipy.io
     self.smat = scipy.io.savemat
     return
  
   def diagnose(self):
     """ diagnose the model variables
     """
     pyOM.diagnose(self)
     self.write_mat()
     return

   def write_mat(self):
     M=self.fortran.pyom_module         # fortran module with model variables
     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_{:0=10d}.mat'.format(self.itt)
        file=('snap_%10d.mat'% self.itt).replace(' ','0')
        print ' writing to file ',file
        d={'xt':M.xt[:],'yt':M.yt[:],'zt':M.zt[:],'xu':M.xu[:],'yu':M.yu[:],'zw':M.zw[:],'time':self.time}
        d['b']=M.b[:,:,:,M.tau-1]
        d['u']=M.u[:,:,:,M.tau-1]
        d['v']=M.v[:,:,:,M.tau-1]
        d['w']=M.w[:,:,:,M.tau-1]
        d['p']=M.p_full[:,:,:,M.tau-1]
        d['taux']=M.surface_taux
        d['tauy']=M.surface_tauy
        d['flux']=M.surface_flux
        d['k_b']=M.k_b
        if M.enable_vert_friction_trm: d['a_trm']=M.a_trm
        self.smat(file,d,appendmat=False)#,oned_as='row')
     return   


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