import cmor,numpy,sys,os
from time import localtime, strftime
today = strftime("%Y%m%d", localtime())


try:
    import cdms2
    cdms2.setNetcdfShuffleFlag(0)
    cdms2.setNetcdfDeflateFlag(0)
    cdms2.setNetcdfDeflateLevelFlag(0)
except:
    print "This test code needs a recent cdms2 interface for i/0"
    sys.exit()

level = 3
shuffle = 1
cmor.setup(inpath="Tables",set_verbosity=cmor.CMOR_NORMAL, netcdf_file_action = cmor.CMOR_REPLACE_4, exit_control = cmor.CMOR_EXIT_ON_MAJOR);
cmor.dataset_json("Test/common_user_input.json")

tables=[]
tables.append(cmor.load_table("CMIP6_Amon.json"))
print 'Tables ids:',tables


## read in data, just one slice
f=cdms2.open('data/tas_ccsr-95a.xml')
s=f("tas",time=slice(0,12),squeeze=1)

ntimes = 12
varout='tas'

myaxes=numpy.arange(10)
myvars=numpy.arange(10)
myaxes[0] = cmor.axis(table_entry = 'latitude', 
                      units = 'degrees_north', 
                      coord_vals = s.getLatitude()[:],cell_bounds=s.getLatitude().getBounds())
myaxes[1] = cmor.axis(table_entry = 'longitude', 
                      units = 'degrees_north', 
                      coord_vals = s.getLongitude()[:],cell_bounds=s.getLongitude().getBounds())



myaxes[2] = cmor.axis(table_entry = 'time',
                      units = s.getTime().units,
                      coord_vals = s.getTime()[:],cell_bounds=s.getTime().getBounds())

pass_axes = [myaxes[2],myaxes[0],myaxes[1]]

myvars[0] = cmor.variable( table_entry = varout,
                           units = s.units,
                           axis_ids = pass_axes,
                           original_name = s.id,
                           history = 'no history',
                           comment = 'testing speed'
                           )
cmor.set_deflate(myvars[0], shuffle, 1, level)

import time,MV2
st = time.time()
totcmor=0
totcdms=0
maxcmor=0
mincmor=1000
maxcdms=0
mincdms=1000
c0=st
f=cdms2.open("Test/crap.nc","w")
#print 'Time:',i
print s.filled().shape
cmor.write(myvars[0],s.filled(),ntimes_passed=ntimes)
c=time.time()
print 'cmor write time:',c-c0
totcmor+=c-c0
if maxcmor<c-c0:
    maxcmor=c-c0
if mincmor>c-c0:
    mincmor=c-c0
c0=c
f.write(s,id=varout)
c=time.time()
print 'cdms time:',c-c0
totcdms+=c-c0
if maxcdms<c-c0:
    maxcdms=c-c0
if mincdms>c-c0:
    mincdms=c-c0
c0=c
f.close()
cmor.close()

import cdtime,os
ltime = cdtime.reltime(ntimes-1,'month since 1980').tocomp()
#lcmor = os.stat("CMIP6/CMIP/CSIRO-BOM/NICAM/piControl/r1i1p1f1/Amon/tas/gn/v%s/tas_Amon_piControl_NICAM_r1i1p1f1_gn_197901-197912.nc"%(today))[6]
lcmor = os.stat("CMIP6/ISMIP6/PCMDI/PCMDI-test-1-0/piControl-withism/r11i1p1f1/Amon/tas/gr/v%s/tas_Amon_piControl-withism_PCMDI-test-1-0_r11i1p1f1_gr_197901-197912.nc"%(today))[6]
print 'level:',level,"shuffle:",shuffle
print 'total cmor:',totcmor,mincmor,totcmor/ntimes,maxcmor,lcmor
lcdms = os.stat("Test/crap.nc")[6]
print 'total cdms:',totcdms,mincdms,totcdms/ntimes,maxcdms,lcdms
print 'Size diff:',float(lcmor)/float(lcdms)
print 'speed diff:', totcmor/totcdms

if os.path.exists("summary.txt"):
    f = open("summary.txt")
    s=f.read()
    f.close()
    dic = eval(s)
else:
    dic = {}

dic[(level,shuffle)]=(float(lcmor)/float(lcdms),totcmor/totcdms)

for i in range(10):
    a = dic.get((i,0),"N/A")
    b = dic.get((i,1),"N/A")
    print 'Level: ',i,"no suffle:",a,"shuffle",b
f=open("summary.txt","w")
f.write(repr(dic))
f.close()

