#!/usr/bin/env python3

# Use the astroalign module to align and stack a series of astro images.

# Suggested Workflow;
# If starting with raw images, process them in darktable:
# - darktable *.cr2 (probably including your dark frame,
#   though this script doesn'tsubtract dark frames yet)
#   or if that doesn't work (darktable tends to freeze if you start it
#   with a list of images), try import/folder... in the lighttable tab.
# - lighttable tab: choose the first image
# - darkroom tab: fiddle with it until it looks like what you want
# - Ctrl-C (this copies the operation history, not the image)
# - lighttable tab: Click on each other image in turn and Ctrl-V
# - lighttable tab: export selected
# This will (by default) export to files under darktable_exported/
#
# You can run this script directly on the raw images, but the import
# with rawpy is one of two extremes: by default it does an extreme
# brightness auto-stretch, which leads to incredibly noisy astro images;
# or with no_auto_bright=True, you get no brightness improvement, so you
# lose a lot of the extra info stored in the raw file.
#
# The output of this script is a sequence of aligned images.
# astroalign doesn't do a perfect job (though it's quite a bit better
# than Siril), so I recommend loading the image(s) into GIMP
# and adjusting the alignment to your liking. Then set all the
# layer modes except the bottom one to Addition (or for fun, you
# might want to try Screen, Pin light, Luminance, or one of the LCh *)
#
# The image will be easier to load into GIMP if you use .ora or .tif as
# the export format from this script (ora is the default).
# For ora, the program will write a single file, layers.ora, which
# puts all the layers except the bottom one in Addition mode.
# For tif, the program will write a layers.tif that includes the
# layers as "pages"; GIMP can import pages as layers, but you'll
# need to set the layer modes yourself.
# For any other image format (e.g. png), starstack will write each
# layer as a separate file.

# Copyright 2020 by Akkana Peck: Share and enjoy under the GPLv2 or later.
#
# Originally based on
# https://share.cocalc.com/share/b66ffe0d5b2bc8ff75ac939486710731c2b030f6/astroalign-124/astroalign-py3.ipynb?viewer=share
# (apparently based on an earlier version of astroalign).
# OpenRaster code adapted from Jon Nordby's GIMP OpenRaster file plug-in.


import astroalign
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import zipfile
import xml.etree.ElementTree as ElementTree
import argparse
import sys, os

# rawpy is optional. PIL can read some raw images (at least Canon cr2),
# but may not get the same results as rawpy does.
try:
    import rawpy
except:
    pass

# If the tifffile module is installed and the image extension is .tif
# save multi-page TIFFs that GIMP can open as layers,
# instead of separate images.
try:
    import tifffile
except:
    print("Can't import tifffile")
    pass


class ORAfile:

    layermodes_map = {
        "NORMAL":            "svg:src-over",
        "ADDITION":          "svg:plus",
        "MULTIPLY":          "svg:multiply",
        "SCREEN":            "svg:screen",
        "OVERLAY":           "svg:overlay",
        "DARKEN_ONLY_MODE":  "svg:darken",
        "LIGHTEN_ONLY_MODE": "svg:lighten",
        "DODGE":             "svg:color-dodge",
        "BURN":              "svg:color-burn",
        "HARDLIGHT":         "svg:hard-light",
        "SOFTLIGHT":         "svg:soft-light",
        "DIFFERENCE":        "svg:difference",
    }

    def __init__(self, filename):
        self.orafilename = filename
        if os.path.exists(self.orafilename):
            os.remove(self.orafilename)
        self.orafile = zipfile.ZipFile(self.orafilename, 'w',
                                       compression=zipfile.ZIP_STORED)

        # First file written must be mimetype:
        self.write_file_str('mimetype', 'image/openraster')

        self.ETimage = ElementTree.Element('image')
        self.ETstack = ElementTree.SubElement(self.ETimage, 'stack')

        self.size = None

    def store_layer(self, img, layerno, layername, layerpath=None,
                    layermode="NORMAL"):
        """img is a PIL image"""
        if not self.size:
            self.size = img.size

            a = self.ETimage.attrib
            a['w'] = str(img.size[0])
            a['h'] = str(img.size[1])

        if not layerpath:
            layerpath = 'data/%d.png' % layerno

        # XXX TEMPORARY: can probably get PIL to write directly to the zip.
        tmppng = os.path.join("/tmp", 'tmp.png')
        img.save(tmppng)

        self.orafile.write(tmppng, layerpath)
        os.remove(tmppng)

        layer = ElementTree.Element('layer')
        # self.ETstack.append(layer)
        # Need to insert layers in reverse order: want the first layer
        # on the bottom, later ones higher.
        self.ETstack.insert(0, layer)
        a = layer.attrib
        a['src'] = layerpath
        a['name'] = layername
        a['x'] = "0"
        a['y'] = "0"
        a['opacity'] = "1.0"
        a['composite-op'] = ORAfile.layermodes_map[layermode]
        a['visibility'] = "visible"

    def save_thumbnail(self, img):
        """Turn img into a thumbnail for the ORA file.
           THIS WILL OVERWRITE img WITH A SMALLER VERSION!
        """
        if not self.size:
            raise RuntimeError("Thumbnail before any images have been added")
        w, h = self.size
        # should be at most 256x256, without changing aspect ratio
        if w > h:
            w, h = 256, max(h*256/w, 1)
        else:
            w, h = max(w*256/h, 1), 256

        img.thumbnail((w, h))
        self.store_layer(img, -1, "Thumbnail",
                         layerpath="Thumbnails/thumbnail.png")

    def write_file_str(self, fname, data):
        # work around a permission bug in the zipfile library:
        # http://bugs.python.org/issue3394
        zi = zipfile.ZipInfo(fname)
        zi.external_attr = 0o100644 << 16
        self.orafile.writestr(zi, data)

    def finish(self):
        """Finish and write the ORA file"""

        xmldata = ElementTree.tostring(self.ETimage, encoding='UTF-8')
        self.write_file_str("stack.xml", xmldata)

        self.orafile.close()


def register_all(images, outdir=".", ext="tif", layermode="NORMAL",
                 darkframe=None):
    """Register a set of images (filenames) to the first image.
       Save each image (including the unchanged first one) as a
       set of png images with a_ prepended to the names.
       Input images may be filenames, or may already be numpy arrays.
    """
    if ext.startswith('.'):
        ext = ext[1:]

    if ext.lower() == "tif" and 'tifffile' in sys.modules:
        tiff_multipage = os.path.join(outdir, "layers.tif")
        orafile = None
    elif ext == "ora":
        print("OpenRaster export")
        orafile = ORAfile(os.path.join(outdir, "layers.ora"))
        tiff_multipage = None
    else:
        tiff_multipage = None
        orafile = None

    if darkframe:
        print("Using dark frame", darkframe)
        darkarr = read_image(darkframe)
    else:
        darkarr = None

    for i, img in enumerate(images):
        if type(img) is str:
            layername = os.path.basename(img)
        else:
            layername = "layer %d" % i

        imgarr = read_image(img)

        # Subtract the dark frame.
        # XXX This doesn't work right yet: for cr2, the result has all kinds
        # of new red and blue pixel noise (fully saturated, not subtle)
        # and I haven't figured out why.
        if darkarr is not None:
            try:
                print("Subtracting dark frame from", img)
                # Couldn't subtract dark frame: operands could not be broadcast together with shapes (4000,6000,3) (4024,6024,3)
                # imgarr = imgarr - darkarr
                imgarr -= darkarr
            except Exception as e:
                print("Couldn't subtract dark frame:", e)

        if i == 0:
            # For the first image, don't align it, just save it
            # and its base layer
            baseimg = imgarr
            baselayer = singlelayer(baseimg)
            aligned_arr = baseimg
        else:
            aligned_arr = register(imgarr, baselayer)

        # Now that the image is registered and has the dark frame subtracted,
        # turn it into a PIL Image so it can be saved.
        # This step isn't actually needed for TIFF.
        aligned_img = Image.fromarray(aligned_arr)

        if aligned_img.mode != 'RGB':
            print("Converting monochrome image to save as PNG")
            aligned_img = aligned_img.convert("RGB")

        if tiff_multipage:
            tifffile.imwrite(tiff_multipage, aligned_arr, append=True)
            print("Adding", layername, "to", tiff_multipage)

        elif orafile:
            if i > 0:
                mode = layermode
            else:
                mode = "NORMAL"
            orafile.store_layer(aligned_img, i, layername, layermode=mode)
            print("Adding", layername, "to", orafile.orafilename,
                  "in mode", mode)

        else:
            if type(img) is str:
                fname = os.path.splitext(os.path.basename(img))[0]
                outfname = f"a_{fname}.{ext}"
            else:
                outfname = f"a_img_{i}.{ext}"

            outfname = os.path.join(outdir, outfname)
            aligned_img.save(outfname)

            if os.path.exists(outfname):
                print("Overwriting", outfname, "with", layername)
            else:
                print("Creating", outfname, "with", layername)

    if orafile:
        # save_thumbnail overwrites its input image, but we're done
        # aligned_img so that's okay.
        # orafile.save_thumbnail(aligned_img)
        # finish up
        orafile.finish()
        print("Wrote", orafile.orafilename)


# Multiple color layers? Use just the green layer for alignment.
def singlelayer(img, layer=1):
    if len(img.shape) >= 3:
        return img[:, :, layer]
    return img


def register(rgbimage, baselayer):
    """Align an image of type numpy.ndarray to a base image.
       Input is normally an rgbimage, shape (width, height, 3)
       but can also be monochrome, (width, height, 1).
       Return the realigned image as a numpy.ndarray.
    """
    img2 = singlelayer(rgbimage)

    # Register the two images
    try:
        img_aligned, footprint = astroalign.register(baselayer, img2)
    except astroalign.MaxIterError:
        print(astroalign._find_sources(baselayer).shape, "sources")
        print(astroalign._find_sources(img2).shape, "sources")
        sys.exit(1)

    # Plot the results
    # plot_three(baselayer, img2, img_aligned)

    transf, (pos_img, pos_img_rot) = astroalign.find_transform(baselayer, img2)

    def print_stats():
        print("Rotation: %2d degrees" % (transf.rotation * 180.0 / np.pi))
        print("\nScale factor: %.2f" % transf.scale)
        print("\nTranslation: (x, y) = (%.2f, %.2f)"
              % tuple(transf.translation))
        print("\nTranformation matrix:\n", transf.params)
        print("\nPoint correspondence:")
        for (x1, y1), (x2, y2) in zip(pos_img, pos_img_rot):
            print("(%.2f, %.2f) in source --> (%.2f, %.2f) in target"
                  % (x1, y1, x2, y2))
    # print_stats()

    # Plot correspondences
    # plot_three(baselayer, img2, img_aligned,
    #            pos_img=pos_img, pos_img_rot=pos_img_rot, transf=transf)

    # Align again using the transform.
    # Will use this to align the other channels after using one
    # channel to register the two images.
    # The documentation doesn't mention a footprint being part of the return,
    # but it is.
    # realigned, footprint = astroalign.apply_transform(transf, baselayer, img2)

    # plot_three(baselayer, img2, realigned)

    newshape = rgbimage.shape
    if len(newshape) == 2:
        newshape = rgbimage.shape + (3,)

    # trying https://stackoverflow.com/a/10445502
    rgbArray = np.zeros(newshape, 'uint8')
    for i in range(newshape[-1]):
        layer = singlelayer(rgbimage, i)
        realigned, footprint = astroalign.apply_transform(transf, baselayer,
                                                          layer)

        rgbArray[..., i] = layer

    return rgbArray


def make_test_images():
    np.random.seed(seed=12)

    h, w = img_shape = (200, 200)
    n_stars = 10
    pos_x = np.random.randint(10, w - 10, n_stars)
    pos_y = np.random.randint(10, h - 10, n_stars)
    fluxes = 200.0 + np.random.rand(n_stars) * 300.0

    img = np.zeros(img_shape)
    for x, y, f in zip(pos_x, pos_y, fluxes):
        img[x, y] = f

    # Generate a test image that's rotated and noisy:

    # rotate and make the image one and a half times as big
    from scipy.ndimage import rotate, zoom
    img_rotated = rotate(img, angle=30.0, reshape=False)
    img_rotated = zoom(img_rotated, 1.5, order=2)

    # add a Gaussian PSF response with different seeing for both images
    from scipy.ndimage.filters import gaussian_filter
    img = gaussian_filter(img, sigma=2.0, mode='constant')
    img_rotated = gaussian_filter(img_rotated, sigma=1.5, mode='constant')

    # add some noise
    noise_dc = 5.0
    noise_std = np.sqrt(noise_dc)
    img += np.random.normal(loc=noise_dc, scale=noise_std, size=img.shape)
    img_rotated += np.random.normal(loc=noise_dc, scale=noise_std,
                                    size=img_rotated.shape)

    return img, img_rotated


def read_image(path):
    """If path is a filename, read an image from it into a numpy array.
       If it's already a numpy array, just return it.
    """
    if type(path) is np.ndarray:
        return path
    if type(path) is not str and type(path) is not bytes:
        raise RuntimeError("read_image can't process type " + str(type(path)))

    # First try reading as raw, if rawpy is loaded.
    # rawpy doesn't do a very good job by default, though.
    # Probably it would need some adjustment of parameters below.
    if 'rawpy' in sys.modules:
        try:
            with rawpy.imread(path) as raw:
                # raw.postprocess() -> numpy.ndarray of shape (2856, 4290, 3)
                # Use only the green channel.
                print("Reading", path, "as raw")

                # Relevant parameters:
                # https://letmaik.github.io/rawpy/api/rawpy.Params.html#rawpy.Params
                # no_auto_bright=True
                # auto_bright_thr=(float): ratio of clipped pixels with auto_bright. Default is 0.01 (1%).
                # bright (float) – brightness scaling
                # exp_shift (float) – exposure shift in linear scale. Usable range from 0.25 (2-stop darken) to 8.0 (3-stop lighter).
                # exp_preserve_highlights (float) – preserve highlights when lightening the image with exp_shift. From 0.0 to 1.0 (full preservation).
                # gamma (tuple) – pair (power,slope), default is (2.222, 4.5) for rec. BT.709
                return raw.postprocess(no_auto_bright=False)
        except rawpy._rawpy.LibRawFileUnsupportedError:
            pass

    image = Image.open(path)
    print("Reading", path, "with PIL")
    return np.asarray(image)


colors = ['r', 'g', 'b', 'y', 'cyan', 'w', 'm']

def plot_three(img1, img2, img3, labels=None,
               pos_img=None, pos_img_rot=None, transf=None):

    DEFLABELS = ["Target Image", "Target Image",
                 "Source Image aligned with Target"]

    def smalldim(im):
        return min(im.shape)

    if not labels:
        labels = DEFLABELS

    fig, axes = plt.subplots(2, 2, figsize=(10, 10))
    axes[0, 0].imshow(img1, cmap='gray', interpolation='none', origin='lower')
    axes[0, 0].axis('off')
    axes[0, 0].set_title(labels[0])
    if pos_img is not None:
        circsize = smalldim(img1) / 100
        for (xp, yp), c in zip(pos_img[:len(colors)], colors):
            circ = plt.Circle((xp, yp), circsize,
                              fill=False, edgecolor=c, linewidth=2)
            axes[0, 0].add_patch(circ)

    axes[0, 1].imshow(img2, cmap='gray', interpolation='none',
                      origin='lower')
    axes[0, 1].axis('off')
    axes[0, 1].set_title(labels[1])
    if transf and pos_img_rot is not None:
        circsize = smalldim(img2) / 100
        for (xp, yp), c in zip(pos_img_rot[:len(colors)], colors):
            circ = plt.Circle((xp, yp), circsize * transf.scale,
                              fill=False, edgecolor=c,
                              linewidth=2)
            axes[0, 1].add_patch(circ)

    axes[1, 1].imshow(img3, cmap='gray', interpolation='none',
                      origin='lower')
    axes[1, 1].axis('off')
    axes[1, 1].set_title(labels[2])
    if transf and pos_img_rot is not None:
        circsize = smalldim(img3) / 100
        for (xp, yp), c in zip(pos_img_rot[:len(colors)], colors):
            circ = plt.Circle((xp, yp), circsize * transf.scale,
                              fill=False, edgecolor=c,
                              linewidth=2)
            axes[1, 1].add_patch(circ)

    axes[1, 0].axis('off')

    plt.tight_layout()
    plt.show()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Stack astronomical images")
    parser.add_argument('-t', "--test", dest="test", default=False,
                        action="store_true",
                help="Test mode: generate images instead of reading files")
    parser.add_argument('-d', action="store", dest="dir", default='.',
                        help='Directory to save files (default: .)')
    parser.add_argument('-e', action="store",  dest="ext", default='ora',
                        help='Output image file extension (default: ora)')
    parser.add_argument('-m', action="store",  dest="layermode",
                        default='ADDITION',
                        help='Layer mode if using ora (default: ADDITION)')
    parser.add_argument('-D', action="store",  dest="darkframe",
                        help='Dark frame')
    parser.add_argument('imagefiles', nargs='*', help="2 or more input images")
    args = parser.parse_args(sys.argv[1:])

    if args.test:
        register_all(make_test_images(), outdir=args.dir,
                     ext=args.ext, layermode=args.layermode)
        sys.exit(0)

    if len(args.imagefiles) < 2:
        parser.print_help()
        sys.exit(1)

    register_all(args.imagefiles, outdir=args.dir, ext=args.ext,
                 darkframe=args.darkframe, layermode=args.layermode)
