from __future__ import print_function
import numpy as np
from orphics import maps,io,stats
from enlib import enmap,resample,bench

def mask_map(imap,iys,ixs,hole_arc,hole_frac=0.6):
    shape,wcs = imap.shape,imap.wcs
    Ny,Nx = shape[-2:]
    px = maps.resolution(shape,wcs)*60.*180./np.pi
    hole_n = int(round(hole_arc/px))
    hole_ny = hole_nx = hole_n
    oshape,owcs = enmap.geometry(pos=(0.,0.),shape=(2*hole_n,2*hole_n),res=px*np.pi/180./60.)
    modrmap = enmap.modrmap(oshape,owcs)
    mask = enmap.ones(shape,wcs)
    
    for iy,ix in zip(iys,ixs):
        if iy<=hole_n or ix<=hole_n or iy>=(Ny-hole_n) or ix>=(Nx-hole_n): continue
        vslice = imap[np.int(iy-hole_ny):np.int(iy+hole_ny),np.int(ix-hole_nx):np.int(ix+hole_nx)]
        vslice[modrmap<(hole_frac*hole_arc)*np.pi/180./60.] = np.nan # !!!! could cause a bias
        mask[np.int(iy-hole_ny):np.int(iy+hole_ny),np.int(ix-hole_nx):np.int(ix+hole_nx)][modrmap<hole_arc*np.pi/180./60.] = 0
        
    return mask



def inpaint_map(imap,ras,decs,radii_tags,radii_dict,tot_power_2d,seed=None):
    """
    Brute-force inpaints a map in circular regions.

    imap -- (Ny,Nx) enmap
    ras  -- list of RA of centers in degrees, length M
    decs -- list of DEC of centers in degrees, length M
    radii -- list of strings specifying radius tag for each object, length M
    radii_dict -- dict mapping radius tag string to float value of radius in arcminutes, length K
                  Most expensive operation scales with K.
    tot_power_2d -- total power in map in physical units (e.g. uK^2 radians)

    """


    class G:
        pass

    geometries = {}
    tags = radii_dict.keys()
    for key in tags:
        geometries[key] = calculate_circular_geometry(shape,wcs)




def prepare_circular_mask(stamp,hole_arcminutes):
    '''This routine accepts a stamp as a 2D array. The region
    within hole_arcminutes is masked, and the mean of the region outside it
    is subtracted from the stamp. It returns the masked stamp and the value
    of the mean of the unmasked region that needs to be added back after
    inpainting.

    '''

    stamp[stamp.modrmap()<hole_arcminutes*np.pi/180./60.] = np.nan
    unmasked_mean = np.nanmean(stamp)
    stamp -= unmasked_mean

    return stamp, unmasked_mean


def get_geometry_shapes(shape,wcs,hole_arcminutes):
    modrmap = enmap.modrmap(shape,wcs)
    m1 = np.where(modrmap.reshape(-1)<hole_arcminutes*np.pi/180./60.)[0]
    m2 = np.where(modrmap.reshape(-1)>=hole_arcminutes*np.pi/180./60.)[0]
    return m1.size,m2.size


def make_circular_geometry(shape,wcs,context_arcmin,hole_arcmin,power2d,buffer_factor = 2,verbose=False):
    '''Makes the circular geometry matrices that need to be pre-calculated for later inpainting.

    Arguments
    ---------

    input2DPower - ndarray containing 2D power spectrum of a map. It need not already have been
                   downsampled to the shape of the stamp cutout
    inputLy
    inputLx      - the fourier wavenumbers corresponding to the y and x axes of the stamp cutout
    stampArc     - the width in arcminutes of the stamp cut out
    stampPxX      - the pixel width in arcminutes of the stamp cut out in the x direction
    stampPxY      - the pixel width in arcminutes of the stamp cut out in the y direction
    holeArc      - the radius of the circular hole in arcminutes
    bufferFactor - the pixel covariance matrix will be calculated on a periodic stamp larger by this
                   factor
    verbose      - True if you want more commentary

    Returns
    -------

    meanMul      - a matrix that has shape (nh,nc) where nh is the number of pixels in the hole and
                   nc is the number of pixels outside (in the "context"). It should be multiplied
                   by a vector (nc) containing pixels outside to get a vector (nh) for the mean
                   value of the pixels inside
    covRoot      - a (nh,nh) sqrt(covariance matrix) that can be used to generate a random realization
                   in the hole. This can be generated by multiplying the sqrt of cov by a vector (nh)
                   of standard normal variables. The generated vector should be added to the mean value
                   obtained using meanMul
    pcov         - the pixel-pixel covariance matrix used in intermediate steps, if you want to re-use it
    targetTemplate - a liteMap template of the stamp cutout
    m1           - a boolean array that can be used to select the hole region
    m2           - a boolean array that can be used to select the context region


    '''

    arc = context_arcmin
    res = maps.resolution(shape,wcs)*60.*180./np.pi
                                                      
    bshape,bwcs = maps.rect_geometry(width_arcmin=arc*buffer_factor,px_res_arcmin=res)
    tshape,twcs = maps.rect_geometry(width_arcmin=arc,px_res_arcmin=res)
    sny,snx = tshape
    bmodlmap = enmap.modlmap(bshape,bwcs)
    modlmap = enmap.modlmap(shape,wcs)

    if verbose: print ("Downsampling...")
    Niy,Nix = shape[-2:]
    Noy,Nox = bshape[-2:]
    
    # print(bshape,tshape,power2d.shape)
    # io.plot_img(np.fft.fftshift(np.log10(power2d)))
    #out_power = resample.resample_fft(power2d,bshape[-2:],axes=[-2,-1])
    #out_power = np.fft.ifftshift(resample.resample_fft(np.fft.fftshift(power2d),bshape[-2:],axes=[-2,-1]))
    out_power = resample.resample_bin(power2d,factors=[float(Noy)/Niy,float(Nox)/Nix],axes=[-2,-1])
    # io.plot_img(np.fft.fftshift(np.log10(out_power)))
    # print(out_power.shape)
    
    if verbose: print ("Starting slow part...")
    d = maps.diagonal_cov(out_power)
    with bench.show("pixcov"):
        pcov = maps.pixcov(bshape,bwcs,d)[0,0,:sny,:snx,:sny,:snx]
    modrmap = enmap.modrmap(tshape,twcs)
    m1 = np.where(modrmap.reshape(-1)<hole_arcmin*np.pi/180./60.)[0]
    m2 = np.where(modrmap.reshape(-1)>=hole_arcmin*np.pi/180./60.)[0]

    with bench.show("geom"):
        meanMul, cov = get_geometry(pcov.reshape(sny*snx,sny*snx),m1,m2)
    
    covRoot = stats.eig_pow(cov,0.5)

    return meanMul, covRoot, pcov, tshape,twcs, m1, m2


def get_geometry(pixcov,m1,m2):
    # m1 is hole
    # m2 is context

    Cinv = np.linalg.inv(pixcov)
    cslice = Cinv[m1][:,m1]
    meanMul1 = np.linalg.inv(cslice)

    mul2 = Cinv[m1][:,m2]
    meanMul = np.dot(-meanMul1,mul2)
    cov = np.linalg.pinv(Cinv[m1][:,m1])
    return meanMul, cov


def fill_hole(masked_stamp,meanMatrix,holeArc,m1,m2,covRoot=None):
    '''Returns the result of an inpaint operation as a 1d unraveled vector

    Arguments
    ---------

    masked_liteMap_stamp - the cutout stamp that contains a masked hole and
                           unmasked context
    meanMatrix           - an (nh,nc) matrix. See docs for make_circular_geometry
    holeArc              - radius of masked hole in arcminutes
    m1                   - a 1d boolean selecting the hole region on an unraveled stamp
    m2                   - a 1d boolean selecting the context region on an unraveled stamp
    covRoot              - the square root of the covariance matrix inside the hole. See
                           docs for make_circular_geometry. If unspecified, the random
                           realization returned is zero.

    Returns
    -------

    mean  -  a 1d (nh) vector containing the mean inpainted value constrained according
             to the context
    rand  -  a 1d (nh) vector containing a random realization inside the hole
    sim   -  a 1d (nh) vector containing the sum of mean and rand

    '''

    mean = np.dot(meanMatrix,masked_stamp.reshape(-1)[m2])
    r = np.random.normal(0.,1.,size=(m1.size))
    if covRoot is not None:
        rand = np.dot(covRoot,r)
    else:
        rand = 0.
    sim = mean + rand
    return mean, rand, sim

                                                      
