import cv2 as cv
import numpy as np

class VisUtil:

    @staticmethod
    def get_red_mask(hsv_img):
        # note that is in [0-180] and s, v in [0, 255]
        mask1 = cv.inRange(hsv_img, (0, 100, 70), (2, 255, 255))
        mask2 = cv.inRange(hsv_img, (170, 100, 70), (180, 255, 255))
        return mask1 + mask2  # and

    @staticmethod
    def noise_down(mask):
        kernel1 = np.ones((4, 4), np.uint8)
        # opposite to erode. Fill red space
        res = cv.dilate(mask, kernel1, iterations=2)
        # if one pix is black center is black
        res = cv.erode(res, kernel1, iterations=1)
        return res

    @staticmethod
    def white_balance(img):
        result = cv.cvtColor(img, cv.COLOR_BGR2LAB)
        avg_a = np.average(result[:, :, 1])
        avg_b = np.average(result[:, :, 2])
        result[:, :, 1] = result[:, :, 1] - ((avg_a - 128) * (result[:, :, 0] / 255.0) * 1.1)
        result[:, :, 2] = result[:, :, 2] - ((avg_b - 128) * (result[:, :, 0] / 255.0) * 1.1)
        result = cv.cvtColor(result, cv.COLOR_LAB2BGR)
        return result

    # frame may be passed only for painting contours
    @staticmethod
    def find_interesting_bound(mask):
        contours, hierarchy = cv.findContours(mask, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
        if len(contours) != 0:
            maxContour = max(contours, key=cv.contourArea)
            # x, y, w, h = cv.boundingRect(maxContour)
            return cv.boundingRect(maxContour), maxContour, contours
        else:
            return (0, 0, 0, 0), None, None

    @staticmethod
    def crop(img, x, y, w, h):
        if w > 50 and h > 50:
            topx = x + w
            topy = y + h
            cropped = img[y:topy, x:topx]
            cropped = cv.resize(cropped, (128, 128))

            return cropped
        else:
            return None


from threading import Thread
from typing import List
import tensorflow as tf
from queue import Queue
import numpy as np
from nboost.model.bert_model import modeling, tokenization
from nboost.model.base import BaseModel
from nboost.types import Choice

class BertModel(BaseModel):

    def __init__(self, verbose=False, **kwargs):
        super().__init__(**kwargs)
        self.download()

        self.output_q = Queue()
        self.input_q = Queue()

        ckpts = list(self.model_dir.glob('*.ckpt*'))
        if not len(ckpts) > 0:
            raise FileNotFoundError("Tensorflow model not found")
        self.checkpoint = str(ckpts[0]).split('.ckpt')[0] + '.ckpt'
        self.vocab_file = str(self.model_dir.joinpath('vocab.txt'))
        self.bert_config_file = str(self.model_dir.joinpath('bert_config.json'))
        if not verbose:
            tf.logging.set_verbosity(tf.logging.ERROR)
        self.model_thread = Thread(target=self.run_model)
        self.model_thread.start()

    @staticmethod
    def create_model(bert_config, input_ids, input_mask, segment_ids,
                     labels, num_labels):
        """Creates a classification model."""
        model = modeling.BertModel(
            config=bert_config,
            is_training=False,
            input_ids=input_ids,
            input_mask=input_mask,
            token_type_ids=segment_ids,
            use_one_hot_embeddings=False)

        output_layer = model.get_pooled_output()
        hidden_size = output_layer.shape[-1].value

        output_weights = tf.get_variable(
            "output_weights", [num_labels, hidden_size],
            initializer=tf.truncated_normal_initializer(stddev=0.02))

        output_bias = tf.get_variable(
            "output_bias", [num_labels], initializer=tf.zeros_initializer())

        with tf.variable_scope("loss"):
            logits = tf.matmul(output_layer, output_weights, transpose_b=True)
            logits = tf.nn.bias_add(logits, output_bias)
            log_probs = tf.nn.log_softmax(logits, axis=-1)

            one_hot_labels = tf.one_hot(labels, depth=num_labels, dtype=tf.float32)

            per_example_loss = -tf.reduce_sum(one_hot_labels * log_probs, axis=-1)
            loss = tf.reduce_mean(per_example_loss)

            return (loss, per_example_loss, log_probs)

    def model_fn_builder(self, bert_config, num_labels, init_checkpoint):
        """Returns `model_fn` closure for TPUEstimator."""

        def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument
            """The `model_fn` for TPUEstimator."""

            input_ids = features["input_ids"]
            input_mask = features["input_mask"]
            segment_ids = features["segment_ids"]
            label_ids = features["label_ids"]

            (total_loss, per_example_loss, log_probs) = self.create_model(
                bert_config, input_ids, input_mask, segment_ids, label_ids,
                num_labels)

            tvars = tf.trainable_variables()

            (assignment_map, initialized_variable_names
             ) = modeling.get_assignment_map_from_checkpoint(tvars, init_checkpoint)

            tf.train.init_from_checkpoint(init_checkpoint, assignment_map)

            output_spec = tf.estimator.EstimatorSpec(
                mode=mode,
                predictions={
                    "log_probs": log_probs,
                    "label_ids": label_ids,
                })
            return output_spec

        return model_fn

    def input_fn(self):
        """The actual input function."""

        output_types = {
            "input_ids": tf.int32,
            "segment_ids": tf.int32,
            "input_mask": tf.int32,
            "label_ids": tf.int32,
        }
        dataset = tf.data.Dataset.from_generator(self.feature_generator, output_types)

        dataset = dataset.padded_batch(
            batch_size=self.batch_size,
            padded_shapes={
                "input_ids": [self.max_seq_len],
                "segment_ids": [self.max_seq_len],
                "input_mask": [self.max_seq_len],
                "label_ids": [],
            },
            padding_values={
                "input_ids": 0,
                "segment_ids": 0,
                "input_mask": 0,
                "label_ids": 0,
            },
            drop_remainder=True)

        return dataset

    def run_model(self):
        bert_config = modeling.BertConfig.from_json_file(self.bert_config_file)
        assert self.max_seq_len <= bert_config.max_position_embeddings

        run_config = tf.estimator.RunConfig(model_dir=str(self.data_dir))

        model_fn = self.model_fn_builder(
            bert_config=bert_config,
            num_labels=2,
            init_checkpoint=self.checkpoint)

        estimator = tf.estimator.Estimator(
            model_fn=model_fn,
            config=run_config)

        result = estimator.predict(input_fn=self.input_fn,
                                   yield_single_examples=True)

        for item in result:
            self.output_q.put((item["log_probs"], item["label_ids"]))

    def feature_generator(self):
        tokenizer = tokenization.FullTokenizer(vocab_file=self.vocab_file, do_lower_case=True)
        while True:
            next = self.input_q.get()
            if not next:
                break
            query, candidates = next
            query = tokenization.convert_to_unicode(query)
            query_token_ids = tokenization.convert_to_bert_input(
                text=query, max_seq_length=self.max_seq_len, tokenizer=tokenizer,
                add_cls=True)

            for i, doc_text in enumerate(candidates):
                doc_token_id = tokenization.convert_to_bert_input(
                    text=tokenization.convert_to_unicode(doc_text),
                    max_seq_length=self.max_seq_len - len(query_token_ids),
                    tokenizer=tokenizer,
                    add_cls=False)

                query_ids = query_token_ids
                doc_ids = doc_token_id
                input_ids = query_ids + doc_ids

                query_segment_id = [0] * len(query_ids)
                doc_segment_id = [1] * len(doc_ids)
                segment_ids = query_segment_id + doc_segment_id

                input_mask = [1] * len(input_ids)

                features = {
                    "input_ids": input_ids,
                    "segment_ids": segment_ids,
                    "input_mask": input_mask,
                    "label_ids": 0
                }
                yield features

    def pad(self, candidates):
        if len(candidates) % self.batch_size == 0:
            return candidates
        else:
            candidates += ['PADDING DOC'] * (self.batch_size - (len(candidates) % self.batch_size))
            return candidates

    def rank(self, query: bytes, choices: List[Choice]) -> List[int]:
        bodies = [choice.body for choice in choices]
        actual_length = len(bodies)
        candidates = self.pad(bodies)
        self.input_q.put((query, bodies))

        results = [self.output_q.get() for _ in range(len(candidates))][:actual_length]
        log_probs, labels = zip(*results)
        log_probs = np.stack(log_probs).reshape(-1, 2)
        scores = log_probs[:, 1]
        assert len(scores) == actual_length
        return scores.argsort()[::-1]

    def close(self):
        self.input_q.put(None)
        self.model_thread.join()


#!/usr/bin/env python

import os
import sys
import time
import pdb
import argparse
import numpy
import crowdsource.psf as psfmod
from astropy.io import fits
import crowdsource
from unwise_psf import unwise_psf
# implicit dependency for WISE runs only
# https://github.com/legacysurvey/unwise_psf
import crowdsource.unwise_primary as unwise_primary
from astropy import wcs
from collections import OrderedDict
from pkg_resources import resource_filename

extrabits = {'crowdsat': 2**25,
             'nebulosity': 2**26,
             'w1brightoffedge': 2**7,
             'w2brightoffedge': 2**8,
             'hyperleda': 2**9}

nodeblend_bits = extrabits['hyperleda']
sharp_bits = (extrabits['w1brightoffedge'] | extrabits['w2brightoffedge'])

def wise_filename(basedir, coadd_id, band, _type, uncompressed=False,
                  drop_first_dir=False, epoch=-1):
    # type should be one of:
    # 'img-u', 'img-m', 'invvar-u', 'invvar-m', 'std-u', 'std-m'
    # 'n-u', 'n-m', 'frames', 'msk'

    # -msk is special because the info for both W1/W2 is in same file

    fname = 'unwise-' + coadd_id
    if _type != 'msk':
        fname += '-w' + str(band)

    fname += ('-' + _type + '.fits')

    path = [basedir, coadd_id[0:3], coadd_id, fname]
    if drop_first_dir:
        del path[1]
    if epoch >= 0:
        epochstr = 'e%03d' % epoch if _type != 'msk' else 'fulldepth'
        path = path[0:1] + [epochstr] + path[1:]
    fname = os.path.join(*path)

    if not uncompressed or _type == 'msk':
        if (_type != 'img-u') and (_type != 'img-m') and (_type != 'frames'):
            fname += '.gz'

    return fname

def read_blist(brightstars, raim, decim, hdr, maxsep):
    from astropy.coordinates.angle_utilities import angular_separation
    sep = angular_separation(numpy.radians(brightstars['ra']),
                             numpy.radians(brightstars['dec']),
                             numpy.radians(raim),
                             numpy.radians(decim))
    sep = numpy.degrees(sep)
    m = (sep < 3) & (brightstars['k_m'] < 5)
    brightstars = brightstars[m]
    wcs0 = wcs.WCS(hdr)
    yy, xx = wcs0.all_world2pix(brightstars['ra'], brightstars['dec'], 0)
    m = (xx > 0) & (xx < hdr['NAXIS1']) & (yy > 0) & (yy < hdr['NAXIS2'])
    xx, yy = xx[m], yy[m]
    mag = brightstars['k_m'][m]
    if not numpy.any(m):
        return None
    else:
        return [xx, yy, mag]

def massage_isig_and_dim(isig, im, flag, band, nm, nu, fac=None):
    """Construct a WISE inverse sigma image and add saturation to flag.

    unWISE provides nice inverse variance maps.  These however have no
    contribution from Poisson noise from sources, and so underestimate
    the uncertainties dramatically in bright regions.  This can pull the
    whole fit awry in bright areas, since the sky model means that every
    pixel feels every other pixel.

    It's not clear what the best solution is.  We make a goofy inverse
    sigma image from the original image and the inverse variance image.  It
    is intended to be sqrt(ivar) for the low count regime and grow like
    sqrt(1/im) for the high count regime.  The constant of proportionality
    should in principle be worked out; here I set it to 0.15, which worked
    once, and it doesn't seem like this should depend much on which
    WISE exposure the image came from?  It's ultimately something like the gain
    or zero point...
    """

    if fac is None:
        bandfacs = {1: 0.15, 2: 0.3}
        bandfloors = {1: 0.5, 2: 2}
        fac = bandfacs[band]
        floor = bandfloors[band]

    satbit = 16 if band == 1 else 32
    satlimit = 85000  # if band == 1 else 130000
    msat = ((flag & satbit) != 0) | (im > satlimit) | ((nm == 0) & (nu > 1))
    from scipy.ndimage import morphology
    # dilate = morphology.iterate_structure(
    #     morphology.generate_binary_structure(2, 1), 3)
    xx, yy = numpy.mgrid[-3:3+1, -3:3+1]
    dilate = xx**2+yy**2 <= 3**2
    msat = morphology.binary_dilation(msat, dilate)
    isig[msat] = 0
    flag = flag.astype('i8')
    # zero out these bits; we claim them for our own purposes.
    massagebits = (extrabits['crowdsat'] | crowdsource.nodeblend_maskbit |
                   crowdsource.sharp_maskbit | extrabits['nebulosity'])
    flag &= ~massagebits
    flag[msat] |= extrabits['crowdsat']
    flag[(flag & nodeblend_bits) != 0] |= crowdsource.nodeblend_maskbit
    flag[(flag & sharp_bits) != 0] |= crowdsource.sharp_maskbit

    sigma = numpy.sqrt(1./(isig + (isig == 0))**2 + floor**2 +
                       fac**2*numpy.clip(im, 0, numpy.inf))
    sigma[msat] = numpy.inf
    sigma[isig == 0] = numpy.inf
    return (1./sigma).astype('f4'), flag

def wise_psf_stamp(band, nosmooth=False):
    # psf noise: ~roughly 0.1 count in outskirts of W1 and W2
    if band >= 3:
        raise ValueError('Need to stare at W3+ PSF more!')
    psfnoise = 0.1
    stampfn = resource_filename('unwise_psf',
                                'data/psf_model_w'+str(band)+'.fits')
    stamp = fits.getdata(stampfn)
    edges = numpy.concatenate([stamp[0, 1:-1], stamp[-1, 1:-1],
                               stamp[1:-1, 0], stamp[1:-1, -1]])
    medval = numpy.median(edges[edges != 0]) / 2
    stamp[stamp == 0] = medval
    stamp -= medval
    from scipy import signal
    stamp[stamp < 0] = 0.
    # suppress spurious warnings in signal.wiener
    olderr = numpy.seterr(invalid='ignore', divide='ignore')
    # update to scipy.signal means that Wiener filter uses an FFT
    # to perform the various convolutions, which causes bad errors
    # here unless we cast to f8.  It's not that hard to do something
    # a bit better than scipy.signal.wiener---morally we really want to do
    # something like smooth in log space on radial lines---but I don't
    # want to go further down that rabbit hole today.
    stamp = signal.wiener(stamp.astype('f8'),  11, psfnoise)
    stamp = stamp.astype('f4')
    numpy.seterr(**olderr)
    # taper linearly over outer 60 pixels?
    stampszo2 = stamp.shape[0] // 2
    xx, yy = numpy.mgrid[-stampszo2:stampszo2+1, -stampszo2:stampszo2+1]
    edgedist = numpy.clip(stampszo2-numpy.abs(xx), 0,
                          stampszo2-numpy.abs(yy))
    stamp = stamp * numpy.clip(edgedist / 60., stamp < 10, 1)
    import psf
    stamp = psf.center_psf(stamp, censize=19)
    stamp = stamp / numpy.sum(stamp)
    return stamp

def wise_psf(band, coadd_id):
    stamp = wise_psf_stamp(band)
    stamp = unwise_psf.rotate_using_rd(stamp, coadd_id)
    psf = psfmod.SimplePSF(stamp)
    from functools import partial
    psf.fitfun = partial(psfmod.wise_psf_fit, psfstamp=stamp)
    return psf

def wise_psf_grid(band, coadd_id, basedir, uncompressed=False,
                  drop_first_dir=False, ngrid=None, epoch=-1):
    imagefn = wise_filename(basedir, coadd_id, band, 'img-m',
                            uncompressed=uncompressed,
                            drop_first_dir=drop_first_dir, epoch=epoch)
    hdr = fits.getheader(imagefn)
    if ngrid is None:
        rr, dd = hdr['CRVAL1'], hdr['CRVAL2']
        from astropy.coordinates import SkyCoord
        from astropy import units as u
        coord = SkyCoord(ra=rr*u.deg, dec=dd*u.deg, frame='icrs')
        coord = coord.geocentrictrueecliptic
        lam, bet = coord.lon.deg, coord.lat.deg
        dlam = 1.4/(numpy.abs(numpy.cos(numpy.radians(bet)))+1e-6)
        ngrid = numpy.floor(numpy.clip(dlam / 1, 4, 16)).astype('i4')
    x = numpy.linspace(0, 2047, ngrid)
    y = numpy.linspace(0, 2047, ngrid)
    wcs0 = wcs.WCS(hdr)
    stamp = wise_psf_stamp(band).astype('f4')
    stamps = numpy.zeros((len(x), len(y))+stamp.shape, dtype=stamp.dtype)
    unwise_psf.rotate_using_convolution.cache = None  # clear cache
    for i in range(len(x)):
        for j in range(len(y)):
            rr, dd = wcs0.all_pix2world(y[j], x[i], 0)
            stamps[i, j, ...] = unwise_psf.rotate_using_rd(
                stamp, coadd_id, ra=rr, dec=dd, cache=True)
    psf = psfmod.GridInterpPSF(stamps, x, y)
    from functools import partial
    psf.fitfun = partial(psfmod.wise_psf_fit, psfstamp=(stamps, x, y),
                         grid=True)
    return psf

def read_wise(coadd_id, band, basedir, uncompressed=False,
              drop_first_dir=False, epoch=-1):
    assert((band == 1) or (band == 2))
    assert(len(coadd_id) == 8)

    imagefn = wise_filename(basedir, coadd_id, band, 'img-m',
                            uncompressed=uncompressed,
                            drop_first_dir=drop_first_dir, epoch=epoch)
    ivarfn = wise_filename(basedir, coadd_id, band, 'invvar-m',
                           uncompressed=uncompressed,
                           drop_first_dir=drop_first_dir, epoch=epoch)
    flagfn = wise_filename(basedir, coadd_id, band, 'msk',
                           uncompressed=uncompressed,
                           drop_first_dir=drop_first_dir, epoch=epoch)
    nmfn = wise_filename(basedir, coadd_id, band, 'n-m',
                         uncompressed=uncompressed,
                         drop_first_dir=drop_first_dir, epoch=epoch)
    nufn = wise_filename(basedir, coadd_id, band, 'n-u',
                         uncompressed=uncompressed,
                         drop_first_dir=drop_first_dir, epoch=epoch)

    im, hdr = fits.getdata(imagefn, header=True)
    sqivar = numpy.sqrt(fits.getdata(ivarfn))
    flag = fits.getdata(flagfn)
    nm = fits.getdata(nmfn)
    nu = fits.getdata(nufn)
    sqivar, flag = massage_isig_and_dim(sqivar, im, flag, band, nm, nu)
    return im, sqivar, flag, hdr

def ivarmap(isig, psfstamp):
    from scipy.signal import fftconvolve
    ivarim = fftconvolve(isig**2., psfstamp[::-1, ::-1]**2., mode='same')
    return ivarim

def brightlist(brightstars, coadd_id, band, basedir, uncompressed=False,
               drop_first_dir=False, epoch=-1):
    imagefn = wise_filename(basedir, coadd_id, band, 'img-m',
                            uncompressed=uncompressed,
                            drop_first_dir=drop_first_dir, epoch=epoch)
    hdr = fits.getheader(imagefn)
    blist = read_blist(brightstars, hdr['CRVAL1'], hdr['CRVAL2'], hdr, 3)
    return blist

def collapse_unwise_bitmask(bitmask, band):
    # 2^0 = bright star core and wings
    # 2^1 = PSF-based diffraction spike
    # 2^2 = optical ghost
    # 2^3 = first latent
    # 2^4 = second latent
    # 2^5 = AllWISE-like circular halo
    # 2^6 = bright star saturation
    # 2^7 = geometric diffraction spike

    assert((band == 1) or (band == 2))

    bits_w1 = OrderedDict([('core_wings', 2**0 + 2**1),
                           ('psf_spike', 2**27),
                           ('ghost', 2**25 + 2**26),
                           ('first_latent', 2**13 + 2**14),
                           ('second_latent', 2**17 + 2**18),
                           ('circular_halo', 2**23),
                           ('saturation', 2**4),
                           ('geom_spike', 2**29)])

    bits_w2 = OrderedDict([('core_wings', 2**2 + 2**3),
                           ('psf_spike', 2**28),
                           ('ghost', 2**11 + 2**12),
                           ('first_latent', 2**15 + 2**16),
                           ('second_latent', 2**19 + 2**20),
                           ('circular_halo', 2**24),
                           ('saturation', 2**5),
                           ('geom_spike', 2**30)])

    bits = (bits_w1 if (band == 1) else bits_w2)

    # hack to handle both scalar and array inputs
    result = 0*bitmask

    for i, feat in enumerate(bits.keys()):
        result += (2**i)*(numpy.bitwise_and(bitmask, bits[feat]) != 0)

    # int8 would be fine here, but astropy.io.fits seems to read this
    # as a boolean... so we waste the extra 8 bits.
    return result.astype('i2')

def collapse_extraflags(bitmask, band):
    bits_w1 = OrderedDict([('bright_off_edge', 2**7),
                           ('resolved_galaxy', 2**9),
                           ('big_object', 2**10),
                           ('possible_bright_star_centroid', 2**21),
                           ('crowdsat', extrabits['crowdsat']),
                           ('nebulosity', extrabits['nebulosity']),
                           ('nodeblend', crowdsource.nodeblend_maskbit),
                           ('sharp', crowdsource.sharp_maskbit)])

    bits_w2 = OrderedDict([('bright_off_edge', 2**8),
                           ('resolved_galaxy', 2**9),
                           ('big_object', 2**10),
                           ('possible_bright_star_centroid', 2**22),
                           ('crowdsat', extrabits['crowdsat']),
                           ('nebulosity', extrabits['nebulosity']),
                           ('nodeblend', crowdsource.nodeblend_maskbit),
                           ('sharp', crowdsource.sharp_maskbit)])

    bits = (bits_w1 if (band == 1) else bits_w2)

    # hack to handle both scalar and array inputs
    result = 0*bitmask

    for i, feat in enumerate(bits.keys()):
        result += (2**i)*(numpy.bitwise_and(bitmask, bits[feat]) != 0)

    # could fit in a byte, but astropy.io.fits reads these as booleans,
    # so we waste the byte...
    return result.astype('i2')

if __name__ == "__main__":
    try:
        print('Running on host: ' + str(os.environ.get('HOSTNAME')))
    except Exception:
        print("Couldn't retrieve hostname!")

    parser = argparse.ArgumentParser(description='Run crowdsource on unWISE coadd image')
    parser.add_argument('coadd_id', type=str, nargs=1)
    parser.add_argument('band', type=int, nargs=1)
    parser.add_argument('outfn', type=str, nargs=1)

    parser.add_argument('basedir', type=str, nargs='?', default='/global/projecta/projectdirs/cosmo/work/wise/outputs/merge/neo4/fulldepth')
    parser.add_argument('--refit-psf', '-r', default=False, action='store_true')
    parser.add_argument('--verbose', '-v', default=False, action='store_true')
    parser.add_argument('--uncompressed', '-u', default=False, action='store_true')
    parser.add_argument('--brightcat', '-b',
                        default=os.environ.get('TMASS_BRIGHT', ''), type=str)
    parser.add_argument('--modelfn', '-m', default='', type=str,
                        help='file name for model image, if desired')
    parser.add_argument('--infoimfn', '-i', default='', type=str,
                        help='file name for info image, if desired')
    parser.add_argument('--masknebulosity', '-n', action='store_true')
    parser.add_argument('--forcecat', type=str, default='')
    parser.add_argument('--startsky', type=str, default='')
    parser.add_argument('--startpsf', type=str, default='')
    parser.add_argument('--noskyfit', default=False, action='store_true')
    parser.add_argument('--threshold', default=5, type=float,
                        help='find sources down to threshold*sigma')
    parser.add_argument('--epoch', type=int, default=-1,
                        help='epoch number of time-resolved WISE coadd')
    parser.add_argument('--release', type=str, default='')

    args = parser.parse_args()

    coadd_id = args.coadd_id[0]
    band = args.band[0]
    basedir = args.basedir

    im, sqivar, flag, hdr = read_wise(coadd_id, band, basedir,
                                      uncompressed=args.uncompressed,
                                      epoch=args.epoch)
    if len(args.startsky) > 0:
        startsky = fits.getdata(args.startsky, 'SKY')
    else:
        startsky = numpy.nan
    flag_orig = fits.getdata(wise_filename(basedir, coadd_id, band, 'msk',
                                           uncompressed=args.uncompressed,
                                           epoch=args.epoch))

    if args.masknebulosity:
        import nebulosity_mask
        nebfn = os.path.join(os.environ['WISE_DIR'], 'dat', 'nebnet',
                             'weights1', '1st_try')
        nebmod = nebulosity_mask.load_model(nebfn)
        nebmask = nebulosity_mask.gen_mask_wise(nebmod, im) == 2
        if numpy.any(nebmask):
            flag |= nebmask * extrabits['nebulosity']
            flag |= nebmask * crowdsource.sharp_maskbit
            print('Masking nebulosity, %5.2f' % (
                numpy.sum(nebmask)/1./numpy.sum(numpy.isfinite(nebmask))))

    psf = wise_psf_grid(band, coadd_id, basedir, epoch=args.epoch)
    if len(args.startpsf) > 0:
        startpsf = fits.getdata(args.startpsf, 'PSF').astype('f4')
        # there can be some endianness issues; astype('f4') converts to native
        modpsf = psf(1024, 1024, stampsz=psf.stamp.shape[-1])
        resid = startpsf-modpsf
        # need not sum to zero.
        newstamps = psf.stamp / psf.normstamp[:, :, None, None]
        newstamps += resid
        psf = psfmod.GridInterpPSF(newstamps, psf.x, psf.y)
        from functools import partial
        psf.fitfun = partial(psfmod.wise_psf_fit,
                             psfstamp=(newstamps, psf.x, psf.y), grid=True)

    if len(args.brightcat) > 0:
        brightstars = fits.getdata(args.brightcat)
        blist = brightlist(brightstars, coadd_id, band, basedir,
                           uncompressed=args.uncompressed, epoch=args.epoch)
    else:
        print('No bright star catalog, not marking bright stars.')
        blist = None

    if args.verbose:
        t0 = time.time()
        print('Starting %s, band %d, at %s' % (coadd_id, band, time.ctime()))
        sys.stdout.flush()

    if len(args.forcecat) == 0:
        res = crowdsource.fit_im(
            im, psf, weight=sqivar, dq=flag, refit_psf=args.refit_psf,
            verbose=args.verbose, ntilex=4, ntiley=4, derivcentroids=True,
            maxstars=30000*16, fewstars=50*16, blist=blist,
            threshold=args.threshold, psfvalsharpcutfac=0.5,
            psfsharpsat=0.8)
    else:
        forcecat = fits.getdata(args.forcecat, 1)
        x, y = forcecat['x'], forcecat['y']
        res = crowdsource.fit_im_force(
            im, x, y, psf, weight=sqivar, dq=flag, refit_psf=args.refit_psf,
            blist=blist, refit_sky=(not args.noskyfit),
            startsky=startsky, psfderiv=False, psfvalsharpcutfac=0.5,
            psfsharpsat=0.8)
    cat, model, sky, psf = res
    print('Finishing %s, band %d; %d sec elapsed.' %
          (coadd_id, band, time.time()-t0))

    outfn = args.outfn[0]

    x = cat['x']
    y = cat['y']

    wcs0 = wcs.WCS(hdr)
    ra, dec = wcs0.all_pix2world(y, x, 0)
    coadd_ids = numpy.zeros(len(ra), dtype='a8')
    bands = numpy.zeros(len(ra), dtype='i4')
    ids = numpy.zeros(len(ra), dtype='U20')
    coadd_ids[:] = coadd_id
    bands[:] = band
    if len(args.release) == 0:
        ids = ['%sw%1do%07d' % (coadd_id, band, num) for num in range(len(ra))]
    else:
        ids = ['%sw%1do%07dr%s' % (coadd_id, band, num, args.release)
               for num in range(len(ra))]

    nmfn = wise_filename(basedir, coadd_id, band, 'n-m',
                         uncompressed=args.uncompressed, epoch=args.epoch)
    nmim = fits.getdata(nmfn)
    nms = crowdsource.extract_im(cat['x'], cat['y'], nmim)
    flags_unwise = crowdsource.extract_im(
        cat['x'], cat['y'], collapse_unwise_bitmask(flag_orig, band))
    flags_infoim = collapse_extraflags(flag, band)
    flags_info = crowdsource.extract_im(cat['x'], cat['y'], flags_infoim)
    # cast to i2; astropy.io.fits seems to fail for bools?
    primary = unwise_primary.is_primary(coadd_id, ra, dec).astype('i2')

    import numpy.lib.recfunctions as rfn
    cat = rfn.drop_fields(cat, ['flags'])
    cat = rfn.append_fields(
        cat, ['ra', 'dec', 'coadd_id', 'band', 'unwise_detid', 'nm',
              'primary', 'flags_unwise', 'flags_info'],
        [ra, dec, coadd_ids, bands, ids, nms, primary, flags_unwise,
         flags_info])

    hdr['EXTNAME'] = 'PRIMARY'
    fits.writeto(outfn, None, hdr)
    fits.append(outfn, cat)
    if len(args.modelfn) > 0:
        hdulist = fits.open(args.modelfn, mode='append')
        compkw = {'compression_type': 'GZIP_1',
                  'quantize_method': 2, 'quantize_level': -0.5,
                  'tile_size': model.shape}
        hdr['EXTNAME'] = 'model'
        hdulist.append(fits.CompImageHDU(model, hdr, **compkw))
        hdr['EXTNAME'] = 'sky'
        hdulist.append(fits.CompImageHDU(sky, hdr, **compkw))
        hdulist.close(closed=True)

    if len(args.infoimfn) > 0:
        psffluxivar = ivarmap(sqivar, psf(1024, 1024, stampsz=59)).astype('f4')
        psfstamp = psf(1024, 1024, stampsz=325)
        hdulist = fits.open(args.infoimfn, mode='append')
        compkw = {'compression_type': 'GZIP_1',
                  'quantize_method': 2,
                  'tile_size': psffluxivar.shape}
        hdr['EXTNAME'] = 'psffluxivar'
        hdulist.append(fits.CompImageHDU(psffluxivar, hdr, **compkw))
        hdr['EXTNAME'] = 'infoflags'
        compkw = {'compression_type': 'GZIP_1',
                  'tile_size': flags_infoim.shape}
        # must recast flags_infoim as a u1; unsigned isn't supported
        # in tables, but signed int8 isn't supported in CompImageHDU.
        # ugh.
        hdulist.append(fits.CompImageHDU(flags_infoim.astype('u1'),
                                         hdr, **compkw))
        hdulist.append(fits.ImageHDU(psfstamp, None, name='psf'))
        hdulist.close(closed=True)


import numpy as np

import pyximport
pyximport.install(setup_args={"include_dirs":np.get_include()},build_dir="build", build_in_temp=False)

from align_fill import cython_fill_table

import sys

import torch

sys.path.append("../../../espnet")
from espnet.asr.asr_utils import get_model_conf

import os
from pathlib import Path
from time import time

import argparse

parser = argparse.ArgumentParser()
parser.add_argument("model_path")
parser.add_argument("data_path")
parser.add_argument("eval_path")
parser.add_argument('--start_win', type=int, default=8000)
args = parser.parse_args()

max_prob = -10000000000.0

def align(lpz, char_list, ground_truth, utt_begin_indices, skip_prob):
    blank = 0
    print("Audio length: " + str(lpz.shape[0]))
    print("Text length: " + str(len(ground_truth)))
    if len(ground_truth) > lpz.shape[0] and skip_prob <= max_prob:
        raise AssertionError("Audio is shorter than text!")
    window_len = args.start_win

    # Try multiple window lengths if it fails
    while True:
        # Create table which will contain alignment probabilities
        table = np.zeros([min(window_len, lpz.shape[0]), len(ground_truth)], dtype=np.float32)
        table.fill(max_prob)
        # Use array to log window offsets per character
        offsets = np.zeros([len(ground_truth)], dtype=np.int)

        # Run actual alignment
        t, c = cython_fill_table(table, lpz.astype(np.float32), np.array(ground_truth), offsets, np.array(utt_begin_indices), blank, skip_prob)

        print("Max prob: " + str(table[:, c].max()) + " at " + str(t))

        # Backtracking
        timings = np.zeros([len(ground_truth)])
        char_probs = np.zeros([lpz.shape[0]])
        char_list = [''] * lpz.shape[0]
        current_prob_sum = 0
        try:
            # Do until start is reached
            while t != 0 or c != 0:
                # Calculate the possible transition probabilities towards the current cell
                min_s = None
                min_switch_prob_delta = np.inf
                max_lpz_prob = max_prob
                for s in range(ground_truth.shape[1]):
                    if ground_truth[c, s] != -1:
                        offset = offsets[c] - (offsets[c - 1 - s] if c - s > 0 else 0)
                        switch_prob = lpz[t + offsets[c], ground_truth[c, s]] if c > 0 else max_prob
                        est_switch_prob = table[t, c] - table[t - 1 + offset, c - 1 - s]
                        if abs(switch_prob - est_switch_prob) < min_switch_prob_delta:
                            min_switch_prob_delta = abs(switch_prob - est_switch_prob)
                            min_s = s

                        max_lpz_prob = max(max_lpz_prob, switch_prob)

                stay_prob = max(lpz[t + offsets[c], blank], max_lpz_prob) if t > 0 else max_prob
                est_stay_prob = table[t, c] - table[t - 1, c]

                # Check which transition has been taken
                if abs(stay_prob - est_stay_prob) > min_switch_prob_delta:
                    # Apply reverse switch transition
                    if c > 0:
                        # Log timing and character - frame alignment
                        for s in range(0, min_s + 1):
                            timings[c - s] = (offsets[c] + t) * 10 * 4 / 1000
                        char_probs[offsets[c] + t] = max_lpz_prob
                        char_list[offsets[c] + t] = train_args.char_list[ground_truth[c, min_s]]
                        current_prob_sum = 0

                    c -= 1 + min_s
                    t -= 1 - offset

                else:
                    # Apply reverse stay transition
                    char_probs[offsets[c] + t] = stay_prob
                    char_list[offsets[c] + t] = "ε"
                    t -= 1
        except IndexError:
            # If the backtracking was not successful this usually means the window was too small
            window_len *= 2
            print("IndexError: Trying with win len: " + str(window_len))
            if window_len < 100000:
                continue
            else:
                raise

        break

    return timings, char_probs, char_list

def prepare_text(text):
    # Prepares the given text for alignment
    # Therefore we create a matrix of possible character symbols to represent the given text

    # Create list of char indices depending on the models char list
    ground_truth = "#"
    utt_begin_indices = []
    for utt in text:
        # Only one space in-between
        if ground_truth[-1] != " ":
            ground_truth += " "

        # Start new utterance remeber index
        utt_begin_indices.append(len(ground_truth) - 1)

        # Add chars of utterance
        for char in utt:
            if char.isspace():
                if ground_truth[-1] != " ":
                    ground_truth += " "
            elif char in train_args.char_list and char not in [ ".", ",", "-", "?", "!", ":", "»", "«", ";", "'", "›", "‹", "(", ")"]:
                ground_truth += char

    # Add space to the end
    if ground_truth[-1] != " ":
        ground_truth += " "
    utt_begin_indices.append(len(ground_truth) - 1)

    # Create matrix where first index is the time frame and the second index is the number of letters the character symbol spans
    max_char_len = max([len(c) for c in train_args.char_list])
    ground_truth_mat = np.ones([len(ground_truth), max_char_len], np.int) * -1
    for i in range(len(ground_truth)):
        for s in range(max_char_len):
            if i-s < 0:
                continue
            span = ground_truth[i-s:i+1]
            span = span.replace(" ", '▁')
            if span in train_args.char_list:
                ground_truth_mat[i, s] = train_args.char_list.index(span)

    return ground_truth_mat, utt_begin_indices

def write_output(out_path, utt_begin_indices, char_probs):
    # Uses char-wise alignments to get utterance-wise alignmentes and writes them into the given file

    with open(str(out_path), 'w') as outfile:
        outfile.write(str(path_wav.name) + '\n')

        def compute_time(index, type):
            # Compute start and end time of utterance.
            middle = (timings[index] + timings[index - 1]) / 2
            if type == "begin":
                return max(timings[index + 1] - 0.5, middle)
            elif type == "end":
                return min(timings[index - 1] + 0.5, middle)

        for i in range(len(text)):
            start = compute_time(utt_begin_indices[i], "begin")
            end = compute_time(utt_begin_indices[i + 1], "end")
            start_t = int(round(start * 1000 / 40))
            end_t = int(round(end * 1000 / 40))

            # Compute confidence score by using the min mean probability after splitting into segments of 30 frames
            n = 30
            if end_t == start_t:
                min_avg = 0
            elif end_t - start_t <= n:
                min_avg = char_probs[start_t:end_t].mean()
            else:
                min_avg = 0
                for t in range(start_t, end_t - n):
                    min_avg = min(min_avg, char_probs[t:t + n].mean())

            outfile.write(str(start) + " " + str(end) + " " + str(min_avg) + " | " + text[i] + '\n')

model_path = args.model_path
model_conf = None

# read training config
idim, odim, train_args = get_model_conf(model_path, model_conf)

space_id = train_args.char_list.index('▁')
train_args.char_list[0] = "ε"
train_args.char_list = [c.lower() for c in train_args.char_list]

data_path = Path(args.data_path)
eval_path = Path(args.eval_path)

for path_wav in data_path.glob("*.wav"):

    chapter_sents = data_path / path_wav.name.replace(".wav", ".txt")
    chapter_prob = eval_path / path_wav.name.replace(".wav", ".npz")
    out_path = eval_path / path_wav.name.replace(".wav", ".txt")

    with open(str(chapter_sents), "r") as f:
        text = [t.strip() for t in f.readlines()]

    lpz = np.load(str(chapter_prob))["arr_0"]

    print("Syncing " + str(path_wav))

    ground_truth_mat, utt_begin_indices = prepare_text(text)

    try:
        timings, char_probs, char_list = align(lpz, train_args.char_list, ground_truth_mat, utt_begin_indices, max_prob)
    except AssertionError:
        print("Skipping: Audio is shorter than text")
        continue

    write_output(out_path, utt_begin_indices, char_probs)


import tensorflow as tf
import numpy as np
import math
import time
import os
import tensornets as nets
import cv2
import dlib
import sys
from age_gender_model import run_test

tf.app.flags.DEFINE_string('img_path', "/images/trump.jpg", 'the path of imgaine')
tf.app.flags.DEFINE_string('model_path', "train_models/age_gender_tensornets_wiki_cropface_model_20180519_123211", 'how often to print training status')
FLAGS = tf.app.flags.FLAGS

def main(_):

    tf.reset_default_graph()

    detector = dlib.get_frontal_face_detector() #获取人脸分类器

    with tf.Session() as sess:
        img = cv2.imread(path, cv2.IMREAD_COLOR)

        dets = detector(img, 1) #使用detector进行人脸检测 dets为返回的结果
        HH, WW, _ = img.shape
        faces_tensor = np.array([])
        for index, face in enumerate(dets):

            # 在图片中标注人脸，并显示
            left = face.left()
            top = face.top()
            right = face.right()
            bottom = face.bottom()
            w = right - left
            h = bottom - top
            left = int(max(0, left - 0.2 * w))
            top = int(max(0, top - 0.4 * h))
            right = int(min(WW, right + 0.2 * w))
            bottom = int(min(HH, bottom + 0.2 * h))
            face_img = img[top:bottom, left:right].copy()
            cv2.rectangle(img, (left, top), (right, bottom), (0, 255, 0), 3)
            face_img = cv2.resize(face_img, (224, 224))
            np.append(faces_tensor, face_img)
        print("faces_tensor.shape", faces_tensor.shape)
        age, gender = run_test(sess, input_tensor, model_path)

        for i in range(len(dets)):
            cv2.rectangle(img, (left, top), (right, bottom), (0, 255, 0), 3)
            font = cv2.FONT_HERSHEY_SIMPLEX
            if top - 10 < 0:
                top = top + 20
            else:
                top = top - 10
            gender_text = "male" if gender[i] == 1 else "female"
            cv2.putText(frame,str(age[i]) + ", " + gender_text,(left,top), font, 1,(0,255,0),2,cv2.LINE_AA)
            #print("age: " + str(age[i]) + "gender: " + str(gender[i]))

        pass

if __name__ == '__main__':
    tf.app.run()


# -*- coding: utf-8 -*-

"""
Copyright 2019 Tae Hwan Jung

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

class XLNet(nn.Module):
    """
        Defines a Transformer-XL computation graph with additional
        support for XLNet.

        Args:

        inp_k: int32 Tensor in shape [len, bsz], the input token IDs.
        seg_id: int32 Tensor in shape [len, bsz], the input segment IDs.
        input_mask: float32 Tensor in shape [len, bsz], the input mask.
          0 for real tokens and 1 for padding.
        mems: a list of float32 Tensors in shape [mem_len, bsz, d_model], memory
          from previous batches. The length of the list equals n_layer.
          If None, no memory is used.
        perm_mask: float32 Tensor in shape [len, len, bsz].
          If perm_mask[i, j, k] = 0, i attend to j in batch k;
          if perm_mask[i, j, k] = 1, i does not attend to j in batch k.
          If None, each position attends to all the others.
        target_mapping: float32 Tensor in shape [num_predict, len, bsz].
          If target_mapping[i, j, k] = 1, the i-th predict in batch k is
          on the j-th token.
          Only used during pretraining for partial prediction.
          Set to None during finetuning.
        inp_q: float32 Tensor in shape [len, bsz].
          1 for tokens with losses and 0 for tokens without losses.
          Only used during pretraining for two-stream attention.
          Set to None during finetuning.

        n_layer: int, the number of layers.
        d_model: int, the hidden size.
        n_head: int, the number of attention heads.
        d_head: int, the dimension size of each attention head.
        d_inner: int, the hidden size in feed-forward layers.
        ff_activation: str, "relu" or "gelu".
        n_token: int, the vocab size.

        dropout: float, dropout rate.
        dropatt: float, dropout rate on attention probabilities.

        mem_len: int, the number of tokens to cache.
        reuse_len: int, the number of tokens in the currect batch to be cached
          and reused in the future.
        bi_data: bool, whether to use bidirectional input pipeline.
          Usually set to True during pretraining and False during finetuning.
        clamp_len: int, clamp all relative distances larger than clamp_len.
          -1 means no clamping.

      """
    def __init__(self, n_token, n_layer, n_head, d_head, d_inner, d_model, dropout, dropatt,
                 attn_type, bi_data, clamp_len, same_length, reuse_len, mem_len):
        super(XLNet, self).__init__()

        self.n_token = n_token
        self.n_layer = n_layer
        self.n_head = n_head
        self.d_head = d_head
        self.d_inner = d_inner
        self.d_model = d_model
        self.dropout = dropout
        self.dropatt = dropatt
        self.attn_type = attn_type
        self.bi_data = bi_data
        self.clamp_len = clamp_len
        self.same_length = same_length
        self.reuse_len = reuse_len
        self.mem_len = mem_len

        self.embedding = nn.Embedding(n_token, d_model)
        self.Dropout = nn.Dropout(p=dropout)
        self.DropAttn = nn.Dropout(p=dropatt)

        self.r_w_bias = nn.Parameter(torch.randn(self.n_layer,
                                                  self.n_head,self.d_head))
        self.r_r_bias = nn.Parameter(torch.randn(self.n_layer,
                                                  self.n_head, self.d_head))

        ##### Segment embedding
        self.r_s_bias = nn.Parameter(torch.randn(self.n_layer,
                                                  self.n_head,self.d_head))

        self.seg_embed = nn.Parameter(torch.randn(self.n_layer, 2,
                                                   self.n_head, self.d_head))

        self.mask_emb = nn.Parameter(torch.randn(1, 1, d_model))

        # post-attention projection (back to `d_model`)
        self.proj_o = nn.Parameter(torch.randn(self.d_model,
                                                self.n_head, self.d_head))

        #### Project hidden states to a specific head with a 4D-shape.
        self.q_proj_weight = nn.Parameter(torch.randn(self.d_model,
                                                       self.n_head, self.d_head))
        self.k_proj_weight = nn.Parameter(torch.randn(self.d_model,
                                                       self.n_head, self.d_head))
        self.v_proj_weight = nn.Parameter(torch.randn(self.d_model,
                                                       self.n_head, self.d_head))
        self.r_proj_weight = nn.Parameter(torch.randn(self.d_model,
                                                       self.n_head, self.d_head))

        self.layer_norm = nn.LayerNorm(d_model)

        self.conv1 = nn.Linear(d_model, d_inner)
        self.conv2 = nn.Linear(d_inner, d_model)
        self.relu = nn.ReLU(inplace=True)

        self.softmax_b = nn.Parameter(torch.zeros(self.n_token))

    def gelu(self, x):
        """Gaussian Error Linear Unit.

        This is a smoother version of the RELU.
        Original paper: https://arxiv.org/abs/1606.08415
        Args:
          x: float Tensor to perform activation.

        Returns:
          `x` with the GELU activation applied.
        """
        cdf = 0.5 * (1.0 + torch.tanh(
            (np.sqrt(2 / np.pi) * (x + 0.044715 * torch.pow(x, 3)))))
        return x * cdf

    def rel_shift(self, x, klen=-1):
        """perform relative shift to form the relative attention score."""
        x_size = x.shape

        x = torch.reshape(x, [x_size[1], x_size[0], x_size[2], x_size[3]])
        x = x[1:, 0:, 0:, 0:] # tf.slice(x, [1, 0, 0, 0], [-1, -1, -1, -1])
        x = torch.reshape(x, [x_size[0], x_size[1] - 1, x_size[2], x_size[3]])
        x = x[0:, 0:klen, 0:, 0:] # tf.slice(x, [0, 0, 0, 0], [-1, klen, -1, -1])

        return x

    def positionwise_ffn(self, inp, activation_type='relu'):

        """Position-wise Feed-forward Network."""
        output = self.conv1(inp)
        output = self.Dropout(output)
        if activation_type == 'relu':
            output = self.relu(output)
        elif activation_type == 'gelu':
            output = self.gelu(output)
        else:
            raise ValueError('Unsupported activation type {}'.format(activation_type))

        output = self.layer_norm(output + inp)
        return output

    def post_attention(self, h, attn_vec, residual=True):
        """Post-attention processing."""

        # post-attention projection (back to `d_model`)
        attn_out = torch.einsum('ibnd,hnd->ibh', attn_vec, self.proj_o)

        attn_out = self.Dropout(attn_out)
        if residual:
            output = self.layer_norm(attn_out + h)
        else:
            output = self.layer_norm(attn_out)

        return output

    def head_projection(self, h, name):
        """Project hidden states to a specific head with a 4D-shape."""
        proj_weight = None
        if name == 'q':
            proj_weight = self.q_proj_weight
        elif name == 'k':
            proj_weight = self.k_proj_weight
        elif name =='v':
            proj_weight = self.v_proj_weight
        elif name == 'r':
            proj_weight = self.r_proj_weight
        else:
            raise ValueError('Unknown `name` {}.'.format(name))

        head = torch.einsum('ibh,hnd->ibnd', h, proj_weight)

        return head

    def rel_attn_core(self, q_head, k_head_h, v_head_h, k_head_r, seg_embed, seg_mat,
                      r_w_bias, r_r_bias, r_s_bias, attn_mask, scale):

        """Core relative positional attention operations."""

        # content based attention score
        ac = torch.einsum('ibnd,jbnd->ijbn', q_head + r_w_bias, k_head_h)

        # position based attention score
        bd = torch.einsum('ibnd,jbnd->ijbn', q_head + r_r_bias, k_head_r)
        bd = self.rel_shift(bd, klen=ac.shape[1])

        # segment based attention score
        if seg_mat is None:
            ef = 0
        else:
            ef = torch.einsum('ibnd,snd->ibns', q_head + r_s_bias, seg_embed)
            ef = torch.einsum('ijbs,ibns->ijbn', seg_mat, ef)

        # merge attention scores and perform masking
        attn_score = (ac + bd + ef) * scale
        if attn_mask is not None:
            # attn_score = attn_score * (1 - attn_mask) - 1e30 * attn_mask
            attn_score = attn_score - 1e30 * attn_mask

        # attention probability
        attn_prob = F.softmax(attn_score, dim=1)
        attn_prob = self.DropAttn(attn_prob)

        # attention output
        attn_vec = torch.einsum('ijbn,jbnd->ibnd', attn_prob, v_head_h)

        return attn_vec

    def rel_multihead_attn(self, h, r, r_w_bias, r_r_bias, seg_mat, r_s_bias, seg_embed,
                           attn_mask, mems, d_model, n_head, d_head, dropout, dropatt):
        """Multi-head attention with relative positional encoding."""

        scale = 1 / (d_head ** 0.5)
        if mems is not None and len(mems.size()) > 1:
            cat = torch.cat([mems, h], dim=0)
        else:
            cat = h

        # content heads
        q_head_h = self.head_projection(h, 'q')
        k_head_h = self.head_projection(cat, 'k')
        v_head_h = self.head_projection(cat, 'v')

        # positional heads
        k_head_r = self.head_projection(r, 'r')

        # core attention ops
        attn_vec = self.rel_attn_core(
            q_head_h, k_head_h, v_head_h, k_head_r, seg_embed, seg_mat, r_w_bias,
            r_r_bias, r_s_bias, attn_mask, scale)

        # post processing
        output = self.post_attention(h, attn_vec)

        return output

    def two_stream_rel_attn(self, h, g, r, mems, r_w_bias, r_r_bias, seg_mat, r_s_bias,
                            seg_embed, attn_mask_h, attn_mask_g, target_mapping):
        scale = 1 / (self.d_head ** 0.5)

        # content based attention score
        if mems is not None and len(mems.size()) > 1:
            cat = torch.cat([mems, h], dim=0)
        else:
            cat = h

        # content-based key head
        k_head_h = self.head_projection(cat, 'k')

        # content-based value head
        v_head_h = self.head_projection(cat, 'v')

        # position-based key head
        k_head_r = self.head_projection(r, 'r')

        ##### h-stream
        # content-stream query head
        q_head_h = self.head_projection(h, 'q')

        # core attention ops
        # hˆ(m)_zt = LayerNorm(h^(m-1)_zt + RelAttn(h^(m-1)_zt + [h~^(m-1), hT(m-1)_z<=t]))
        attn_vec_h = self.rel_attn_core(
            q_head_h, k_head_h, v_head_h, k_head_r, seg_embed, seg_mat, r_w_bias,
            r_r_bias, r_s_bias, attn_mask_h, scale)

        # post processing
        output_h = self.post_attention(h, attn_vec_h)

        ##### g-stream
        # query-stream query head
        q_head_g = self.head_projection(g, 'q')

        # core attention ops
        # gˆ(m)_zt = LayerNorm(g^(m-1)_zt + RelAttn(g^(m-1)_zt + [h~^(m-1), hT(m-1)_z<=t]))
        if target_mapping is not None:
            q_head_g = torch.einsum('mbnd,mlb->lbnd', q_head_g, target_mapping)
            attn_vec_g = self.rel_attn_core(
                q_head_g, k_head_h, v_head_h, k_head_r, seg_embed, seg_mat, r_w_bias,
                r_r_bias, r_s_bias, attn_mask_g, scale)
            attn_vec_g = torch.einsum('lbnd,mlb->mbnd', attn_vec_g, target_mapping)
        else:
            attn_vec_g = self.rel_attn_core(
                q_head_g, k_head_h, v_head_h, k_head_r, seg_embed, seg_mat, r_w_bias,
                r_r_bias, r_s_bias, attn_mask_g, scale)

        # post processing
        output_g = self.post_attention(g, attn_vec_g)

        return output_h, output_g

    def _create_mask(self, qlen, mlen, dtype, same_length=False):
        """create causal attention mask."""
        # [[0,1,1],
        #  [0,0,1],
        #  [0,0,0]]
        attn_mask = torch.ones([qlen, qlen], dtype=dtype)
        mask_u = torch.triu(attn_mask) # Upper triangular part.
        mask_dia = torch.tril(attn_mask) & torch.triu(attn_mask) # Diagonal. Figure 2(c)
        attn_mask_pad = torch.zeros([qlen, mlen], dtype=dtype)
        ret = torch.cat([attn_mask_pad, mask_u - mask_dia], dim=1) # [qlen, mlen]
        if same_length:
            # [[0,1,1],
            #  [1,0,1],
            #  [1,1,0]]
            mask_l = torch.tril(attn_mask) # Lower triangular part.
            ret = torch.cat([ret[:, :qlen] + mask_l - mask_dia, ret[:, qlen:]], dim=1)

        return ret.type(dtype=torch.float32) # [qlen, qlen]

    def positional_embedding(self, pos_seq, inv_freq):
        sinusoid_inp = torch.einsum('i,d->id', pos_seq, inv_freq)
        pos_emb = torch.cat([torch.sin(sinusoid_inp), torch.cos(sinusoid_inp)], dim=-1)
        pos_emb = pos_emb[:, None, :]

        return pos_emb

    def _cache_mem(self, curr_out, prev_mem, mem_len, reuse_len=None):
        """cache hidden states into memory."""

        with torch.no_grad():
            if mem_len is None or mem_len == 0:
                return None
            else:
                if reuse_len is not None and reuse_len > 0:
                    curr_out = curr_out[:reuse_len]

                if prev_mem is None:
                    new_mem = curr_out[-mem_len:]
                else:
                    new_mem = torch.cat([prev_mem, curr_out], dim=0)[-mem_len:]

            return new_mem

    def relative_positional_encoding(self, qlen, klen, d_model, clamp_len, attn_type,
                                     bi_data, bsz=None, dtype=None):
        """create relative positional encoding."""

        freq_seq = torch.arange(0, d_model, 2.0)
        if dtype is not None and dtype != torch.float32:
            freq_seq = freq_seq.type(dtype)
        inv_freq = 1 / (10000 ** (freq_seq / d_model))

        if attn_type == 'bi':
            # beg, end = klen - 1, -qlen
            beg, end = klen, -qlen
        elif attn_type == 'uni':
            # beg, end = klen - 1, -1
            beg, end = klen, -1
        else:
            raise ValueError('Unknown `attn_type` {}.'.format(attn_type))

        if bi_data and bsz%2 is 0:
            fwd_pos_seq = torch.arange(beg, end, -1.0)
            bwd_pos_seq = torch.arange(-beg, -end, 1.0)

            if dtype is not None and dtype != torch.float32:
                fwd_pos_seq = fwd_pos_seq.type(dtype=dtype)
                bwd_pos_seq = bwd_pos_seq.type(dtype=dtype)

            if clamp_len > 0:
                fwd_pos_seq = torch.clamp(fwd_pos_seq, -clamp_len, clamp_len)
                bwd_pos_seq = torch.clamp(bwd_pos_seq, -clamp_len, clamp_len)

            fwd_pos_emb = self.positional_embedding(fwd_pos_seq, inv_freq)
            bwd_pos_emb = self.positional_embedding(bwd_pos_seq, inv_freq)

            pos_emb = torch.cat([fwd_pos_emb, bwd_pos_emb], dim=1)
        else:
            fwd_pos_seq = torch.arange(beg, end, -1.0)
            if dtype is not None and dtype != torch.float32:
                fwd_pos_seq = fwd_pos_seq.type(dtype=dtype)
            if clamp_len > 0:
                fwd_pos_seq = torch.clamp(fwd_pos_seq, -clamp_len, clamp_len)
            pos_emb = self.positional_embedding(fwd_pos_seq, inv_freq)

        return pos_emb

    def forward(self, inp_k, seg_id, input_mask, mems, perm_mask, target_mapping, inp_q):
        new_mems = []

        bsz = inp_k.shape[1]
        qlen = inp_k.shape[0]
        mlen = mems[0].size(0) if mems is not None else 0
        klen = mlen + qlen

        ##### Attention mask
        # causal attention mask
        if self.attn_type == 'uni':
            attn_mask = self._create_mask(qlen, mlen, torch.int64, self.same_length)
            attn_mask = attn_mask[:, :, None, None]
        elif self.attn_type == 'bi':
            attn_mask = None
        else:
            raise ValueError('Unsupported attention type: {}'.format(self.attn_type))

        # data mask: input mask & perm mask
        if input_mask is not None and perm_mask is not None:
            data_mask = input_mask[None] + perm_mask
        elif input_mask is not None and perm_mask is None:
            data_mask = input_mask[None]
        elif input_mask is None and perm_mask is not None:
            data_mask = perm_mask
        else:
            data_mask = None

        if data_mask is not None:
            # all mems can be attended to
            mems_mask = torch.zeros([data_mask.shape[0], mlen, bsz],
                                 dtype=torch.float32)
            data_mask = torch.cat([mems_mask, data_mask], dim=1)
            if attn_mask is None:
                attn_mask = data_mask[:, :, :, None]
            else:
                attn_mask += data_mask[:, :, :, None]

        if attn_mask is not None:
            attn_mask = attn_mask.gt(0).type(torch.float32)

        if attn_mask is not None:
            non_tgt_mask = -torch.eye(qlen, dtype=torch.float32) # [qlen, qlen]
            non_tgt_mask = torch.cat([torch.zeros([qlen, mlen], dtype=torch.float32), # [qlen, klen]
                                        non_tgt_mask],
                                        dim=-1)
            non_tgt_mask = (attn_mask +
                            non_tgt_mask[:, :, None, None]).gt(0).type(dtype=torch.float32)
        else:
            non_tgt_mask = None

        ##### Word embedding
        lookup_table = self.embedding
        word_emb_k = lookup_table(inp_k)

        if inp_q is not None:
            if target_mapping is not None:
                word_emb_q = self.mask_emb.repeat(target_mapping.shape[0], bsz, 1)
            else:
                inp_q_ext = inp_q[:, :, None]
                word_emb_q = inp_q_ext * self.mask_emb + (1 - inp_q_ext) * word_emb_k

        #### Figure 2(a), Content Stream(Original Attention), h^(0)_t = e(x_i) = e(inp_k)
        output_h = self.Dropout(word_emb_k)
        if inp_q is not None:
            #### Query Stream, g^(0)_t = w
            #### the first layer query stream is initialized with a trainable vector
            output_g = self.Dropout(word_emb_q)

        ##### Segment embedding
        # paper
        # Given a pair of positions i and j in the sequence, if
        # i and j are from the same segment
        if seg_id is not None:
            # Convert `seg_id` to one-hot `seg_mat`
            mem_pad = torch.zeros([mlen, bsz], dtype=torch.int32)
            cat_ids = torch.cat([mem_pad, seg_id], dim=0)

            # `1` indicates not in the same segment [qlen x klen x bsz]
            seg_mat = (~torch.eq(seg_id[:, None], cat_ids[None, :])).type(torch.long)
            seg_mat = torch.eye(2, dtype=torch.float32)[seg_mat]
        else:
            seg_mat = None

        ##### Positional encoding
        pos_emb = self.relative_positional_encoding(
            qlen, klen, self.d_model, self.clamp_len, self.attn_type, self.bi_data,
            bsz=bsz, dtype=torch.float32)
        pos_emb = self.Dropout(pos_emb)

        ##### Attention layers
        if mems is None:
            mems = [None] * self.n_layer

        for i in range(self.n_layer):
            # cache new mems
            new_mems.append(self._cache_mem(output_h, mems[i], self.mem_len, self.reuse_len))

            # segment bias
            if seg_id is None:
                r_s_bias_i = None
                seg_embed_i = None
            else:
                r_s_bias_i = self.r_s_bias[i]
                seg_embed_i = self.seg_embed[i]

            if inp_q is not None:
                output_h, output_g = self.two_stream_rel_attn(
                    h=output_h,
                    g=output_g,
                    r=pos_emb,
                    r_w_bias= self.r_w_bias[i],
                    r_r_bias= self.r_r_bias[i],
                    seg_mat=seg_mat,
                    r_s_bias=r_s_bias_i,
                    seg_embed=seg_embed_i,
                    attn_mask_h=non_tgt_mask,
                    attn_mask_g=attn_mask,
                    mems=mems[i],
                    target_mapping=target_mapping)
            else:
                output_h = self.rel_multihead_attn(
                    h=output_h,
                    r=pos_emb,
                    r_w_bias=self.r_w_bias[i],
                    r_r_bias=self.r_r_bias[i],
                    seg_mat=seg_mat,
                    r_s_bias=r_s_bias_i,
                    seg_embed=seg_embed_i,
                    attn_mask=non_tgt_mask,
                    mems=mems[i])

            if inp_q is not None:
                output_g = self.positionwise_ffn(inp=output_g)

            output_h = self.positionwise_ffn(inp=output_h)

        if inp_q is not None:
            output = self.Dropout(output_g)
        else:
            output = self.Dropout(output_h)

        logits = torch.einsum('ibd,nd->ibn', output, lookup_table.weight) + self.softmax_b

        return logits, new_mems


import yaml
import numpy as np
from easydict import EasyDict as edict

config = edict()

#test image path
config.TEST_YAML = '/../experiments/fcis/cfgs/fcis_coco_demo.yaml'
config.FG = 'baby1.jpg'
config.BG = 'bg.jpg'
config.EX_MXNET = '../external/mxnet'
config.TEST_CLASS_NAME = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat',
                   'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog',
                   'horse',
                   'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie',
                   'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
                   'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon',
                   'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut',
                   'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse',
                   'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book',
                   'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']
config.PREMODEL = '/../model/fcis_coco'
#
config.MXNET_VERSION = ''
config.output_path = ''
config.symbol = ''
config.gpus = ''
config.CLASS_AGNOSTIC = True
config.SCALES = [(600, 1000)]  # first is scale (the shorter side); second is max size
config.MASK_SIZE = 21
config.BINARY_THRESH = 0.4

# default training
config.default = edict()
config.default.frequent = 20
config.default.kvstore = 'device'

# network related params
config.network = edict()
config.network.pretrained = './model/pretrained_model/resnet_v2_101'
config.network.pretrained_epoch = 0
config.network.PIXEL_MEANS = np.array([0, 0, 0])
config.network.IMAGE_STRIDE = 0
config.network.RPN_FEAT_STRIDE = 16
config.network.RCNN_FEAT_STRIDE = 16
config.network.FIXED_PARAMS = ['conv0', 'stage1', 'gamma', 'beta']
config.network.FIXED_PARAMS_SHARED = ['conv0', 'stage1', 'stage2', 'stage3', 'gamma', 'beta']
config.network.ANCHOR_SCALES = (8, 16, 32)
config.network.ANCHOR_RATIOS = (0.5, 1, 2)
config.network.NUM_ANCHORS = len(config.network.ANCHOR_SCALES) * len(config.network.ANCHOR_RATIOS)

# dataset related params
config.dataset = edict()
config.dataset.dataset = 'PascalVOC'
config.dataset.image_set = 'SDS_train'
config.dataset.test_image_set = 'SDS_val'
config.dataset.root_path = './data'
config.dataset.dataset_path = './data/VOCdevkit'
config.dataset.NUM_CLASSES = 21

# Training configurations
config.TRAIN = edict()

config.TRAIN.lr = 0.0005
config.TRAIN.lr_step = ''
config.TRAIN.warmup = False
config.TRAIN.warmup_lr = 0
config.TRAIN.warmup_step = 0
config.TRAIN.momentum = 0.9
config.TRAIN.wd = 0.0005
config.TRAIN.begin_epoch = 0
config.TRAIN.end_epoch = 0
config.TRAIN.model_prefix = ''

# whether resume training
config.TRAIN.RESUME = False
# whether flip image
config.TRAIN.FLIP = True
# whether shuffle image
config.TRAIN.SHUFFLE = True
# whether use OHEM
config.TRAIN.ENABLE_OHEM = False
# size of images for each device, 2 for rcnn, 1 for rpn and e2e
config.TRAIN.BATCH_IMAGES = 2
# e2e changes behavior of anchor loader and metric
config.TRAIN.END2END = False
# group images with similar aspect ratio
config.TRAIN.ASPECT_GROUPING = True

# R-CNN
# rcnn rois batch size
config.TRAIN.BATCH_ROIS = 128
config.TRAIN.BATCH_ROIS_OHEM = 128
# rcnn rois sampling params
config.TRAIN.FG_FRACTION = 0.25
config.TRAIN.FG_THRESH = 0.5
config.TRAIN.BG_THRESH_HI = 0.5
config.TRAIN.BG_THRESH_LO = 0.0
# rcnn bounding box regression params
config.TRAIN.BBOX_REGRESSION_THRESH = 0.5
config.TRAIN.BBOX_WEIGHTS = np.array([1.0, 1.0, 1.0, 1.0])

# RPN anchor loader
# rpn anchors batch size
config.TRAIN.RPN_BATCH_SIZE = 256
# rpn anchors sampling params
config.TRAIN.RPN_FG_FRACTION = 0.5
config.TRAIN.RPN_POSITIVE_OVERLAP = 0.7
config.TRAIN.RPN_NEGATIVE_OVERLAP = 0.3
config.TRAIN.RPN_CLOBBER_POSITIVES = False
# rpn bounding box regression params
config.TRAIN.RPN_BBOX_WEIGHTS = (1.0, 1.0, 1.0, 1.0)
config.TRAIN.RPN_POSITIVE_WEIGHT = -1.0

# used for end2end training
# RPN proposal
config.TRAIN.CXX_PROPOSAL = True
config.TRAIN.RPN_NMS_THRESH = 0.7
config.TRAIN.RPN_PRE_NMS_TOP_N = 12000
config.TRAIN.RPN_POST_NMS_TOP_N = 2000
config.TRAIN.RPN_MIN_SIZE = config.network.RPN_FEAT_STRIDE
config.TRAIN.RPN_ALLOWED_BORDER = 0
# whether select from all rois or bg rois
config.TRAIN.GAP_SELECT_FROM_ALL = True
config.TRAIN.IGNORE_GAP = False
# binary_threshold for proposal annotator
config.TRAIN.BINARY_THRESH = 0.4
# approximate bounding box regression
config.TRAIN.BBOX_NORMALIZATION_PRECOMPUTED = False
config.TRAIN.BBOX_MEANS = (0.0, 0.0, 0.0, 0.0)
config.TRAIN.BBOX_STDS = (0.1, 0.1, 0.2, 0.2)
# loss weight for cls, bbox and mask
config.TRAIN.LOSS_WEIGHT = (1.0, 10.0, 1.0)
config.TRAIN.CONVNEW3 = False

config.TEST = edict()

# R-CNN testing
# use rpn to generate proposal
config.TEST.HAS_RPN = False
# size of images for each device
config.TEST.BATCH_IMAGES = 1

# RPN proposal
config.TEST.CXX_PROPOSAL = True
config.TEST.RPN_NMS_THRESH = 0.7
config.TEST.RPN_PRE_NMS_TOP_N = 6000
config.TEST.RPN_POST_NMS_TOP_N = 300
config.TEST.RPN_MIN_SIZE = config.network.RPN_FEAT_STRIDE

# RPN generate proposal
config.TEST.PROPOSAL_NMS_THRESH = 0.7
config.TEST.PROPOSAL_PRE_NMS_TOP_N = 20000
config.TEST.PROPOSAL_POST_NMS_TOP_N = 2000
config.TEST.PROPOSAL_MIN_SIZE = config.network.RPN_FEAT_STRIDE

# RCNN nms
config.TEST.NMS = 0.3

# Test Model Epoch
config.TEST.test_epoch = 0

# TEST iteration
config.TEST.ITER = 1
config.TEST.MIN_DROP_SIZE = 16

# mask merge
config.TEST.USE_MASK_MERGE = False
config.TEST.USE_GPU_MASK_MERGE = False
config.TEST.MASK_MERGE_THRESH = 0.5

def update_config(config_file):
    exp_config = None
    with open(config_file) as f:
        exp_config = edict(yaml.load(f))
        for k, v in exp_config.items():
            if k in config:
                if isinstance(v, dict):
                    if k == 'TRAIN':
                        if 'BBOX_WEIGHTS' in v:
                            v['BBOX_WEIGHTS'] = np.array(v['BBOX_WEIGHTS'])
                    elif k == 'network':
                        if 'PIXEL_MEANS' in v:
                            v['PIXEL_MEANS'] = np.array(v['PIXEL_MEANS'])
                    for vk, vv in v.items():
                        config[k][vk] = vv
                else:
                    if k == 'SCALES':
                        config[k][0] = (tuple(v))
                    else:
                        config[k] = v
            else:
                raise ValueError("key must exist in config.py")


import copy
import glob
import os
from collections import OrderedDict
from datetime import datetime

import netCDF4 as nc
import logging
import numpy as np
from netCDF4 import Dataset
from smrf.utils import utils

from awsm import __version__

C_TO_K = 273.16
FREEZE = C_TO_K

class StateUpdater():
    """
    Class to initialize the updates, perform the updates, and store the needed
    updating info
    """

    def __init__(self, myawsm):
        self.update_fp = myawsm.update_file
        # read in updates and store both dates and images
        update_info, x, y = self.initialize_aso_updates(myawsm, self.update_fp)

        self._logger = logging.getLogger(__name__)
        self.end_date = myawsm.end_date
        self.x = x
        self.y = y
        self.update_info = update_info

        self.tzinfo = myawsm.tzinfo
        # check to see if we're outputting the changes resulting from each update
        self.update_change_file = myawsm.config['update depth']['update_change_file']
        if self.update_change_file is not None:
            start_date = myawsm.config['time']['start_date']
            time_zone = myawsm.config['time']['time_zone']
            self.delta_ds = self.initialize_update_output(
                start_date,
                time_zone,
                __version__,
                myawsm.smrf_version)

        # calculate offset for each section of the run and filter updates
        # update_info, runsteps, offsets, firststeps = self.calc_offsets_nsteps(myawsm, update_info)
        #
        # self.runsteps = runsteps
        # self.offsets = offsets
        # self.firststeps = firststeps

        # save the dates of each update if there are any updates in the time frame
        self.update_dates = []
        if len(self.update_info) > 0:
            self.update_dates = [self.update_info[k]['date_time']
                                 for k in self.update_info.keys()]

        # get necessary variables from awsm class
        self.active_layer = myawsm.active_layer
        # Buffer size (in cells) for the interpolation to search overself.
        self.update_buffer = myawsm.update_buffer

        self.ny = myawsm.topo.ny
        self.nx = myawsm.topo.nx
        self.topo = myawsm.topo

    def do_update_pysnobal(self, output_rec, dt):
        """
        Function to update a time step of a pysnobal run by updating the
        output_rec

        Args:
            output_rec:     iPySnobal state variables
            dt:             iPySnobal datetime of timestep
        """
        self._logger.debug('Preparing to update pysnobal')
        # find the correct update number
        ks = self.update_info.keys()
        update_num = [k for k in ks if self.update_info[k]['date_time'] == dt]
        if len(update_num) > 1:
            raise ValueError(
                'Something wrong in pysnobal updating date compare')

        un = update_num[0]

        # get parameters from PySnobal
        m_s = output_rec['m_s']
        T_s_0 = output_rec['T_s_0'] - FREEZE
        T_s_l = output_rec['T_s_l'] - FREEZE
        T_s = output_rec['T_s'] - FREEZE
        h2o_sat = output_rec['h2o_sat']
        z_s = output_rec['z_s']
        density = output_rec['rho']

        # do the updating
        updated_fields = self.hedrick_updating_procedure(m_s, T_s_0, T_s_l, T_s,
                                                         h2o_sat, density, z_s,
                                                         self.x, self.y,
                                                         self.update_info[un])

        # calculate the change from the update
        no_update = np.isnan(self.update_info[un]['depth'])
        # difference in depth
        diff_z = updated_fields['D'] - output_rec['z_s']
        diff_z[no_update] = np.nan
        # difference in density
        diff_rho = updated_fields['rho'] - output_rec['rho']
        diff_rho[no_update] = np.nan
        # difference in SWE
        diff_swe = updated_fields['D'] * \
            updated_fields['rho'] - output_rec['m_s']
        diff_swe[no_update] = np.nan

        # check to see if last update for the time period

        islast = False
        update_numbers = list(self.update_info.keys())
        if update_numbers[-1] == un:
            islast = True
        else:
            next_un_date = self.update_info[un+1]['date_time']
            if next_un_date > self.end_date.replace(tzinfo=self.tzinfo):
                islast = True

        # write a file to show the change in updates
        if self.update_change_file is not None:
            self.output_update_changes(diff_z, diff_rho, diff_swe, dt, islast)

        # save the fields
        output_rec['m_s'] = updated_fields['D'] * updated_fields['rho']
        output_rec['T_s_0'] = updated_fields['T_s_0'] + FREEZE
        output_rec['T_s_l'] = updated_fields['T_s_l'] + FREEZE
        output_rec['T_s'] = updated_fields['T_s'] + FREEZE
        output_rec['h2o_sat'] = updated_fields['h2o_sat']
        output_rec['z_s'] = updated_fields['D']
        output_rec['rho'] = updated_fields['rho']

        return output_rec

    def initialize_aso_updates(self, myawsm, update_fp):
        """
        Read in the ASO update file and parse images by date
        Argument:
                myawsm: instantiated awsm class
                update_fp: file pointer to netCDF with all flights in it
        Return:
                update_info: dictionary of updates
        """

        # Update the snow depths in the initialization file using ASO lidar:
        ds = Dataset(update_fp, 'r')
        ds.set_always_mask(False)

        # get all depths, x, y, time
        D_all = ds.variables['depth'][:]
        D_all = D_all.filled(fill_value=np.nan)
        D_all[np.isinf(D_all)] = np.nan
        D_all[D_all > 200.0] = np.nan
        if np.any(D_all) > 100:
            print('Check D_all')
            D_all[D_all > 100.0] = np.nan

        x = ds.variables['x'][:]
        y = ds.variables['y'][:]
        times = ds.variables['time']
        ts = times[:]
        # convert time index to dates
        t = nc.num2date(ts, times.units, times.calendar,
                        only_use_cftime_datetimes=False)

        # find wyhr of dates
        t_wyhr = []
        for t1 in t:
            tmp_date = t1.replace(tzinfo=myawsm.tzinfo)
            # get wyhr
            tmpwyhr = int(utils.water_day(tmp_date)[0]*24)
            t_wyhr.append(tmpwyhr)

        t_wyhr = np.array(t_wyhr)

        # make dictionary of updates
        update_info = OrderedDict()
        keys = range(1, len(t_wyhr)+1)
        for idk, k in enumerate(keys):
            # make dictionary for each update
            update_info[k] = {}
            # set update number
            update_info[k]['number'] = k
            update_info[k]['date_time'] = t[idk].replace(tzinfo=myawsm.tzinfo)
            update_info[k]['wyhr'] = t_wyhr[idk]
            # set depth
            update_info[k]['depth'] = D_all[idk, :]

        return update_info, x, y

    def calc_offsets_nsteps(self, myawsm, update_info):
        """
        Function to calculate the offset for each update run and the number of steps
        the iSnobal run needs to run
        Args:
            myawsm: awsm class
            update_info: pandas dataframe of the update netCDF info
        Returns:
            update_info: update info dataframe filtered to desired updates
            runsteps:    numpy array of runsteps for each section
            offsets:     offset wyhr for each update
            firststeps:  number of steps to run before first update, if any
        """

        #t_wyhr = update_info['wyhr'].values
        update_number = update_info.keys()
        # filter to desired flights if user input
        if myawsm.flight_numbers is not None:
            filter_number = myawsm.flight_numbers
        else:
            filter_number = update_number
            # # make list if not list
            # if not isinstance(filter_number, list):
            #     filter_number = [filter_number]
            #     update_number = [update_number]

        myawsm._logger.debug(
            'Will update with flights {}'.format(filter_number))
        # filter to correct hours
        for un in update_number:
            if un not in filter_number:
                # delete update if not in desired update inputs
                update_info.pop(un)

        # check if a first run with no update is needed to get us up to the first update
        if myawsm.restart_crash:
            test_start_wyhr = myawsm.restart_hr+1
        else:
            test_start_wyhr = myawsm.start_wyhr

        # filter so we are in the dates
        update_info_copy = copy.deepcopy(update_info)
        for un in update_info_copy.keys():
            tw = update_info[un]['wyhr']
            # get rid of updates more than a day before start date
            if tw < test_start_wyhr - 24:
                update_info.pop(un)
            # get rid of updates past the end of the run
            elif tw > myawsm.end_wyhr:
                update_info.pop(un)

        # now we are down to the correct wyhrs and update numbers
        update_number = update_info.keys()
        t_wyhr = [update_info[k]['wyhr'] for k in update_number]

        # this is where each run will start from
        offsets = t_wyhr

        if len(offsets) == 0:
            raise IOError('No update dates in this run')

        # do we need to do that run?
        if offsets[0] - test_start_wyhr <= 0:
            firststeps = 0
        else:
            firststeps = offsets[0] - test_start_wyhr

        runsteps = np.zeros_like(offsets)
        for ido, offs in enumerate(offsets):
            if ido == len(offsets) - 1:
                runsteps[ido] = myawsm.end_wyhr - offs
            else:
                runsteps[ido] = offsets[ido+1] - offs

        myawsm._logger.debug('Filtered to updates on: {}'.format(t_wyhr))

        return update_info, runsteps, offsets, firststeps

    def find_update_snow(self, myawsm, offset):
        """
        Function to find the nearest, lower ouptut file to perform the update_info
        """
        # get output files so far
        d = sorted(glob.glob("%s/snow*" % myawsm.pathro),
                   key=os.path.getmtime)

        d.sort(key=lambda f: os.path.splitext(f))

        hr = []
        for idf, f in enumerate(d):
            # get the hr
            nm = os.path.basename(f)
            head = os.path.dirname(f)
            hr.append(int(nm.split('.')[1]))

        hr = np.array(hr)
        # filter to outputs less than offset
        hr = hr[hr < offset]
        #hr = [h for h in hr if h < offset]
        # find closest
        idx = (np.abs(hr - offset)).argmin()
        offset_hr = int(hr[idx])

        # make sure we are within a day of the update
        if np.abs(offset_hr - offset) > 24:
            raise ValueError('No output in output directory within'
                             ' a day of {} update'.format(offset))

        # set the file to update
        update_snow = os.path.join(myawsm.pathro, 'snow.%04d' % (offset_hr))
        if not os.path.exists(update_snow):
            raise ValueError(
                'Update snow file {} does not exist'.format(update_snow))

        return update_snow

    def hedrick_updating_procedure(self, m_s, T_s_0, T_s_l, T_s, h2o_sat, density, z_s,
                                   x, y, update_info):
        """
        This function performs the direct insertion procedure and returns the
        updated fields.

        Argument:
            m_s:    swe array to be updated
            T_s_0:  surface layer temperature array to be updated
            T_s_l:  lower layer temperature array to be updated
            T_s:    bulk temperature array to be updated
            h2o_sat: h2o_sat array to be updated
            density: density array to be updated
            z_s:     snow height array to be updated
            x:       x vector
            y:       y vector
            logger:  awsm logger
            update_info: necessary update info
        Returns:
            updated_fields:  dictionary of updated fields including dem, z0, D,
                             rho, T_s_0, T_s_l, T_s, h2o_sat
        """

        # keep a copy of the original inputs
        original_fields = {}
        original_fields['m_s'] = m_s.copy()
        original_fields['T_s_0'] = T_s_0.copy()
        original_fields['T_s_l'] = T_s_l.copy()
        original_fields['T_s'] = T_s.copy()
        original_fields['h2o_sat'] = h2o_sat.copy()
        original_fields['density'] = density.copy()
        original_fields['z_s'] = z_s.copy()

        activeLayer = self.active_layer
        # Buffer size (in cells) for the interpolation to search over.
        Buf = self.update_buffer
        # get dem and roughness
        dem = self.topo.dem
        z0 = self.topo.roughness

        # New depth field
        D = update_info['depth']

        # make mask
        # D[self.topo.mask == 0.0] = np.nan
        mask = np.ones_like(D)
        mask[np.isnan(D)] = 0.0

        XX, YY = np.meshgrid(x, y)

        nrows = len(y)
        ncols = len(x)

        # Special case - 20160607
        # I am trying an update with only Tuolumne Basin data where I will mask in
        # Cherry and Eleanor to create a hybrid iSnobal/ASO depth image.
        tempASO = D.copy()
        tempASO[np.isnan(D)] = 0.0
        tempiSnobal = z_s.copy()
        tuolx_mask = mask
        tempASO[tuolx_mask == 1.0] = 0.0
        #tempiSnobal[tuolx_mask == 1] = 0.0
        I_ASO = (tempASO == 0.0)
        tempASO[I_ASO] = tempiSnobal[I_ASO]
        tempASO[tuolx_mask == 1.0] = D[tuolx_mask == 1.0]
        D = tempASO.copy()

        # Address problem of bit resolution where cells have mass and density,
        # but no depth is reported (less than minimum depth above zero).
        u_depth = np.unique(z_s)

        id_m_s = m_s == 0.0
        z_s[id_m_s] = 0.0
        density[id_m_s] = 0.0
        T_s[id_m_s] = -75.0
        T_s_l[id_m_s] = -75.0
        T_s_0[id_m_s] = -75.0
        h2o_sat[id_m_s] = 0.0

        z_s[(density > 0.0) & (z_s == 0.0)] = u_depth[1]

        rho = density.copy()
        D[D < 0.05] = 0.0  # Set shallow snow (less than 5cm) to 0.
        D[mask == 0.0] = np.nan  # Set out of watershed cells to NaN
        rho[mask == 0.0] = np.nan  # Set out of watershed cells to NaN
        tot_pix = ncols * nrows  # Get number of pixels in domain.

        I_model = np.where(z_s == 0)  # Snow-free pixels in the model.
        modelDepth = tot_pix - len(I_model[0])  # of pixels with snow (model).
        # Snow-free pixels from lidar.
        I_lidar = np.where((D == 0) | (np.isnan(D)))
        lidarDepth = tot_pix - len(I_lidar[0])  # of pixels with snow (lidar).
        I_rho = np.where(density == 0)  # Snow-free pixels upon importing.
        # of pixels with density (model).
        modelDensity = tot_pix - len(I_rho[0])

        self._logger.debug('\nJust After Importing.\n \
                              Number of modeled cells with snow depth: {}\n \
                              Number of modeled cells with density: {}\n \
                              Number of lidar cells measuring snow: {}'.format(modelDepth, modelDensity, lidarDepth))
        id0 = D == 0.0

        # Find cells without lidar snow and set the modeled density to zero.
        rho[id0] = 0.0
        rho[rho == 0.0] = np.nan  # Set all cells with no density to NaN.

        # Find cells without lidar snow and set the active layer temp to NaN.
        T_s_0[id0] = np.nan
        T_s_0[T_s_0 <= -75.0] = np.nan  # Change isnobal no-values to NaN.

        # Find cells without lidar snow and set the lower layer temp to NaN.
        T_s_l[id0] = np.nan
        T_s_l[T_s_l <= -75.0] = np.nan  # Change isnobal no-values to NaN.

        # Find cells without lidar snow and set the snow temp to np.nan.
        T_s[id0] = np.nan
        T_s[T_s <= -75.0] = np.nan  # Change isnobal no-values to NaN.

        # Find cells without lidar snow and set the h2o saturation to NaN.
        h2o_sat[id0] = np.nan
        h2o_sat[mask == 0.0] = np.nan
        # h2o_sat[h2o_sat == -75.0] = np.nan # Change isnobal no-values to NaN.

        # Snow-free pixels before interpolation
        I_rho = np.where(np.isnan(rho))
        #modelDensity = tot_pix - size(I_rho, 1)
        modelDensity = tot_pix - len(I_rho[0])

        self._logger.debug('\nBefore Interpolation.\n \
                              Number of modeled cells with snow depth: {0}\n \
                              Number of modeled cells with density: {1}\n \
                              Number of lidar cells measuring snow: {2}'.format(modelDepth, modelDensity, lidarDepth))

        # Now find cells where lidar measured snow, but Isnobal simulated no snow:
        I = np.where((np.isnan(rho)) & (D > 0.0))
        I_25 = np.where((z_s <= (activeLayer * 1.20)) &
                        (D >= activeLayer))  # find cells with lidar
        # depth greater than, and iSnobal depths less than, the active layer
        # depth. Lower layer temperatures of these cells will need to be
        # interpolated from surrounding cells with lower layer temperatures.
        # This happens AFTER the interpolation of all the other variables
        # below.  If cell has snow depth 120% of set active layer depth, then
        # the lower layer temp will be replaced by an areal interpolated value
        # from surrounding cells with lower layer temps and depths greater than
        # 120% of active layer.

        # Interpolate over these cells to come up with values for them.
        X, Y = np.meshgrid(range(ncols), range(nrows))

        # make matrices to use in following commands
        tmp1 = np.ones((nrows+2*Buf, Buf))
        tmp1[:] = np.nan
        tmp2 = np.ones((Buf, ncols))
        tmp2[:] = np.nan
        # create buffer
        rho_buf = np.concatenate((tmp1, np.concatenate(
            (tmp2, rho, tmp2), axis=0), tmp1), axis=1)
        T_s_0_buf = np.concatenate((tmp1, np.concatenate(
            (tmp2, T_s_0, tmp2), axis=0), tmp1), axis=1)
        T_s_l_buf = np.concatenate((tmp1, np.concatenate(
            (tmp2, T_s_l, tmp2), axis=0), tmp1), axis=1)
        T_s_buf = np.concatenate((tmp1, np.concatenate(
            (tmp2, T_s, tmp2), axis=0), tmp1), axis=1)
        h2o_buf = np.concatenate((tmp1, np.concatenate(
            (tmp2, h2o_sat, tmp2), axis=0), tmp1), axis=1)

        # hopefully fixed for loop logic below

        # Loop through cells with D > 0 and no iSnobal density,
        for idx, (ix, iy) in enumerate(zip(I[0], I[1])):
            # active layer temp, snow temp, and h2o saturation.
            xt = X[ix, iy]+Buf  # Add the buffer to the x coords.
            yt = Y[ix, iy]+Buf  # Add the buffer to the y coords.
            n = range(11, Buf+2, 10)  # Number of cells in averaging window
            for n1 in n:  # Loop through changing buffer windows until enough
                # cells are found to calculate an average.
                xl = xt - int((n1 - 1) / 2)
                xh = xt + int((n1 - 1) / 2)
                yl = yt - int((n1 - 1) / 2)
                yh = yt + int((n1 - 1) / 2)
                window = rho_buf[yl:yh, xl:xh]
                # find number of pixels with a value.
                qq = np.where(np.isfinite(window))
                if (len(qq[0]) > 10):
                    val = np.nanmean(window[:])
                    # Interpolate for density (just a windowed mean)
                    rho[ix, iy] = val
                    window = T_s_0_buf[yl:yh, xl:xh]
                    val = np.nanmean(window[:])
                    # Interpolate for active snow layer temp
                    T_s_0[ix, iy] = val
                    # Handle the lower layer temp in the following for-loop.
                    window = T_s_buf[yl:yh, xl:xh]
                    val = np.nanmean(window[:])
                    T_s[ix, iy] = val  # Interpolate for avg snow temp
                    window = h2o_buf[yl:yh, xl:xh]
                    val = np.nanmean(window[:])
                    # Interpolate for liquid water saturation
                    h2o_sat[ix, iy] = val

                # check to see if you found a value
                if n1 == n[-1] and np.isnan(rho[ix, iy]):
                    self._logger.error('Failed to find desnity wihtin buffer')
                # if we found a value, move on
                if np.isfinite(rho[ix, iy]):
                    break

        # ##################### hopefully fixed for loop logic below
        self._logger.debug('Done with loop 1')
        # Now loop over cells with D > activelayer > z_s.  These cells were being
        # assigned no temperature in their lower layer (-75) when they needed to
        # have a real temperature.  Solution is to interpolate from nearby cells
        # using an expanding moving window search.
        for ix, iy in zip(I_25[0], I_25[1]):
            xt = X[ix, iy] + Buf  # Add the buffer to the x coords.
            yt = Y[ix, iy] + Buf  # Add the buffer to the y coords.
            n = range(11, Buf+2, 10)
            for jj in n:  # Loop through changing buffer windows until enough
                # cells are found to calculate an average.
                xl = xt - int((jj-1)/2)
                xh = xt + int((jj-1)/2)
                yl = yt - int((jj-1)/2)
                yh = yt + int((jj-1)/2)
                window = T_s_l_buf[yl:yh, xl:xh]
                val = np.nanmean(window[:])
                T_s_l[ix, iy] = val  # Interpolate for lower layer temp
                # fix this to be pyton logic
                # if np.isnan(T_s_l[ii]) == False:
                if not np.any(np.isnan(T_s_l[ix, iy])):
                    break

        self._logger.debug('Done with loop 2')

        iq = (np.isnan(D)) & (np.isfinite(rho))
        # Once more, change cells with no lidar snow to have np.nan density.
        rho[iq] = np.nan

        # Find occurance where cell has depth and density but no temperature.
        # Delete snowpack from this cell.
        iq2 = (np.isnan(T_s)) & (np.isfinite(rho))
        D[iq2] = 0.0
        rho[iq2] = np.nan

        # Snow-free pixels from lidar.
        I_lidar = np.where((D == 0.0) | (np.isnan(D)))
        lidarDepth = tot_pix - len(I_lidar[0])  # of pixels with snow (lidar).
        I_rho = np.where(np.isnan(rho))  # Snow-free pixels upon importing.
        # of pixels with density (model).
        modelDensity = tot_pix - len(I_rho[0])

        I_lidaridx = (D == 0.0) | (np.isnan(D))  # Snow-free pixels from lidar.
        I_rhoidx = np.isnan(rho)  # Snow-free pixels upon importing.

        self._logger.debug('\nAfter Interpolation.\n \
                              Number of modeled cells with snow depth: {}\n \
                              Number of modeled cells with density: {}\n \
                              Number of lidar cells measuring snow: {}'.format(modelDepth, modelDensity, lidarDepth))

        # Reset NaN's to the proper values for Isnobal:
        # if size(I_lidar, 1) ~= size(I_rho, 1)
        if len(I_lidar[0]) != len(I_rho[0]):
            raise ValueError(
                'Lidar depths do not match interpolated model densities.  Try changing buffer parameters.')

        rho[I_rhoidx] = 0.0  # rho is the updated density map.
        # D is lidar snow depths, I_rho is where no snow exists.
        D[rho == 0.0] = 0.0
        I_25_new = D <= activeLayer  # find cells with lidar depth less than 25 cm
        # These cells will have the corresponding lower layer temp changed to
        # -75 (no value) and the upper layer temp will be set to equal the
        # average snowpack temp in that cell.
        # T_s is the average snow temperature in a cell. Is NaN (-75) for all rho = 0.
        T_s[rho == 0.0] = -75.0
        # T_s_0 is the updated active (upper) layer.  Is >0 for everywhere rho is >0.
        T_s_0[rho == 0.0] = -75.0
        # If lidar depth <= 25cm, set active layer temp to average temp of cell
        T_s_0[I_25_new] = T_s[I_25_new]
        T_s_l[I_25_new] = -75.0
        T_s_l[np.isnan(T_s_l)] = -75.0
        h2o_sat[rho == 0.0] = 0.0

        # grab unmasked cells again
        nmask = mask == 0
        # Make sure non-updated cells stay the same as original
        m_s[nmask] = original_fields['m_s'][nmask]  # Get SWE image.
        # Get active snow layer temperature image
        T_s_0[nmask] = original_fields['T_s_0'][nmask]
        # Get lower snow layer temperature image
        T_s_l[nmask] = original_fields['T_s_l'][nmask]
        # Get average snowpack temperature image
        T_s[nmask] = original_fields['T_s'][nmask]
        # Get liquid water saturation image
        h2o_sat[nmask] = original_fields['h2o_sat'][nmask]
        D[nmask] = original_fields['z_s'][nmask]
        rho[nmask] = original_fields['density'][nmask]

        # create dictionary to return with updated arrays
        updated_fields = {}
        updated_fields['dem'] = dem
        updated_fields['z0'] = z0
        updated_fields['D'] = D
        updated_fields['rho'] = rho
        updated_fields['T_s_0'] = T_s_0
        updated_fields['T_s_l'] = T_s_l
        updated_fields['T_s'] = T_s
        updated_fields['h2o_sat'] = h2o_sat

        return updated_fields

    def output_update_changes(self, diff_z, diff_rho, diff_swe, dt, islast):
        """
        Output the effect of the depth update each time that the state is
        updated.

        Args:
            diff_z: numpy array of change in depth
            diff_rho: numpy array of change in rho
            diff_swe: numpy array of change in swe
            dt: PySnobal time step
            islast: boolean describing if it is the last update to process
        """
        variable_list = ['depth_change', 'rho_change', 'swe_change']
        # now find the correct index
        # the current time integer
        times = self.delta_ds.variables['time']
        # convert to index
        t = nc.date2num(dt.replace(tzinfo=None), times.units, times.calendar)
        # find index in context of file
        if len(times) != 0:
            index = np.where(times[:] == t)[0]
            if index.size == 0:
                index = len(times)
            else:
                index = index[0]
        else:
            index = len(times)

        # insert the time and data
        self.delta_ds.variables['time'][index] = t
        self.delta_ds.variables['depth_change'][index, :] = diff_z
        self.delta_ds.variables['rho_change'][index, :] = diff_rho
        self.delta_ds.variables['swe_change'][index, :] = diff_swe

        self.delta_ds.sync()

        # close file if we're done
        if islast:
            self.delta_ds.close()

    def initialize_update_output(self, start_date, time_zone, awsm_version,
                                 smrf_version):
        """
        Initialize the output files to track the changes in the state Variables
        resulting from an update in snow depth.

        Args:
            start_date: start_date from config
            time_zone: time zone from config
            awsm_version: version of awsm being used
            smrf_version: version of smrf being used

        """
        fmt = '%Y-%m-%d %H:%M:%S'
        # chunk size
        cs = (6, 10, 10)

        variable_dict = {
            'depth_change': {
                'units': 'meters',
                'description': 'change in depth from update'
            },
            'rho_change': {
                'units': 'kg*m^-2',
                'description': 'change in density from update'
            },
            'swe_change': {
                'units': 'mm',
                'description': 'change in SWE from update'
            }
        }
        # determine the x,y vectors for the netCDF file
        x = self.x
        y = self.y
        # self.mask = topo.mask
#         dimensions = ('Time','dateStrLen','y','x')
        dimensions = ('time', 'y', 'x')
        self.date_time = {}

        if os.path.isfile(self.update_change_file):
            self._logger.warning(
                'Opening {}, data may be overwritten!'.format(self.update_change_file))
            ds = nc.Dataset(self.update_change_file, 'a')
            h = '[{}] Data added or updated'.format(
                datetime.now().strftime(fmt))
            setattr(ds, 'last_modified', h)

        else:
            ds = nc.Dataset(self.update_change_file, 'w')

            dimensions = ('time', 'y', 'x')

            # create the dimensions
            ds.createDimension('time', None)
            ds.createDimension('y', len(y))
            ds.createDimension('x', len(x))

            # create some variables
            ds.createVariable('time', 'f', dimensions[0])
            ds.createVariable('y', 'f', dimensions[1])
            ds.createVariable('x', 'f', dimensions[2])

            # setattr(em.variables['time'], 'units', 'hours since %s' % options['time']['start_date'])
            setattr(ds.variables['time'], 'units',
                    'hours since %s' % start_date)
            setattr(ds.variables['time'], 'calendar', 'standard')
            #     setattr(em.variables['time'], 'time_zone', time_zone)
            ds.variables['y'][:] = y
            ds.variables['x'][:] = x

            # em image
            for v, f in variable_dict.items():
                ds.createVariable(v, 'f', dimensions[:3], chunksizes=cs)
                setattr(ds.variables[v], 'units', f['units'])
                setattr(ds.variables[v], 'description', f['description'])

            # define some global attributes
            ds.setncattr_string('Conventions', 'CF-1.6')
            ds.setncattr_string('dateCreated', datetime.now().strftime(fmt))
            ds.setncattr_string('created_with', 'Created with awsm {} and smrf {}'.format(
                awsm_version, smrf_version))
            ds.setncattr_string('history', '[{}] Create netCDF4 file'.format(
                datetime.now().strftime(fmt)))
            ds.setncattr_string('institution',
                                'USDA Agricultural Research Service, Northwest Watershed Research Center')

        # save the open dataset so we can write to it
        ds.sync()

        return ds