import numpy as np
import fitsio
from astrometry.util.fits import fits_table
from astrometry.util.resample import resample_with_wcs, OverlapError
from legacypipe.bits import DQ_BITS
from legacypipe.survey import tim_get_resamp
from legacypipe.utils import copy_header_with_wcs

import logging
logger = logging.getLogger('legacypipe.coadds')
def info(*args):
    from legacypipe.utils import log_info
    log_info(logger, args)
def debug(*args):
    from legacypipe.utils import log_debug
    log_debug(logger, args)

class SimpleCoadd(object):
    '''A class for handling coadds of unWISE (and GALEX) images.
    '''
    def __init__(self, ra, dec, W, H, pixscale, bands):
        from legacypipe.survey import wcs_for_brick, BrickDuck
        self.W = W
        self.H = H
        self.bands = bands
        brick = BrickDuck(ra, dec, 'quack')
        self.wcs = wcs_for_brick(brick, W=self.W, H=self.H,
                                        pixscale=pixscale)
        # images
        self.co_images = dict([(band, np.zeros((self.H,self.W), np.float32))
                               for band in bands])
        self.co_nobs = dict([(band, np.zeros((self.H,self.W), np.uint16))
                             for band in bands])
        # models
        self.co_models = dict([(band, np.zeros((self.H,self.W), np.float32))
                               for band in bands])
        # invvars
        self.co_invvars = dict([(band, np.zeros((self.H,self.W), np.float32))
                                for band in bands])

    def add(self, models, unique=False):
        for name, band, wcs, img, mod, ie in models:
            debug('Accumulating tile', name, 'band', band)
            try:
                Yo,Xo,Yi,Xi,resam = resample_with_wcs(self.wcs, wcs,
                                                      [img, mod], intType=np.int16)
            except OverlapError:
                debug('No overlap between tile', name, 'and coadd')
                continue
            rimg,rmod = resam
            debug('Adding', len(Yo), 'pixels from tile', name, 'to coadd')
            iv = ie[Yi,Xi]**2
            if unique:
                K = np.flatnonzero((self.co_nobs[band][Yo,Xo] == 0) * (iv>0))
                iv = iv[K]
                rimg = rimg[K]
                rmod = rmod[K]
                Yo = Yo[K]
                Xo = Xo[K]
                debug('Cut to', len(Yo), 'unique pixels w/ iv>0')

            debug('Tile:', np.sum(iv>0), 'of', len(iv), 'pixels have IV')
            self.co_images [band][Yo,Xo] += rimg * iv
            self.co_models [band][Yo,Xo] += rmod * iv
            self.co_nobs   [band][Yo,Xo] += 1
            self.co_invvars[band][Yo,Xo] += iv
            debug('Band', band, ': now', np.sum(self.co_nobs[band]>0), 'pixels are set in image coadd')

    def finish(self, survey, brickname, version_header,
               apradec=None, apertures=None):
        # apradec = (ra,dec): aperture photometry locations
        # apertures: RADII in PIXELS
        if apradec is not None:
            assert(apertures is not None)
            (ra,dec) = apradec
            ok,xx,yy = self.wcs.radec2pixelxy(ra, dec)
            assert(np.all(ok))
            del ok
            apxy = np.vstack((xx - 1., yy - 1.)).T
            ap_iphots = [np.zeros((len(ra), len(apertures)), np.float32)
                         for band in self.bands]
            ap_dphots = [np.zeros((len(ra), len(apertures)), np.float32)
                         for band in self.bands]
            ap_rphots = [np.zeros((len(ra), len(apertures)), np.float32)
                         for band in self.bands]

        coimgs = []
        comods = []
        for iband,band in enumerate(self.bands):
            coimg = self.co_images[band]
            comod = self.co_models[band]
            coiv  = self.co_invvars[band]
            con   = self.co_nobs[band]
            with np.errstate(divide='ignore', invalid='ignore'):
                coimg /= coiv
                comod /= coiv
            coimg[coiv == 0] = 0.
            comod[coiv == 0] = 0.
            coimgs.append(coimg)
            comods.append(comod)

            hdr = copy_header_with_wcs(version_header, self.wcs)
            self.add_to_header(hdr, band)
            self.write_coadds(survey, brickname, hdr, band, coimg, comod, coiv, con)

            if apradec is not None:
                import photutils
                mask = (coiv == 0)
                with np.errstate(divide='ignore'):
                    imsigma = 1.0/np.sqrt(coiv)
                imsigma[mask] = 0.
                for irad,rad in enumerate(apertures):
                    aper = photutils.CircularAperture(apxy, rad)
                    p = photutils.aperture_photometry(coimg, aper, error=imsigma,
                                                      mask=mask)
                    ap_iphots[iband][:,irad] = p.field('aperture_sum')
                    ap_dphots[iband][:,irad] = p.field('aperture_sum_err')
                    p = photutils.aperture_photometry(coimg - comod, aper, mask=mask)
                    ap_rphots[iband][:,irad] = p.field('aperture_sum')

        self.write_color_image(survey, brickname, coimgs, comods)

        if apradec is not None:
            return ap_iphots, ap_dphots, ap_rphots

    def add_to_header(self, hdr, band):
        pass

    def write_coadds(self, survey, brickname, hdr, band, coimg, comod, coiv, con):
        pass

    def write_color_image(self, survey, brickname, coimgs, comods):
        pass

class UnwiseCoadd(SimpleCoadd):
    def __init__(self, ra, dec, W, H, pixscale):
        super().__init__(ra, dec, W, H, pixscale, [1,2,3,4])

    def add_to_header(self, hdr, band):
        hdr.add_record(dict(name='TELESCOP', value='WISE'))
        hdr.add_record(dict(name='FILTER', value='W%i' % band, comment='WISE band'))
        hdr.add_record(dict(name='MAGZERO', value=22.5,
                            comment='Magnitude zeropoint'))
        hdr.add_record(dict(name='MAGSYS', value='Vega',
                            comment='This WISE image is in Vega fluxes'))

    def write_coadds(self, survey, brickname, hdr, band, coimg, comod, coiv, con):
        with survey.write_output('image', brick=brickname, band='W%i'%band,
                                 shape=coimg.shape) as out:
            out.fits.write(coimg, header=hdr)
        with survey.write_output('model', brick=brickname, band='W%i'%band,
                                 shape=comod.shape) as out:
            out.fits.write(comod, header=hdr)
        with survey.write_output('invvar', brick=brickname, band='W%i'%band,
                                 shape=coiv.shape) as out:
            out.fits.write(coiv, header=hdr)

    def write_color_image(self, survey, brickname, coimgs, comods):
        from legacypipe.survey import imsave_jpeg
        # W1/W2 color jpeg
        rgb = _unwise_to_rgb(coimgs[:2])
        with survey.write_output('wise-jpeg', brick=brickname) as out:
            imsave_jpeg(out.fn, rgb, origin='lower')
            info('Wrote', out.fn)
        rgb = _unwise_to_rgb(comods[:2])
        with survey.write_output('wisemodel-jpeg', brick=brickname) as out:
            imsave_jpeg(out.fn, rgb, origin='lower')
            info('Wrote', out.fn)
        coresids = [coimg - comod for coimg, comod in zip(coimgs[:2], comods[:2])]
        rgb = _unwise_to_rgb(coresids)
        with survey.write_output('wiseresid-jpeg', brick=brickname) as out:
            imsave_jpeg(out.fn, rgb, origin='lower')
            info('Wrote', out.fn)

def _unwise_to_rgb(imgs):
    img = imgs[0]
    H,W = img.shape
    ## FIXME
    w1,w2 = imgs
    rgb = np.zeros((H, W, 3), np.uint8)
    scale1 = 50.
    scale2 = 50.
    mn,mx = -1.,100.
    arcsinh = 1.
    img1 = w1 / scale1
    img2 = w2 / scale2
    if arcsinh is not None:
        def nlmap(x):
            return np.arcsinh(x * arcsinh) / np.sqrt(arcsinh)
        mean = (img1 + img2) / 2.
        I = nlmap(mean)
        with np.errstate(divide='ignore'):
            img1 = I * img1 / mean
            img2 = I * img2 / mean
        img1[mean == 0] = 0.
        img2[mean == 0] = 0.
        mn = nlmap(mn)
        mx = nlmap(mx)
    img1 = (img1 - mn) / (mx - mn)
    img2 = (img2 - mn) / (mx - mn)
    rgb[:,:,2] = (np.clip(img1, 0., 1.) * 255).astype(np.uint8)
    rgb[:,:,0] = (np.clip(img2, 0., 1.) * 255).astype(np.uint8)
    rgb[:,:,1] = rgb[:,:,0]/2 + rgb[:,:,2]/2
    return rgb

def make_coadds(tims, bands, targetwcs,
                coweights=True,
                mods=None, blobmods=None,
                xy=None, apertures=None, apxy=None,
                ngood=False, detmaps=False, psfsize=False,
                allmasks=True, anymasks=False,
                get_max=False, sbscale=True,
                psf_images=False, nsatur=None,
                callback=None, callback_args=None,
                plots=False, ps=None,
                lanczos=True, mp=None,
                satur_val=10.):
    from astrometry.util.ttime import Time
    t0 = Time()

    if callback_args is None:
        callback_args = []

    class Duck(object):
        pass
    C = Duck()

    W = int(targetwcs.get_width())
    H = int(targetwcs.get_height())

    # always, for patching SATUR, etc pixels?
    unweighted=True

    C.coimgs = []
    if coweights:
        # the pixelwise inverse-variances (weights) of the "coimgs".
        C.cowimgs = []
    if detmaps:
        C.galdetivs = []
        C.psfdetivs = []
    if mods is not None:
        C.comods = []
        C.coresids = []
    if blobmods is not None:
        C.coblobmods = []
        C.coblobresids = []
    if apertures is not None:
        C.AP = fits_table()
    if allmasks:
        C.allmasks = []
    if anymasks:
        C.anymasks = []
    if max:
        C.maximgs = []
    if psf_images:
        C.psf_imgs = []
    if nsatur:
        C.satmaps = []

    if xy:
        ix,iy = xy
        C.T = fits_table()
        if ngood:
            C.T.ngood   = np.zeros((len(ix), len(bands)), np.int16)
        C.T.nobs    = np.zeros((len(ix), len(bands)), np.int16)
        C.T.anymask = np.zeros((len(ix), len(bands)), np.int16)
        C.T.allmask = np.zeros((len(ix), len(bands)), np.int16)
        if psfsize:
            C.T.psfsize = np.zeros((len(ix), len(bands)), np.float32)
        if detmaps:
            C.T.psfdepth = np.zeros((len(ix), len(bands)), np.float32)
            C.T.galdepth = np.zeros((len(ix), len(bands)), np.float32)

    if lanczos:
        debug('Doing Lanczos resampling')

    for tim in tims:
        # surface-brightness correction
        tim.sbscale = (targetwcs.pixel_scale() / tim.subwcs.pixel_scale())**2

    # We create one iterator per band to do the tim resampling.  These all run in
    # parallel when multi-processing.
    imaps = []
    for band in bands:
        args = []
        for itim,tim in enumerate(tims):
            if tim.band != band:
                continue
            if mods is None:
                mo = None
            else:
                mo = mods[itim]
            if blobmods is None:
                bmo = None
            else:
                bmo = blobmods[itim]
            args.append((itim,tim,mo,bmo,lanczos,targetwcs,sbscale))
        if mp is not None:
            imaps.append(mp.imap_unordered(_resample_one, args))
        else:
            imaps.append(map(_resample_one, args))

    # Args for aperture photometry
    apargs = []

    if xy:
        # To save the memory of 2 x float64 maps, we instead do arg min/max maps

        # append a 0 to the list of mjds so that mjds[-1] gives 0.
        mjds = np.array([tim.time.toMjd() for tim in tims] + [0])
        mjd_argmins = np.empty((H,W), np.int16)
        mjd_argmaxs = np.empty((H,W), np.int16)
        mjd_argmins[:,:] = -1
        mjd_argmaxs[:,:] = -1

    if plots:
        allresids = []

    tinyw = 1e-30
    for iband,(band,timiter) in enumerate(zip(bands, imaps)):
        debug('Computing coadd for band', band)

        # coadded weight map (moo)
        cow    = np.zeros((H,W), np.float32)
        # coadded weighted image map
        cowimg = np.zeros((H,W), np.float32)

        kwargs = dict(cowimg=cowimg, cow=cow)

        if detmaps:
            # detection map inverse-variance (depth map)
            psfdetiv = np.zeros((H,W), np.float32)
            C.psfdetivs.append(psfdetiv)
            kwargs.update(psfdetiv=psfdetiv)
            # galaxy detection map inverse-variance (galdepth map)
            galdetiv = np.zeros((H,W), np.float32)
            C.galdetivs.append(galdetiv)
            kwargs.update(galdetiv=galdetiv)

        if mods is not None:
            # model image
            cowmod = np.zeros((H,W), np.float32)
            # chi-squared image
            cochi2 = np.zeros((H,W), np.float32)
            kwargs.update(cowmod=cowmod, cochi2=cochi2)

        if blobmods is not None:
            # model image
            cowblobmod = np.zeros((H,W), np.float32)
            kwargs.update(cowblobmod=cowblobmod)

        if unweighted:
            # unweighted image
            coimg  = np.zeros((H,W), np.float32)
            if mods is not None:
                # unweighted model
                comod  = np.zeros((H,W), np.float32)
            if blobmods is not None:
                coblobmod  = np.zeros((H,W), np.float32)
            # number of exposures
            con    = np.zeros((H,W), np.int16)
            kwargs.update(coimg=coimg)

        # We have *three* counters for the number of pixels
        # overlapping each coadd brick pixel:
        #
        # - "con" counts the pixels included in the unweighted coadds.
        #   This map is not passed outside this function or used
        #   anywhere else.
        #
        # - "congood" counts pixels with (iv > 0).  This gets passed
        #   to the *write_coadd_images* function, where it gets
        #   written to the *nexp* maps.
        #
        # - "nobs" counts all pixels, regardless of masks.  This gets
        #   sampled at *xy* positions, and ends up in the tractor
        #   catalog "nobs" column.
        #
        # (you want to know the number of observations within the
        # source footprint, not just the peak pixel which may be
        # saturated, etc.)

        if ngood:
            congood = np.zeros((H,W), np.int16)
            kwargs.update(congood=congood)

        if xy or allmasks or anymasks:
            # These match the type of the "DQ" images.
            # "any" mask
            ormask  = np.zeros((H,W), np.int16)
            # "all" mask
            andmask = np.empty((H,W), np.int16)
            from functools import reduce
            allbits = reduce(np.bitwise_or, DQ_BITS.values())
            andmask[:,:] = allbits
            kwargs.update(ormask=ormask, andmask=andmask)
        if xy or allmasks:
            # number of observations
            nobs = np.zeros((H,W), np.int16)

        if nsatur:
            satmap = np.zeros((H,W), np.int16)

        if psfsize:
            psfsizemap = np.zeros((H,W), np.float32)
            # like "cow", but constant invvar per-CCD;
            # only required for psfsizemap
            flatcow = np.zeros((H,W), np.float32)
            kwargs.update(psfsize=psfsizemap)

        if max:
            maximg = np.zeros((H,W), np.float32)
            C.maximgs.append(maximg)

        if psf_images:
            psf_img = 0.

        for R in timiter:
            if R is None:
                continue
            itim,Yo,Xo,iv,im,mo,bmo,dq = R
            tim = tims[itim]

            if plots:
                _make_coadds_plots_1(im, band, mods, mo, iv, unweighted,
                                     dq, satur_val, allresids, ps, H, W,
                                     tim, Yo, Xo)
            # invvar-weighted image
            cowimg[Yo,Xo] += iv * im
            cow   [Yo,Xo] += iv

            goodpix = None
            if unweighted:
                if dq is None:
                    goodpix = 1
                else:
                    # include SATUR pixels if no other pixels exists
                    okbits = 0
                    for bitname in ['satur']:
                        okbits |= DQ_BITS[bitname]
                    brightpix = ((dq & okbits) != 0)
                    if satur_val is not None:
                        # HACK -- force SATUR pix to be bright
                        im[brightpix] = satur_val
                    # Include these pixels if none other exist??
                    for bitname in ['interp']:
                        okbits |= DQ_BITS[bitname]
                    goodpix = ((dq & ~okbits) == 0)

                coimg[Yo,Xo] += goodpix * im
                con  [Yo,Xo] += goodpix

            if xy or allmasks or anymasks:
                if dq is not None:
                    ormask [Yo,Xo] |= dq
                    andmask[Yo,Xo] &= dq
            if xy or allmasks:
                # raw exposure count
                nobs[Yo,Xo] += 1
            if xy:
                # mjd_min/max
                update = np.logical_or(mjd_argmins[Yo,Xo] == -1,
                                       (mjd_argmins[Yo,Xo] > -1) *
                                       (mjds[itim] < mjds[mjd_argmins[Yo,Xo]]))
                mjd_argmins[Yo[update],Xo[update]] = itim
                update = np.logical_or(mjd_argmaxs[Yo,Xo] == -1,
                                       (mjd_argmaxs[Yo,Xo] > -1) *
                                       (mjds[itim] > mjds[mjd_argmaxs[Yo,Xo]]))
                mjd_argmaxs[Yo[update],Xo[update]] = itim
                del update

            if nsatur and dq is not None:
                satmap[Yo,Xo] += (1*((dq & DQ_BITS['satur'])>0))

            if psfsize:
                # psfnorm is in units of 1/pixels.
                # (eg, psfnorm for a gaussian is 1./(2.*sqrt(pi) * psf_sigma) )
                # Neff is in pixels**2
                neff = 1./tim.psfnorm**2
                # Narcsec is in arcsec**2
                narcsec = neff * tim.wcs.pixel_scale()**2
                # Make smooth maps -- don't ignore CRs, saturated pix, etc
                iv1 = 1./tim.sig1**2
                psfsizemap[Yo,Xo] += iv1 * (1. / narcsec)
                flatcow   [Yo,Xo] += iv1

            if psf_images:
                from astrometry.util.util import lanczos3_interpolate
                h,w = tim.shape
                patch = tim.psf.getPointSourcePatch(w//2, h//2).patch
                patch /= np.sum(patch)
                # In case the tim and coadd have different pixel scales,
                # resample the PSF stamp.
                ph,pw = patch.shape
                pscale = tim.imobj.pixscale / targetwcs.pixel_scale()
                coph = int(np.ceil(ph * pscale))
                copw = int(np.ceil(pw * pscale))
                coph = 2 * (coph//2) + 1
                copw = 2 * (copw//2) + 1
                # want input image pixel coords that change by 1/pscale
                # and are centered on pw//2, ph//2
                cox = np.arange(copw) * 1./pscale
                cox += pw//2 - cox[copw//2]
                coy = np.arange(coph) * 1./pscale
                coy += ph//2 - coy[coph//2]
                fx,fy = np.meshgrid(cox,coy)
                fx = fx.ravel()
                fy = fy.ravel()
                ix = (fx + 0.5).astype(np.int32)
                iy = (fy + 0.5).astype(np.int32)
                dx = (fx - ix).astype(np.float32)
                dy = (fy - iy).astype(np.float32)
                copsf = np.zeros(coph*copw, np.float32)
                lanczos3_interpolate(ix, iy, dx, dy, [copsf], [patch])
                copsf = copsf.reshape((coph,copw))
                copsf /= copsf.sum()
                if plots:
                    _make_coadds_plots_2(patch, copsf, psf_img, tim, band, ps)

                psf_img += copsf / tim.sig1**2

            if detmaps:
                # point-source depth
                detsig1 = tim.sig1 / tim.psfnorm
                psfdetiv[Yo,Xo] += (iv > 0) * (1. / detsig1**2)
                # Galaxy detection map
                gdetsig1 = tim.sig1 / tim.galnorm
                galdetiv[Yo,Xo] += (iv > 0) * (1. / gdetsig1**2)

            if ngood:
                congood[Yo,Xo] += (iv > 0)

            if mods is not None:
                # straight-up
                comod[Yo,Xo] += goodpix * mo
                # invvar-weighted
                cowmod[Yo,Xo] += iv * mo
                # chi-squared
                cochi2[Yo,Xo] += iv * (im - mo)**2
                del mo

            if blobmods is not None:
                # straight-up
                coblobmod[Yo,Xo] += goodpix * bmo
                # invvar-weighted
                cowblobmod[Yo,Xo] += iv * bmo
                del bmo
            del goodpix

            if max:
                maximg[Yo,Xo] = np.maximum(maximg[Yo,Xo], im * (iv>0))

            del Yo,Xo,im,iv
            # END of loop over tims
        # Per-band:
        cowimg /= np.maximum(cow, tinyw)
        C.coimgs.append(cowimg)
        if coweights:
            C.cowimgs.append(cow)
        if mods is not None:
            cowmod  /= np.maximum(cow, tinyw)
            C.comods.append(cowmod)
            coresid = cowimg - cowmod
            coresid[cow == 0] = 0.
            C.coresids.append(coresid)

        if blobmods is not None:
            cowblobmod  /= np.maximum(cow, tinyw)
            C.coblobmods.append(cowblobmod)
            coblobresid = cowimg - cowblobmod
            coblobresid[cow == 0] = 0.
            C.coblobresids.append(coblobresid)

        if xy or allmasks:
            # If these was no coverage, don't set ALLMASK
            andmask[nobs == 0] = 0
        if allmasks:
            C.allmasks.append(andmask)
        if anymasks:
            C.anymasks.append(ormask)
        if nsatur:
            C.satmaps.append(satmap >= nsatur)
        if psf_images:
            C.psf_imgs.append(psf_img / np.sum(psf_img))

        if unweighted:
            coimg  /= np.maximum(con, 1)
            del con

            if plots:
                _make_coadds_plots_3(cowimg, cow, coimg, band, ps)

            # Patch pixels with no data in the weighted coadd.
            cowimg[cow == 0] = coimg[cow == 0]
            del coimg
            if mods is not None:
                cowmod[cow == 0] = comod[cow == 0]
                del comod
            if blobmods is not None:
                cowblobmod[cow == 0] = coblobmod[cow == 0]
                del coblobmod
        if xy:
            if ngood:
                C.T.ngood  [:,iband] = congood[iy,ix]
            C.T.nobs   [:,iband] = nobs   [iy,ix]
            C.T.anymask[:,iband] = ormask [iy,ix]
            C.T.allmask[:,iband] = andmask[iy,ix]
            if detmaps: