from __future__ import division, print_function
import numpy as np, glob, re, sys, os, pixell
from pixell import utils, enmap, bunch, mpi, enplot

default_pathformat = "tile%(y)03d_%(x)03d.fits"

def leaftile(idir, odir, tsize=675, comm=None, verbose=False, lrange=[0,-6],
        monolithic=False, slice=None):
    """Given a input directory containing a tiled dmap in standard
    ordering, outputs a leaflet-compatible hierarchy of tiles in
    odir with tile size tsize."""
    # First create our base tiles. These have opposite y ordering than
    # dmap tiles, and may be different-sized.
    otilename = "tile_%(y)d_%(x)d.fits"
    if not monolithic:
        itilename = idir + "/tile%(y)03d_%(x)03d.fits"
        itile1, itile2 = None, None
    else:
        itilename = idir
        itile1, itile2 = (0,0), (1,1)
    retile(itilename,
            "%s/%d/%s" % (odir,lrange[0],otilename),
            ocorner=(np.pi/2,-np.pi), otilesize=(-tsize,tsize),
            comm=comm, verbose=verbose, itile1=itile1, itile2=itile2, slice=slice)
    # Then loop over the smaller levels
    for level in range(lrange[0]-1,lrange[1],-1):
        if comm: comm.barrier()
        combine_tiles("%s/%d/%s" % (odir, level+1, otilename),
                "%s/%d/%s" % (odir, level, otilename), tyflip=True,
                pad_to=tsize, comm=comm, verbose=verbose)

def combine_tiles(ipathfmt, opathfmt, combine=2, downsample=2,
        itile1=(None,None), itile2=(None,None), tyflip=False, txflip=False,
        pad_to=None, comm=None, verbose=False):
    """Given a set of tiles on disk at locaiton ipathfmt % {"y":...,"x"...},
    combine them into larger tiles, downsample and write the result to
    opathfmt % {"y":...,"x":...}. x and y must be contiguous and start at 0.
    
    reftile[2] indicates the tile coordinates of the first valid input tile.
    This needs to be specified if not all tiles of the logical tiling are
    physically present.

    tyflip and txflip indicate if the tiles coordinate system is reversed
    relative to the pixel coordinates or not."
    """
    # Expand combine and downsample to 2d
    combine    = np.zeros(2,int)+combine
    downsample = np.zeros(2,int)+downsample
    if pad_to is not None:
        pad_to = np.zeros(2,int)+pad_to
    # Handle optional mpi
    rank, size = (comm.rank, comm.size) if comm is not None else (0, 1)
    # Find the range of input tiles
    itile1, itile2 = find_tile_range(ipathfmt, itile1, itile2)
    # Read the first tile to get its size information
    ibase = enmap.read_map(ipathfmt % {"y":itile1[0],"x":itile1[1]})*0
    # Find the set of output tiles we need to consider
    otile1 = itile1//combine
    otile2 = (itile2-1)//combine+1
    # And loop over them
    oyx = [(oy,ox) for oy in range(otile1[0],otile2[0]) for ox in range(otile1[1],otile2[1])]
    for i in range(rank, len(oyx), size):
        oy, ox = oyx[i]
        # Read in all associated tiles into a list of lists
        rows = []
        for dy in range(combine[0]):
            iy = oy*combine[0] + dy
            if iy >= itile2[0]: continue
            cols = []
            for dx in range(combine[1]):
                ix = ox*combine[1] + dx
                if ix >= itile2[1]: continue
                if iy < itile1[0] or ix < itile1[1]:
                    # The first tiles are missing on disk, but are
                    # logically a part of the tiling. Use ibase,
                    # which has been zeroed out.
                    cols.append(ibase)
                else:
                    itname = ipathfmt % {"y": iy, "x": ix}
                    cols.append(enmap.read_map(itname))
            if txflip: cols = cols[::-1]
            rows.append(cols)
        # Stack them next to each other into a big tile
        if tyflip: rows = rows[::-1]
        omap = enmap.tile_maps(rows)
        # Downgrade if necessary
        if np.any(downsample>1):
            omap = enmap.downgrade(omap, downsample)
        if pad_to is not None:
            # Padding happens towards the end of the tiling,
            # which depends on the flip status
            padding = np.array([[0,0],[pad_to[0]-omap.shape[-2],pad_to[1]-omap.shape[-1]]])
            if tyflip: padding[:,0] = padding[::-1,0]
            if txflip: padding[:,1] = padding[::-1,1]
            omap = enmap.pad(omap, padding)
        # And output
        otname = opathfmt % {"y": oy, "x": ox}
        utils.mkdir(os.path.dirname(otname))
        enmap.write_map(otname, omap)
        if verbose: print(otname)

def retile(ipathfmt, opathfmt, itile1=(None,None), itile2=(None,None),
        otileoff=(0,0), otilenum=(None,None), ocorner=(-np.pi/2,-np.pi),
        otilesize=(675,675), comm=None, verbose=False, slice=None, wrap=True):
    """Given a set of tiles on disk with locations ipathfmt % {"y":...,"x":...},
    retile them into a new tiling and write the result to opathfmt % {"y":...,"x":...}.
    The new tiling will have tile size given by otilesize[2]. Negative size means the
    tiling will to down/left instead of up/right. The corner of the tiling will
    be at sky coordinates ocorner[2] in radians. The new tiling will be pixel-
    compatible with the input tiling - w.g. the wcs will only differ by crpix.

    The output tiling will logically cover the whole sky, but only output tiles
    that overlap with input tiles will actually be written. This can be modified
    by using otileoff[2] and otilenum[2]. otileoff gives the tile indices of the
    corner tile, while otilenum indicates the number of tiles to write."""
    # Set up mpi
    rank, size = (comm.rank, comm.size) if comm is not None else (0, 1)
    # Expand any scalars
    if otilesize is None: otilesize=(675,675)
    otilesize = np.zeros(2,int)+otilesize
    otileoff  = np.zeros(2,int)+otileoff
    # Find the range of input tiles
    itile1, itile2 = find_tile_range(ipathfmt, itile1, itile2)
    # To fill in the rest of the information we need to know more
    # about the input tiling, so read the first tile
    ibase = enmap.read_map(ipathfmt % {"y":itile1[0],"x":itile1[1]})
    if slice: ibase = eval("ibase"+slice)
    itilesize = ibase.shape[-2:]
    ixres = ibase.wcs.wcs.cdelt[0]
    nphi  = utils.nint(360/np.abs(ixres))
    ntile_wrap = nphi//otilesize[1]
    # Find the pixel position of our output corners according to the wcs.
    # This is the last place we need to do a coordinate transformation.
    # All the rest can be done in pure pixel logic.
    pixoff = np.round(ibase.sky2pix(ocorner)).astype(int)
    # Find the range of output tiles
    def pix2otile(pix, ioff, osize): return (pix-ioff)//osize
    otile1 = pix2otile(itile1*itilesize,   pixoff, otilesize)
    otile2 = pix2otile(itile2*itilesize-1, pixoff, otilesize)
    otile1, otile2 = np.minimum(otile1,otile2), np.maximum(otile1,otile2)
    otile2 += 1
    # We can now loop over output tiles
    cache = [None,None,None]
    oyx = [(oy,ox) for oy in range(otile1[0],otile2[0]) for ox in range(otile1[1],otile2[1])]
    for i in range(rank, len(oyx), size):
        otile = np.array(oyx[i])
        # Find out which input tiles overlap with this output tile.
        # Our tile stretches from opix1:opix2 relative to the global input pixels
        opix1 = otile*otilesize + pixoff
        opix2 = (otile+1)*otilesize + pixoff
        # output tiles and input tiles may increase in opposite directions
        opix1, opix2 = np.minimum(opix1,opix2), np.maximum(opix1,opix2)
        try: omap = read_area(ipathfmt, [opix1,opix2],itile1=itile1, itile2=itile2,cache=cache, slice=slice)
        except (IOError, OSError): continue
        x = otile[1]+otileoff[1]
        if wrap: x %= ntile_wrap
        oname = opathfmt % {"y":otile[0]+otileoff[0],"x":x}
        utils.mkdir(os.path.dirname(oname))
        enmap.write_map(oname, omap)
        if verbose: print(oname)

def read_monolithic(idir, verbose=True, slice=None, dtype=None):
    # Find the range of input tiles
    ipathfmt = idir + "/tile%(y)03d_%(x)03d.fits"
    itile1, itile2 = find_tile_range(ipathfmt)
    def read(fname):
        m = enmap.read_map(fname)
        if slice: m = eval("m" + slice)
        return m
    # Read the first and last tile to get the total dimensions
    m1 = read(ipathfmt % {"y":itile1[0],"x":itile1[1]})
    m2 = read(ipathfmt % {"y":itile2[0]-1,"x":itile2[1]-1})
    wy,wx  = m1.shape[-2:]
    oshape = tuple(np.array(m1.shape[-2:])*(itile2-itile1-1) + np.array(m2.shape[-2:]))
    if dtype is None: dtype = m1.dtype
    omap  = enmap.zeros(m1.shape[:-2] + oshape, m1.wcs, dtype)
    del m1, m2
    # Now loop through all tiles and copy them in to the correct position
    for ty in range(itile1[0],itile2[0]):
        for tx in range(itile1[1],itile2[1]):
            m = read(ipathfmt % {"y":ty,"x":tx})
            oy = ty - itile1[0]
            ox = tx - itile1[1]
            omap[...,oy*wy:(oy+1)*wy,ox*wx:(ox+1)*wx] = m
            if verbose: print(ipathfmt % {"y":ty,"x":tx})
    return omap

def monolithic(idir, ofile, verbose=True, slice=None, dtype=None):
    omap = read_monolithic(idir, verbose=verbose, slice=slice, dtype=dtype)
    enmap.write_map(ofile, omap)

def range_overlap(a,b):
    return np.array([np.maximum(a[0],b[0]),np.minimum(a[1],b[1])])

def find_tile_range(pathfmt, tile1=(None,None), tile2=(None,None)):
    """Given a path format with with labeled formats including y
    and x (like %(y)d), Return the range of tiles available, in
    the form of [{start,end},{y,x}] tile indices.
    If tile1[2] is specified, then this will override the start
    of the result. If tile2[2] is specified, then it will override
    the end of the result."""
    if tile1 is None: tile1 = (None,None)
    if tile2 is None: tile2 = (None,None)
    # If both offset and ntile are specified, we don't need any
    # complicated disk search.
    if tile1[0] is not None and tile1[1] is not None and tile2[0] is not None and tile2[1] is not None: return np.asarray(tile1), np.asarray(tile2)
    # Find the min/max on disk. We do that by constructing a glob to
    # roughly match them, and them filtering them with a regex.
    ranges = [None,None]
    gstr  = utils.format_to_glob(pathfmt)
    regex = utils.format_to_regex(pathfmt)
    files = glob.glob(gstr)
    if len(files) == 0:
        raise ValueError("Found no files matching path format!")
    for file in files:
        m = re.match(regex, file)
        if not m: continue
        try: yx = [int(m.group(name)) for name in ["y","x"]]
        except IndexError: yx = [0,0]
        for i in range(2):
            if ranges[i] is None: ranges[i] = [yx[i],yx[i]+1]
            ranges[i] = [min(ranges[i][0],yx[i]),max(ranges[i][1],yx[i]+1)]
    # Override if needed
    for i in range(2):
        if tile1[i] is not None:
            ranges[i][0] = tile1[i]
        if tile2[i]  is not None:
            ranges[i][1] = tile2[i]
    ranges = np.array(ranges)
    return ranges[:,0], ranges[:,1]

def read_tileset_geometry(ipathfmt, itile1=(None,None), itile2=(None,None)):
    itile1, itile2 = find_tile_range(ipathfmt, itile1, itile2)
    mfile1 = ipathfmt % {"y":itile1[0],"x":itile1[1]}
    mfile2 = ipathfmt % {"y":itile2[0]-1,"x":itile2[1]-1}
    m1 = enmap.read_map(mfile1)
    m2 = m1 if mfile1 == mfile2 else enmap.read_map(mfile2)
    wy,wx  = m1.shape[-2:]
    oshape = tuple(np.array(m1.shape[-2:])*(itile2-itile1-1) + np.array(m2.shape[-2:]))
    return bunch.Bunch(shape=m1.shape[:-2]+oshape, wcs=m1.wcs, dtype=m1.dtype,
            tshape=m1.shape[-2:])

def read_area(ipathfmt, opix, itile1=(None,None), itile2=(None,None), verbose=False,
        cache=None, slice=None, wrap=True):
    """Given a set of tiles on disk with locations ipathfmt % {"y":...,"x":...},
    read the data corresponding to the pixel range opix[{from,to},{y,x}] in
    the full map."""
    opix = np.asarray(opix)
    # Find the range of input tiles
    itile1, itile2 = find_tile_range(ipathfmt, itile1, itile2)
    # To fill in the rest of the information we need to know more
    # about the input tiling, so read the first tile
    if cache is None or cache[2] is None:
        geo = read_tileset_geometry(ipathfmt, itile1=itile1, itile2=itile2)
    else: geo = cache[2]
    if cache is not None: cache[2] = geo
    # Determine tile wrapping
    npix_phi  = np.abs(360./geo.wcs.wcs.cdelt[0])
    ntile_phi = utils.nint(npix_phi/geo.tshape[-1])

    isize = geo.tshape
    osize = opix[1]-opix[0]
    omap  = enmap.zeros(geo.shape[:-2]+tuple(osize), geo.wcs, geo.dtype)
    # Find out which input tiles overlap with this output tile.
    # Our tile stretches from opix1:opix2 relative to the global input pixels
    it1 = opix[0]//isize
    it2 = (opix[1]-1)//isize+1
    noverlap = 0
    for ity in range(it1[0],it2[0]):
        if ity < itile1[0] or ity >= itile2[0]: continue
        # Start/end of this tile in global input pixels
        ipy1, ipy2 = ity*isize[0], (ity+1)*isize[0]
        overlap = range_overlap(opix[:,0],[ipy1,ipy2])
        oy1,oy2 = overlap-opix[0,0]
        iy1,iy2 = overlap-ipy1
        for itx in range(it1[1],it2[1]):
            if wrap: itx_wrap = itx % ntile_phi
            if itx_wrap < itile1[1] or itx_wrap >= itile2[1]: continue
            ipx1, ipx2 = itx*isize[1], (itx+1)*isize[1]
            overlap = range_overlap(opix[:,1],[ipx1,ipx2])
            ox1,ox2 = overlap-opix[0,1]
            ix1,ix2 = overlap-ipx1
            # Read the input tile and copy over
            iname = ipathfmt % {"y":ity,"x":itx_wrap}
            if cache is None or cache[0] != iname:
                imap  = enmap.read_map(iname)
                if slice: imap = eval("imap"+slice)
            else: imap = cache[1]
            if cache is not None:
                cache[0], cache[1] = iname, imap
            if verbose: print(iname)
            # Edge input tiles may be smaller than the standard
            # size.
            ysub = isize[0]-imap.shape[-2]
            xsub = isize[1]-imap.shape[-1]
            # If the input map is too small, there may actually be
            # zero overlap.
            if oy2-ysub <= oy1 or ox2-xsub <= ox1: continue
            omap[...,oy1:oy2-ysub,ox1:ox2-xsub] = imap[...,iy1:iy2-ysub,ix1:ix2-xsub]
            noverlap += 1
    if noverlap == 0:
        raise IOError("No tiles for tiling %s in range %s" % (ipathfmt, ",".join([":".join([str(p) for p in r]) for r in opix.T])))
    # Set up the wcs for the output tile
    omap.wcs.wcs.crpix -= opix[0,::-1]
    return omap

def read_retile(ipathfmt, tpos, otilesize=None, pixoff=(0,0), margin=0,
        itile1=(None,None), itile2=(None,None), verbose=False):
    """Read a single tile from the tiling at ipathfmt % {"y":tpos[0],"x":tpos[1]},
    returning it as an enmap. If otilesize or pixoff are specified, then
    the tiling will be from a retiling with that tile size and pixel offset,
    which means that behind the scenes multiple tiles will be read and stitched
    together. If margin[{left,right},{y,x}] is not zero, then it specifies the
    number of pixels to extend the tile by. The tile will be extended with
    data from the tiling, i.e. from the neighboring tiles."""
    if otilesize is None:
        geo = read_tileset_geometry(ipathfmt, itile1, itile2)
        otilesize = geo.tshape
    tpos   = np.zeros(2,int)+tpos
    pixoff = np.zeros(2,int)+pixoff
    margin = np.zeros((2,2),int)+margin
    pixbox = np.array([tpos*otilesize,(tpos+1)*otilesize])+pixoff
    pixbox[0] -= margin[0]
    pixbox[1] += margin[1]
    return read_area(ipathfmt, pixbox, itile1=itile1, itile2=itile2, verbose=verbose)

def retile_iterator(ipathfmt, otilesize=None, pixoff=(0,0), margin=0,
        itile1=(None,None), itile2=(None,None), comm=None, verbose=False):
    """Iterator that yields a series of tiles from the tileset given by
    ipathfmt. See read_retile for how margin, otilesize and pixoff
    allow you to iterate over a modified tiling. This currently just
    calls read_area repeatedly, so it is not as efficient as it could have
    been."""
    # Handle mpi
    rank, nproc = (0,1) if comm is None else (comm.rank, comm.size)
    # Find the number of tiles to iterate over
    geo = read_tileset_geometry(ipathfmt, itile1, itile2)
    if otilesize is None: otilesize = geo.tshape
    otilesize = np.array(otilesize)
    notile = (np.array(geo.shape[-2:])-pixoff+otilesize-1)//otilesize
    tyx = [(y,x) for y in range(notile[0]) for x in range(notile[1])]
    for i in range(rank, len(tyx), nproc):
        tpos = tyx[i]
        yield tpos, read_retile(ipathfmt, tpos, otilesize=otilesize, pixoff=pixoff,
                margin=margin, itile1=itile1, itile2=itile2)




