from __future__ import print_function
import numpy as np
from ..utils import Util2d, Util3d, Transient2d, MfList, \
    HeadFile, CellBudgetFile, UcnFile, FormattedHeadFile
from ..mbase import BaseModel
from ..pakbase import Package
from . import NetCdf,netcdf
from . import shapefile_utils



NC_UNITS_FORMAT = {"hk": "{0}/{1}", "sy": "", "ss": "1/{0}", "rech": "{0}/{1}", "strt": "{0}",
                   "wel_flux": "{0}^3/{1}", "top": "{0}", "botm": "{0}", "thickness": "{0}",
                   "ghb_cond": "{0}/{1}^2", "ghb_bhead": "{0}", "transmissivity": "{0}^2/{1}",
                   "vertical_conductance": "{0}/{1}^2", "primary_storage_coefficient": "1/{1}",
                   "horizontal_hydraulic_conductivity": "{0}/{1}", "riv_cond": "1/{1}",
                   "riv_stage": "{0}", "riv_rbot": "{0}", "head":"{0}",
                   "drawdown": "{0}", "cell_by_cell_flow": "{0}^3/{1}", "sy":"{1}/{1}",
                   "prsity": "{1}/{1}", "hani": "{0}/{0}", "al": "{0}/{0}", "drn_elev": "{0}",
                   "drn_cond": "1/{1}", "dz": "{0}", "subsidence": "{0}",
                   "chd_shead": "{0}", "chd_ehead": "{0}", "2D_cumulative_well_flux": "{0}^3/{1}",
                   "3D_cumulative_well_flux": "{0}^3/{1}","vka":"{0}/{1}"}
NC_PRECISION_TYPE = {np.float32: "f4", np.int: "i4", np.int64: "i4", np.int32:"i4"}

NC_LONG_NAMES = {"hk": "horizontal hydraulic conductivity",
                 "vka":"vertical hydraulic conductivity",
                 "sy": "specific yield",
                 "ss": "specific storage",
                 "rech": " recharge",
                 "strt": "starting heads",
                 "wel_flux": "well flux",
                 "top": "model top",
                 "botm": "layer bottom",
                 "thickness": "layer thickness",
                 "ghb_cond": "GHB boundary conductance",
                 "ghb_bhead": "GHB boundary head",
                 "riv_cond": "river bed conductance",
                 "riv_stage": "river stage",
                 "riv_rbot": "river bottom elevation",
                 "drn_elev": "drain elevation",
                 "drn_cond": "drain conductance",
                 "hani": "horizontal anisotropy",
                 "prsity": "porosity",
                 "sconc1": "starting concentration",
                 "ibound": "flow model active array",
                 "icbund": "transport model active array"
                 }


def get_var_array_dict(m):
    vdict = {}
    #for vname in f.var_attr_dict.keys():
    #    vdict[vname] = f.nc.variables[vname][:]
    for attr in m:
        if hasattr(attr,"stress_period_data"):
            array_dict = attr.stress_period_data.array

    return vdict


def ensemble_helper(inputs_filename,outputs_filename,models,add_reals=True,**kwargs):
    """ helper to export an ensemble of model instances.  Assumes
    all models have same dis and sr, only difference is properties and
    boundary conditions.  Assumes model.nam.split('_')[-1] is the
    realization suffix to use in the netcdf variable names
    """
    f_in,f_out = None,None
    for m in models[1:]:
        assert m.get_nrow_ncol_nlay_nper() == models[0].get_nrow_ncol_nlay_nper()
    if inputs_filename is not None:
        f_in = models[0].export(inputs_filename,**kwargs)
        vdict = {}
        vdicts = [models[0].export(vdict,**kwargs)]
        i = 1
        for m in models[1:]:
            suffix = m.name.split('.')[0].split('_')[-1]
            vdict = {}
            m.export(vdict,**kwargs)
            vdicts.append(vdict)
            if add_reals:
                f_in.append(vdict,suffix=suffix)
            i += 1
        mean, stdev = {},{}
        for vname in vdict.keys():
            alist = []
            for vd in vdicts:
                alist.append(vd[vname])
            alist = np.array(alist)
            mean[vname] = alist.mean(axis=0)
            stdev[vname] = alist.std(axis=0)
            mean[vname][vdict[vname]==netcdf.FILLVALUE] = netcdf.FILLVALUE
            stdev[vname][vdict[vname]==netcdf.FILLVALUE] = netcdf.FILLVALUE
            mean[vname][np.isnan(vdict[vname])] = netcdf.FILLVALUE
            stdev[vname][np.isnan(vdict[vname])] = netcdf.FILLVALUE

        if i >= 2:
            if not add_reals:
                f_in.write()
                f_in = NetCdf.empty_like(mean,output_filename=inputs_filename)
                f_in.append(mean,suffix="**mean**")
                f_in.append(stdev,suffix="**stdev**")
            else:
                f_in.append(mean,suffix="**mean**")
                f_in.append(stdev,suffix="**stdev**")
        f_in.add_global_attributes({"namefile":''})

    if outputs_filename is not None:
        f_out = output_helper(outputs_filename,models[0],models[0].\
                          load_results(as_dict=True),**kwargs)
        vdict = {}
        vdicts = [output_helper(vdict,models[0],models[0].\
                                load_results(as_dict=True),**kwargs)]
        i = 1
        for m in models[1:]:
            suffix = m.name.split('.')[0].split('_')[-1]
            oudic = m.load_results(as_dict=True)
            vdict = {}
            output_helper(vdict,m,oudic,**kwargs)
            vdicts.append(vdict)
            if add_reals:
                f_out.append(vdict,suffix=suffix)
            i += 1

        mean, stdev = {},{}
        for vname in vdict.keys():
            alist = []
            for vd in vdicts:
                alist.append(vd[vname])
            alist = np.array(alist)
            mean[vname] = alist.mean(axis=0)
            stdev[vname] = alist.std(axis=0)
            mean[vname][np.isnan(vdict[vname])] = netcdf.FILLVALUE
            stdev[vname][np.isnan(vdict[vname])] = netcdf.FILLVALUE
            mean[vname][vdict[vname]==netcdf.FILLVALUE] = netcdf.FILLVALUE
            stdev[vname][vdict[vname]==netcdf.FILLVALUE] = netcdf.FILLVALUE
        if i >= 2:
            if not add_reals:
                f_out.write()
                f_out = NetCdf.empty_like(mean,output_filename=outputs_filename)
                f_out.append(mean,suffix="**mean**")
                f_out.append(stdev,suffix="**stdev**")

            else:
                f_out.append(mean,suffix="**mean**")
                f_out.append(stdev,suffix="**stdev**")
        f_out.add_global_attributes({"namefile":''})
    return f_in,f_out


def _add_output_nc_variable(f,times,shape3d,out_obj,var_name,logger=None,text='',
                            mask_vals=[],mask_array3d=None):
    if logger:
        logger.log("creating array for {0}".format(
                var_name))

    array = np.zeros((len(times),shape3d[0],shape3d[1],shape3d[2]),
                     dtype=np.float32)
    array[:] = np.NaN
    for i,t in enumerate(times):
        if t in out_obj.recordarray["totim"]:
            try:
                if text:
                    a = out_obj.get_data(totim=t,full3D=True,text=text)
                    if isinstance(a,list):
                        a = a[0]
                else:
                    a = out_obj.get_data(totim=t)
            except Exception as e:
                estr = "error getting data for {0} at time {1}:{2}".format(
                        var_name+text.decode().strip().lower(),t,str(e))
                if logger:
                    logger.warn(estr)
                else:
                    print(estr)
                continue
            if mask_array3d is not None and a.shape == mask_array3d.shape:
                a[mask_array3d] = np.NaN
            try:
                array[i,:,:,:] = a.astype(np.float32)
            except Exception as e:
                estr = "error assigning {0} data to array for time {1}:{2}".format(
                    var_name+text.decode().strip().lower(),t,str(e))
                if logger:
                    logger.warn(estr)
                else:
                    print(estr)
                continue

    if logger:
        logger.log("creating array for {0}".format(
                var_name))

    for mask_val in mask_vals:
        array[np.where(array==mask_val)] = np.NaN
    mx,mn = np.nanmax(array),np.nanmin(array)
    array[np.isnan(array)] = netcdf.FILLVALUE

    if isinstance(f,dict):
        if text:
            var_name = text.decode().strip().lower()
        f[var_name] = array
        return f

    units = None
    if var_name in NC_UNITS_FORMAT:
        units = NC_UNITS_FORMAT[var_name].format(
                f.grid_units, f.time_units)
    precision_str = "f4"

    if text:
        var_name = text.decode().strip().lower()
    attribs = {"long_name": var_name}
    attribs["coordinates"] = "time layer latitude longitude"
    attribs["min"] = mn
    attribs["max"] = mx
    if units is not None:
        attribs["units"] = units
    try:
        var = f.create_variable(var_name, attribs,
                                precision_str=precision_str,
                                dimensions=("time", "layer", "y", "x"))
    except Exception as e:
        estr = "error creating variable {0}:\n{1}".format(
                var_name, str(e))
        if logger:
            logger.lraise(estr)
        else:
            raise Exception(estr)

    try:
        var[:] = array
    except Exception as e:
        estr = "error setting array to variable {0}:\n{1}".format(
                var_name, str(e))
        if logger:
            logger.lraise(estr)
        else:
            raise Exception(estr)


def output_helper(f,ml,oudic,**kwargs):
    """export model outputs using the model spatial reference
    info.
    Parameters
    ----------
        f : filename for output - must have .shp or .nc extension
        ml : BaseModel derived type
        oudic : dict {output_filename,flopy datafile/cellbudgetfile instance}
    Returns
    -------
        None
    Note:
    ----
        casts down double precision to single precision for netCDF files

    """
    assert isinstance(ml,BaseModel)
    assert len(oudic.keys()) > 0
    logger = kwargs.pop("logger",None)
    stride = kwargs.pop("stride",1)
    suffix = kwargs.pop("suffix",None)
    forgive = kwargs.pop("forgive",False)
    if len(kwargs) > 0 and logger is not None:
        str_args = ','.join(kwargs)
        logger.warn("unused kwargs: "+str_args)
    # this sucks!  need to round the totims in each output file instance so
    # that they will line up
    for key,out in oudic.items():
       times = [float("{0:15.6f}".format(t)) for t in out.recordarray["totim"]]
       out.recordarray["totim"] = times

    times = []
    for filename,df in oudic.items():
        [times.append(t) for t in df.recordarray["totim"] if t not in times]
    assert len(times) > 0
    times.sort()

    # rectify times - only use times that are common to every output file
    common_times = []
    skipped_times = []
    for t in times:
        keep = True
        for filename,df in oudic.items():
            if t not in df.recordarray["totim"]:
                keep = False
                break
        if keep:
            common_times.append(t)
        else:
            skipped_times.append(t)
    assert len(common_times) > 0
    if len(skipped_times) > 0:
        if logger:
            logger.warn("the following output times are not common to all" +\
                        " output files and are being skipped:\n" +\
                        "{0}".format(skipped_times))
        else:
            print("the following output times are not common to all" +\
                        " output files and are being skipped:\n" +\
                        "{0}".format(skipped_times))
    times = [t for t in common_times[::stride]]
    if isinstance(f, str) and f.lower().endswith(".nc"):
        f = NetCdf(f, ml, time_values=times,logger=logger,
                   forgive=forgive)
    elif isinstance(f,NetCdf):
        otimes = list(f.nc.variables["time"][:])
        assert otimes == times
    if isinstance(f,NetCdf) or isinstance(f,dict):
        shape3d = (ml.nlay,ml.nrow,ml.ncol)
        mask_vals = []
        mask_array3d = None
        if ml.bas6:
            mask_vals.append(ml.bas6.hnoflo)
            mask_array3d = ml.bas6.ibound.array == 0
        if ml.bcf:
            mask_vals.append(ml.bcf.hdry)
        if ml.lpf:
            mask_vals.append(ml.lpf.hdry)

        for filename,out_obj in oudic.items():
            filename = filename.lower()

            if isinstance(out_obj,UcnFile):
                _add_output_nc_variable(f,times,shape3d,out_obj,
                                        "concentration",logger=logger,
                                        mask_vals=mask_vals,
                                        mask_array3d=mask_array3d)

            elif isinstance(out_obj,HeadFile):
                _add_output_nc_variable(f,times,shape3d,out_obj,
                                        out_obj.text.decode(),logger=logger,
                                        mask_vals=mask_vals,
                                        mask_array3d=mask_array3d)

            elif isinstance(out_obj,FormattedHeadFile):
                _add_output_nc_variable(f,times,shape3d,out_obj,
                                        out_obj.text,logger=logger,
                                        mask_vals=mask_vals,
                                        mask_array3d=mask_array3d)

            elif isinstance(out_obj,CellBudgetFile):
                var_name = "cell_by_cell_flow"
                for text in out_obj.textlist:
                    _add_output_nc_variable(f,times,shape3d,out_obj,
                                            var_name,logger=logger,text=text,
                                            mask_vals=mask_vals,
                                            mask_array3d=mask_array3d)

            else:
                estr = "unrecognized file extention:{0}".format(filename)
                if logger:
                    logger.lraise(estr)
                else:
                    raise Exception(estr)

    else:
        if logger:
            logger.lraise("unrecognized export argument:{0}".format(f))
        else:
            raise NotImplementedError("unrecognized export argument" +\
                                      ":{0}".format(f))
    return f


def model_helper(f, ml, **kwargs):
    assert isinstance(ml,BaseModel)
    package_names = kwargs.get("package_names", None)
    if package_names is None:
        package_names = [pak.name[0] for pak in ml.packagelist]

    if isinstance(f, str) and f.lower().endswith(".nc"):
        f = NetCdf(f, ml, **kwargs)

    if isinstance(f, str) and f.lower().endswith(".shp"):
        shapefile_utils.model_attributes_to_shapefile(f, ml,
                                      package_names=package_names,
                                      **kwargs)

    elif isinstance(f,NetCdf):

        for pak in ml.packagelist:
            if pak.name[0] in package_names:
                f = pak.export(f,**kwargs)
        return f

    elif isinstance(f,dict):
        for pak in ml.packagelist:
            f = pak.export(f,**kwargs)

    else:
        raise NotImplementedError("unrecognized export argument:{0}".format(f))

    return f

def package_helper(f, pak, **kwargs):
    assert isinstance(pak,Package)
    if isinstance(f, str) and f.lower().endswith(".nc"):
        f = NetCdf(f, pak.parent)

    if isinstance(f, str) and f.lower().endswith(".shp"):
        shapefile_utils.model_attributes_to_shapefile(f, pak.parent,
                                                      package_names=pak.name,
                                                      **kwargs)

    elif isinstance(f, NetCdf) or isinstance(f,dict):
        attrs = dir(pak)
        if 'sr' in attrs:
            attrs.remove('sr')
        if 'start_datetime' in attrs:
            attrs.remove('start_datetime')
        for attr in attrs:
            if '__' in attr:
                continue
            a = pak.__getattribute__(attr)
            if isinstance(a, Util2d) and len(a.shape) == 2 and a.shape[1] > 0:
                try:
                    f = util2d_helper(f, a, **kwargs)
                except:
                    f.logger.warn("error adding {0} as variable".format(a.name))
            elif isinstance(a, Util3d):
                f = util3d_helper(f, a, **kwargs)
            elif isinstance(a, Transient2d):
                f = transient2d_helper(f, a, **kwargs)
            elif isinstance(a, MfList):
                f = mflist_helper(f, a, **kwargs)
            elif isinstance(a, list):
                for v in a:
                    if isinstance(v, Util3d):
                        f = util3d_helper(f, v, **kwargs)
        return f

    else:
        raise NotImplementedError("unrecognized export argument:{0}".format(f))


def generic_array_helper(f, array, var_name="generic_array",
                         dimensions = ("time", "layer", "y", "x"),
                         precision_str="f4",units="unitless",**kwargs):
    assert isinstance(f,NetCdf),"generic_array_helper() can only be used " +\
                                "with instantiated netCDfs"
    assert array.ndim == len(dimensions),"generic_array_helper() "+\
                                         "array.ndim != dimensions"
    coords_dims = {"time":"time", "layer":"layer", "y":"latitude","x":"longitude"}
    coords = ' '.join([coords_dims[d] for d in dimensions])
    mn = kwargs.pop("min",-1.0e+9)
    mx = kwargs.pop("max",1.0e+9)
    long_name = kwargs.pop("long_name",var_name)
    if len(kwargs) > 0:
        f.logger.warn("generic_array_helper(): unrecognized kwargs:" +\
                        ",".join(kwargs.keys()))
    attribs = {"long_name": long_name}
    attribs["coordinates"] = coords
    attribs["units"] = units
    attribs["min"] = mn
    attribs["max"] = mx
    try:
        var = f.create_variable(var_name, attribs, precision_str=precision_str,
                                dimensions=dimensions)
    except Exception as e:
        estr = "error creating variable {0}:\n{1}".format(var_name, str(e))
        f.logger.warn(estr)
        raise Exception(estr)
    try:
        var[:] = array
    except Exception as e:
        estr = "error setting array to variable {0}:\n{1}".format(var_name, str(e))
        f.logger.warn(estr)
        raise Exception(estr)
    return f


def mflist_helper(f, mfl, **kwargs):
    """ export helper for MfList instances

    Parameters
    -----------
        f : string (filename) or existing export instance type (NetCdf only for now)
        mfl : MfList instance

    """
    assert isinstance(mfl, MfList) \
        , "mflist_helper only helps MfList instances"

    if isinstance(f, str) and f.lower().endswith(".nc"):
        f = NetCdf(f, mfl.model)

    if isinstance(f, str) and f.lower().endswith(".shp"):
        kper = kwargs.get("kper",None)
        if mfl.sr is None:
            raise Exception("MfList.to_shapefile: SpatialReference not set")
        import flopy.utils.flopy_io as fio
        if kper is None:
            keys = mfl.data.keys()
            keys.sort()
        else:
            keys = [kper]
        array_dict = {}
        for kk in keys:
            arrays = mfl.to_array(kk)
            for name, array in arrays.items():
                for k in range(array.shape[0]):
                    #aname = name+"{0:03d}_{1:02d}".format(kk, k)
                    n = fio.shape_attr_name(name, length=4)
                    aname = "{}{:03d}{:03d}".format(n, k+1, int(kk)+1)
                    array_dict[aname] = array[k]
        shapefile_utils.write_grid_shapefile(f, mfl.sr, array_dict)


    elif isinstance(f, NetCdf) or isinstance(f,dict):
        base_name = mfl.package.name[0].lower()
        #f.log("getting 4D masked arrays for {0}".format(base_name))
        #m4d = mfl.masked_4D_arrays
        #f.log("getting 4D masked arrays for {0}".format(base_name))

        #for name, array in m4d.items():
        for name, array in mfl.masked_4D_arrays_itr():
            var_name = base_name + '_' + name
            if isinstance(f,dict):
                f[var_name] = array
                continue
            f.log("processing {0} attribute".format(name))

            units = None
            if var_name in NC_UNITS_FORMAT:
                units = NC_UNITS_FORMAT[var_name].format(f.grid_units, f.time_units)
            precision_str = NC_PRECISION_TYPE[mfl.dtype[name].type]
            if var_name in NC_LONG_NAMES:
                attribs = {"long_name":NC_LONG_NAMES[var_name]}
            else:
                attribs = {"long_name":var_name}
            attribs["coordinates"] = "time layer latitude longitude"
            attribs["min"] = np.nanmin(array)
            attribs["max"] = np.nanmax(array)
            if units is not None:
                attribs["units"] = units
            try:
                var = f.create_variable(var_name, attribs, precision_str=precision_str,
                                        dimensions=("time", "layer", "y", "x"))
            except Exception as e:
                estr = "error creating variable {0}:\n{1}".format(var_name, str(e))
                f.logger.warn(estr)
                raise Exception(estr)

            array[np.isnan(array)] = f.fillvalue
            try:
                var[:] = array
            except Exception as e:
                estr = "error setting array to variable {0}:\n{1}".format(var_name, str(e))
                f.logger.warn(estr)
                raise Exception(estr)
            f.log("processing {0} attribute".format(name))

        return f
    else:
        raise NotImplementedError("unrecognized export argument:{0}".format(f))

def transient2d_helper(f, t2d, **kwargs):
    """ export helper for Transient2d instances

    Parameters
    -----------
        f : string (filename) or existing export instance type (NetCdf only for now)
        t2d : Transient2d instance
        min_valid : minimum valid value
        max_valid : maximum valid value

    """

    assert isinstance(t2d, Transient2d)\
        , "transient2d_helper only helps Transient2d instances"

    min_valid = kwargs.get("min_valid", -1.0e+9)
    max_valid = kwargs.get("max_valid", 1.0e+9)

    if isinstance(f, str) and f.lower().endswith(".nc"):
        f = NetCdf(f, t2d.model)

    if isinstance(f, str) and f.lower().endswith(".shp"):
        array_dict = {}
        for kper in range(t2d.model.nper):
            u2d = t2d[kper]
            name = '{}_{:03d}'.format(shapefile_utils.shape_attr_name(u2d.name), kper + 1)
            array_dict[name] = u2d.array
        shapefile_utils.write_grid_shapefile(f, t2d.model.sr, array_dict)

    elif isinstance(f, NetCdf) or isinstance(f,dict):
        # mask the array is defined by any row col with at lease
        # one active cell
        mask = None
        if t2d.model.bas6 is not None:
            ibnd = np.abs(t2d.model.bas6.ibound.array).sum(axis=0)
            mask = ibnd == 0
        elif t2d.model.btn is not None:
            ibnd = np.abs(t2d.model.btn.icbund.array).sum(axis=0)
            mask = ibnd == 0

        #f.log("getting 4D array for {0}".format(t2d.name_base))
        array = t2d.array
        #f.log("getting 4D array for {0}".format(t2d.name_base))
        with np.errstate(invalid="ignore"):
            if array.dtype not in [int,np.int,np.int32,np.int64]:
                if mask is not None:
                    array[:,0,mask] = np.NaN
                array[array <= min_valid] = np.NaN
                array[array >= max_valid] = np.NaN
                mx, mn = np.nanmax(array), np.nanmin(array)
            else:
                mx, mn = np.nanmax(array), np.nanmin(array)
                array[array <= min_valid] = netcdf.FILLVALUE
                array[array >= max_valid] = netcdf.FILLVALUE
                #if t2d.model.bas6 is not None:
                #    array[:, 0, t2d.model.bas6.ibound.array[0] == 0] = \
                #        f.fillvalue
                #elif t2d.model.btn is not None:
                #    array[:, 0, t2d.model.btn.icbund.array[0] == 0] = \
                #        f.fillvalue

        var_name = t2d.name_base.replace('_', '')
        if isinstance(f,dict):
            array[array==netcdf.FILLVALUE] = np.NaN
            f[var_name] = array
            return f

        array[np.isnan(array)] = f.fillvalue
        units = "unitless"

        if var_name in NC_UNITS_FORMAT:
            units = NC_UNITS_FORMAT[var_name].format(f.grid_units, f.time_units)
        try:
            precision_str = NC_PRECISION_TYPE[t2d.dtype]
        except:
            precision_str = NC_PRECISION_TYPE[t2d.dtype.type]
        if var_name in NC_LONG_NAMES:
            attribs = {"long_name": NC_LONG_NAMES[var_name]}
        else:
            attribs = {"long_name": var_name}
        attribs["coordinates"] = "time layer latitude longitude"
        attribs["units"] = units
        attribs["min"] = mn
        attribs["max"] = mx
        try:
            var = f.create_variable(var_name, attribs, precision_str=precision_str,
                                    dimensions=("time", "layer", "y", "x"))
        except Exception as e:
            estr = "error creating variable {0}:\n{1}".format(var_name, str(e))
            f.logger.warn(estr)
            raise Exception(estr)
        try:
            var[:,0] = array
        except Exception as e:
            estr = "error setting array to variable {0}:\n{1}".format(var_name, str(e))
            f.logger.warn(estr)
            raise Exception(estr)
        return f

    else:
        raise NotImplementedError("unrecognized export argument:{0}".format(f))


def util3d_helper(f, u3d, **kwargs):
    """ export helper for Transient2d instances

    Parameters
    -----------
        f : string (filename) or existing export instance type (NetCdf only for now)
        u3d : Util3d instance
        min_valid : minimum valid value
        max_valid : maximum valid value

    """

    assert isinstance(u3d, Util3d), "util3d_helper only helps Util3d instances"
    assert len(u3d.shape) == 3, "util3d_helper only supports 3D arrays"

    min_valid = kwargs.get("min_valid", -1.0e+9)
    max_valid = kwargs.get("max_valid", 1.0e+9)

    if isinstance(f, str) and f.lower().endswith(".nc"):
        f = NetCdf(f, u3d.model)

    if isinstance(f, str) and f.lower().endswith(".shp"):
        array_dict = {}
        for ilay in range(u3d.model.nlay):
            u2d = u3d[ilay]
            name = '{}_{:03d}'.format(shapefile_utils.shape_attr_name(u2d.name), ilay + 1)
            array_dict[name] = u2d.array
        shapefile_utils.write_grid_shapefile(f, u3d.model.sr,
                             array_dict)

    elif isinstance(f, NetCdf) or isinstance(f,dict):
        var_name = u3d.name[0].replace(' ', '_').lower()
        #f.log("getting 3D array for {0}".format(var_name))
        array = u3d.array

        # this is for the crappy vcont in bcf6
        # if isinstance(f,NetCdf) and array.shape != f.shape:
        #     f.log("broadcasting 3D array for {0}".format(var_name))
        #     full_array = np.empty(f.shape)
        #     full_array[:] = np.NaN
        #     full_array[:array.shape[0]] = array
        #     array = full_array
        #     f.log("broadcasting 3D array for {0}".format(var_name))
        #f.log("getting 3D array for {0}".format(var_name))
            #
        mask = None
        if u3d.model.bas6 is not None and "ibound" not in var_name:
            mask = u3d.model.bas6.ibound.array == 0
        elif u3d.model.btn is not None and 'icbund' not in var_name:
            mask = u3d.model.btn.icbund.array == 0

        if mask is not None and array.shape != mask.shape:
            #f.log("broadcasting 3D array for {0}".format(var_name))
            full_array = np.empty(mask.shape)
            full_array[:] = np.NaN
            full_array[:array.shape[0]] = array
            array = full_array
            #f.log("broadcasting 3D array for {0}".format(var_name))


        # runtime warning issued in some cases - need to track down cause
        # happens when NaN is already in array
        with np.errstate(invalid="ignore"):
            if array.dtype not in [int,np.int,np.int32,np.int64]:
                #if u3d.model.bas6 is not None and "ibound" not in var_name:
                #    array[u3d.model.bas6.ibound.array == 0] = np.NaN
                #elif u3d.model.btn is not None and 'icbund' not in var_name:
                #    array[u3d.model.btn.icbund.array == 0] = np.NaN
                array[mask] = np.NaN
                array[array <= min_valid] = np.NaN
                array[array >= max_valid] = np.NaN
                mx, mn = np.nanmax(array), np.nanmin(array)
            else:
                mx, mn = np.nanmax(array), np.nanmin(array)
                array[mask] = netcdf.FILLVALUE
                array[array <= min_valid] = netcdf.FILLVALUE
                array[array >= max_valid] = netcdf.FILLVALUE
                if u3d.model.bas6 is not None and "ibound" not in var_name:
                    array[u3d.model.bas6.ibound.array == 0] = netcdf.FILLVALUE
                elif u3d.model.btn is not None and 'icbund' not in var_name:
                    array[u3d.model.btn.icbund.array == 0] = netcdf.FILLVALUE

        if isinstance(f,dict):
            f[var_name] = array
            return f

        array[np.isnan(array)] = f.fillvalue
        units = "unitless"
        if var_name in NC_UNITS_FORMAT:
            units = NC_UNITS_FORMAT[var_name].format(f.grid_units, f.time_units)
        precision_str = NC_PRECISION_TYPE[u3d.dtype]
        if var_name in NC_LONG_NAMES:
            attribs = {"long_name": NC_LONG_NAMES[var_name]}
        else:
            attribs = {"long_name": var_name}
        attribs["coordinates"] = "layer latitude longitude"
        attribs["units"] = units
        attribs["min"] = mn
        attribs["max"] = mx
        try:
            var = f.create_variable(var_name, attribs, precision_str=precision_str,
                                    dimensions=("layer", "y", "x"))
        except Exception as e:
            estr = "error creating variable {0}:\n{1}".format(var_name, str(e))
            f.logger.warn(estr)
            raise Exception(estr)
        try:
            var[:] = array
        except Exception as e:
            estr = "error setting array to variable {0}:\n{1}".format(var_name, str(e))
            f.logger.warn(estr)
            raise Exception(estr)
        return f

    else:
        raise NotImplementedError("unrecognized export argument:{0}".format(f))


def util2d_helper(f, u2d, **kwargs):
    """ export helper for Util2d instances

    Parameters
    ----------
        f : string (filename) or existing export instance type (NetCdf only for now)
        u2d : Util2d instance
        min_valid : minimum valid value
        max_valid : maximum valid value

    """
    assert isinstance(u2d, Util2d), "util2d_helper only helps Util2d instances"
    assert len(u2d.shape) == 2, "util2d_helper only supports 2D arrays"

    min_valid = kwargs.get("min_valid", -1.0e+9)
    max_valid = kwargs.get("max_valid", 1.0e+9)

    if isinstance(f, str) and f.lower().endswith(".nc"):
        f = NetCdf(f, u2d.model)

    if isinstance(f, str) and f.lower().endswith(".shp"):
        name = shapefile_utils.shape_attr_name(u2d.name, keep_layer=True)
        shapefile_utils.write_grid_shapefile(f, u2d.model.sr, {name: u2d.array})
        return

    elif isinstance(f, NetCdf) or isinstance(f,dict):

        # try to mask the array - assume layer 1 ibound is a good mask
        #f.log("getting 2D array for {0}".format(u2d.name))
        array = u2d.array
        #f.log("getting 2D array for {0}".format(u2d.name))

        with np.errstate(invalid="ignore"):
            if array.dtype not in [int,np.int,np.int32,np.int64]:
                if u2d.model.bas6 is not None and \
                                "ibound" not in u2d.name.lower():
                    array[u2d.model.bas6.ibound.array[0, :, :] == 0] = np.NaN
                elif u2d.model.btn is not None and \
                                "icbund" not in u2d.name.lower():
                    array[u2d.model.btn.icbund.array[0, :, :] == 0] = np.NaN
                array[array <= min_valid] = np.NaN
                array[array >= max_valid] = np.NaN
                mx, mn = np.nanmax(array), np.nanmin(array)
            else:
                mx, mn = np.nanmax(array), np.nanmin(array)
                array[array <= min_valid] = netcdf.FILLVALUE
                array[array >= max_valid] = netcdf.FILLVALUE
                if u2d.model.bas6 is not None and \
                                "ibound" not in u2d.name.lower():
                    array[u2d.model.bas6.ibound.array[0, :, :] == 0] = \
                        netcdf.FILLVALUE
                elif u2d.model.btn is not None and \
                                "icbund" not in u2d.name.lower():
                    array[u2d.model.btn.icbund.array[0, :, :] == 0] = \
                        netcdf.FILLVALUE
        var_name = u2d.name
        if isinstance(f,dict):
            f[var_name] = array
            return f

        array[np.isnan(array)] = f.fillvalue
        units = "unitless"

        if var_name in NC_UNITS_FORMAT:
            units = NC_UNITS_FORMAT[var_name].format(f.grid_units, f.time_units)
        precision_str = NC_PRECISION_TYPE[u2d.dtype]
        if var_name in NC_LONG_NAMES:
            attribs = {"long_name":NC_LONG_NAMES[var_name]}
        else:
            attribs = {"long_name": var_name}
        attribs["coordinates"] = "latitude longitude"
        attribs["units"] = units
        attribs["min"] = mn
        attribs["max"] = mx

        try:
            var = f.create_variable(var_name, attribs, precision_str=precision_str,
                                    dimensions=("y", "x"))
        except Exception as e:
            estr = "error creating variable {0}:\n{1}".format(var_name, str(e))
            f.logger.warn(estr)
            raise Exception(estr)
        try:
            var[:] = array
        except Exception as e:
            estr = "error setting array to variable {0}:\n{1}".format(var_name, str(e))
            f.logger.warn(estr)
            raise Exception(estr)
        return f

    else:
        raise NotImplementedError("unrecognized export argument:{0}".format(f))
