from __future__ import print_function

import time
from netCDF4 import Dataset
from wrf import getvar, ALL_TIMES, extract_vars

wrf_filenames = ["/Users/ladwig/Documents/wrf_files/wrf_vortex_multi/moving_nest/wrfout_d02_2005-08-28_00:00:00",
                 "/Users/ladwig/Documents/wrf_files/wrf_vortex_multi/moving_nest/wrfout_d02_2005-08-28_12:00:00", 
                 "/Users/ladwig/Documents/wrf_files/wrf_vortex_multi/moving_nest/wrfout_d02_2005-08-29_00:00:00"]

wrfin = [Dataset(x) for x in wrf_filenames]

my_cache = extract_vars(wrfin, ALL_TIMES, ("P", "PB", "PH", "PHB", "T", "QVAPOR", "HGT", "U", "V", "W", "PSFC"))

start = time.time()
for var in ("avo", "eth", "cape_2d", "cape_3d", "ctt", "dbz", "mdbz", 
            "geopt", "helicity", "lat", "lon", "omg", "p", "pressure", 
            "pvo", "pw", "rh2", "rh", "slp", "ter", "td2", "td", "tc", 
            "theta", "tk", "tv", "twb", "updraft_helicity", "ua", "va", 
            "wa", "uvmet10", "uvmet", "z", "cfrac", "zstag", "geopt_stag"):
    v = getvar(wrfin, var, ALL_TIMES)
end = time.time()

print ("Time taken without variable cache: ", (end-start))

start = time.time()
for var in ("avo", "eth",  "cape_2d", "cape_3d", "ctt", "dbz", "mdbz", 
            "geopt", "helicity", "lat", "lon", "omg", "p", "pressure", 
            "pvo", "pw", "rh2", "rh", "slp", "ter", "td2", "td", "tc", 
            "theta", "tk", "tv", "twb", "updraft_helicity", "ua", "va", 
            "wa", "uvmet10", "uvmet", "z", "cfrac", "zstag", "geopt_stag"):
    v = getvar(wrfin, var, ALL_TIMES, cache=my_cache)
end = time.time()

print ("Time taken with variable cache: ", (end-start))