#!/usr/bin/env python
# encoding: utf-8

# A library of AFW (Application FrameWork), related to images
import lsst.afw.image as afwImage
# A library of AFW, related to geometry
import lsst.afw.geom  as afwGeom

import sys

def main():
    """
    Arrange patches to make a large image, and save it in "large.fits"

    make-large-image PATCHES...
        (e.g. make-large-image *,*.fits)
    """

    if len(sys.argv) <= 1:
        print("Input 1 patch or more.")
        return 1

    largeImage = makeLargeImageFromPatches(sys.argv[2:])

    print("Saving (Wait a while)...")
    largeImage.writeFits(sys.argv[1])


def makeLargeImageFromPatches(filenames):
    """
    @param filenames List of file names
    @return afwImage.ExposureF
    """

    # Loading all image would be dangerous by memory limitation.
    # Load metric information only.
    bboxes = []
    for filename in filenames:
        print("Loading a bounding box:", filename, "...")
        bboxes.append(
            # Get the bounding box of this patch
            # in the parent (= the tract)
            afwImage.ExposureF(filename).getBBox(afwImage.PARENT)
        )

    # Calculate the bbox that includes all the bboxes.
    #     Return values of getMin/Max method family are inclusive
    #     (range = [min, max]) cf. getBegin/End (used below)
    x0 = min(bbox.getMinX() for bbox in bboxes)
    y0 = min(bbox.getMinY() for bbox in bboxes)
    x1 = max(bbox.getMaxX() for bbox in bboxes)
    y1 = max(bbox.getMaxY() for bbox in bboxes)

    width  = x1 - x0 + 1
    height = y1 - y0 + 1

    # Prepare a large image
    print("Preparing a large image: ({}, {})".format(width, height))
    exposure = afwImage.ExposureF(width, height)

    # Set the position of this image in the tract
    exposure.setXY0(afwGeom.Point2I(x0, y0))

    # Extract the image data from the exposure
    largeImage = exposure.getMaskedImage()

    # Paste patches into largeImage
    for filename in filenames:
        print("Loading a patch:", filename, "...")

        # Load a patch again
        patch = afwImage.ExposureF(filename)
        bbox = patch.getBBox(afwImage.PARENT)

        # The position of this patch within largeImage:
        #     Return values of getBegin/End method family are one past the end
        #     (range = [begin,end) ) cf. getMin/Max (above)
        xp0 = bbox.getBeginX() - x0
        yp0 = bbox.getBeginY() - y0
        xp1 = bbox.getEndX  () - x0
        yp1 = bbox.getEndY  () - y0

        # Paste the patch (As for the notation in [], see numpy's document)
        largeImage[xp0:xp1, yp0:yp1] = patch.getMaskedImage()[:]

        # Blindly trust that all the patches have the same WCS information etc,
        # and set arbitrary one of them in the large image.
        if not exposure.hasWcs() and patch.hasWcs():
            exposure.setCalib   (patch.getCalib   ())
            exposure.setFilter  (patch.getFilter  ())
            exposure.setMetadata(patch.getMetadata())
            exposure.setWcs     (patch.getWcs     ())

    return exposure


if __name__ == "__main__":
    main()