from __future__ import print_function, division

import os

import argparse

import numpy as np

from netCDF4 import Dataset

from wrf import (getvar, interplevel, interpline, vertcross, vinterp, py2round,

CoordPair, ll_to_xy, xy_to_ll, to_np)

VARS_TO_KEEP = ("Times", "XLAT", "XLONG", "XLAT_U", "XLAT_V", "XLONG_U",

"XLONG_V", "U", "V", "W", "PH", "PHB", "T", "P", "PB", "Q2",

"T2", "PSFC", "U10", "V10", "XTIME", "QVAPOR", "QCLOUD",

"QGRAUP", "QRAIN", "QSNOW", "QICE", "MAPFAC_M", "MAPFAC_U",

"MAPFAC_V", "F", "HGT", "RAINC", "RAINSH", "RAINNC")

DIMS_TO_TRIM = ("west_east", "south_north", "bottom_top", "bottom_top_stag",

"west_east_stag", "south_north_stag")

WRF_DIAGS = ["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"]

INTERP_METHS = ["interplevel", "vertcross", "interpline", "vinterp"]

LATLON_METHS = ["xy", "ll"]

def copy_and_reduce(opts):

infilename = opts.filename

outfilename = os.path.expanduser(

os.path.join(opts.outdir, "ci_test_file.nc"))

with Dataset(infilename) as infile, Dataset(outfilename, "w") as outfile:

# Copy the global attributes

outfile.setncatts(infile.__dict__)

# Copy Dimensions

for name, dimension in infile.dimensions.iteritems():

if name in DIMS_TO_TRIM:

if name.find("_stag") > 0:

dimsize = (len(dimension) / 2

if len(dimension) % 2 == 0

else (len(dimension) + 1) / 2)

else:

dimsize = (len(dimension) / 2

if len(dimension) % 2 == 0

else (len(dimension) - 1) / 2)

else:

dimsize = len(dimension)

outfile.createDimension(name, dimsize)

# Copy Variables

for name, variable in infile.variables.iteritems():

if name not in VARS_TO_KEEP:

continue

outvar = outfile.createVariable(name, variable.datatype,

variable.dimensions, zlib=True)

in_slices = tuple(slice(0, dimsize) for dimsize in outvar.shape)

outvar[:] = variable[in_slices]

outvar.setncatts(infile.variables[name].__dict__)

def add_to_ncfile(outfile, var, varname):

dim_d = dict(zip(var.dims, var.shape))

for dimname, dimsize in dim_d.items():

if dimname not in outfile.dimensions:

outfile.createDimension(dimname, dimsize)

fill_value = None

if "_FillValue" in var.attrs:

fill_value = var.attrs["_FillValue"]

ncvar = outfile.createVariable(varname, var.dtype, var.dims,

zlib=True, fill_value=fill_value)

var_vals = to_np(var)

ncvar[:] = var_vals[:]

def make_result_file(opts):

infilename = os.path.expanduser(

os.path.join(opts.outdir, "ci_test_file.nc"))

outfilename = os.path.expanduser(

os.path.join(opts.outdir, "ci_result_file.nc"))

with Dataset(infilename) as infile, Dataset(outfilename, "w") as outfile:

for varname in WRF_DIAGS:

var = getvar(infile, varname)

add_to_ncfile(outfile, var, varname)

for interptype in INTERP_METHS:

if interptype == "interpline":

hts = getvar(infile, "z")

p = getvar(infile, "pressure")

hts_850 = interplevel(hts, p, 850.)

add_to_ncfile(outfile, hts_850, "interplevel")

if interptype == "vertcross":

hts = getvar(infile, "z")

p = getvar(infile, "pressure")

pivot_point = CoordPair(hts.shape[-1] // 2, hts.shape[-2] // 2)

ht_cross = vertcross(hts, p, pivot_point=pivot_point,

angle=90.)

add_to_ncfile(outfile, ht_cross, "vertcross")

if interptype == "interpline":

t2 = getvar(infile, "T2")

pivot_point = CoordPair(t2.shape[-1] // 2, t2.shape[-2] // 2)

t2_line = interpline(t2, pivot_point=pivot_point, angle=90.0)

add_to_ncfile(outfile, t2_line, "interpline")

if interptype == "vinterp":

tk = getvar(infile, "temp", units="k")

interp_levels = [200, 300, 500, 1000]

field = vinterp(infile,

field=tk,

vert_coord="theta",

interp_levels=interp_levels,

extrapolate=True,

field_type="tk",

log_p=True)

add_to_ncfile(outfile, field, "vinterp")

for latlonmeth in LATLON_METHS:

if latlonmeth == "xy":

# Hardcoded values from other unit tests

lats = [22.0, 25.0, 27.0]

lons = [-90.0, -87.5, -83.75]

xy = ll_to_xy(infile, lats[0], lons[0])

add_to_ncfile(outfile, xy, "xy")

else:

# Hardcoded from other unit tests

x_s = np.asarray([10, 50, 90], int)

y_s = np.asarray([10, 50, 90], int)

ll = xy_to_ll(infile, x_s[0], y_s[0])

add_to_ncfile(outfile, ll, "ll")

def main(opts):

copy_and_reduce(opts)

make_result_file(opts)

if __name__ == "__main__":

DEFAULT_FILE = ("/Users/ladwig/Documents/wrf_files/"

"wrf_vortex_multi/moving_nest/"

"wrfout_d02_2005-08-28_12:00:00")

parser = argparse.ArgumentParser(description="Generate conda test files "

"for unit testing.")

parser.add_argument("-f", "--filename", required=False,

default=DEFAULT_FILE,

help="the WRF test file")

parser.add_argument("-o", "--outdir", required=False, default="./",

help="the location for the output files")

opts = parser.parse_args()

main(opts)