#!/usr/bin/env python
"""
Created on May 23, 2013, last modified on June 19, 2013
@author: Cheng Chang (cheng.chang.ece@gmail.com)
         Computer Science Group, Computational Science Center
         Brookhaven National Laboratory

This code is for Differential Phase Contrast (DPC) imaging based on Fourier-shift fitting
implementation.

Reference: Yan, H. et al. Quantitative x-ray phase imaging at the nanoscale by multilayer
           Laue lenses. Sci. Rep. 3, 1307; DOI:10.1038/srep01307 (2013).

Test data is available at:
https://docs.google.com/file/d/0B3v6W1bQwN_AdjZwWmE3WTNqVnc/edit?usp=sharing
"""
from __future__ import print_function
import os
import numpy as np
import matplotlib.pyplot as plt
import PIL

from scipy.misc import imsave
from scipy.optimize import minimize
import time
import cStringIO as StringIO
import load_timepix

rss_cache = {}

rss_iters = 0


def get_beta(xdata):
    length = len(xdata)
    try:
        beta = rss_cache[length]
    except Exception:
        # beta = 1j * (np.arange(length) + 1 - (np.floor(length / 2.0) + 1))
        beta = 1j * (np.arange(length) - np.floor(length / 2.0))
        rss_cache[length] = beta

    return beta


def rss(v, xdata, ydata, beta):
    """Function to be minimized in the Nelder Mead algorithm"""
    fitted_curve = xdata * v[0] * np.exp(v[1] * beta)
    return np.sum(np.abs(ydata - fitted_curve) ** 2)


def pil_load(fn):
    im = PIL.Image.open(fn)

    def toarray(im, dtype=np.uint8):
        x_str = im.tostring("raw", im.mode)
        return np.fromstring(x_str, dtype)

    assert im.mode.startswith("I;16")
    if im.mode.endswith("B"):
        x = toarray(im, ">u2")
    else:
        x = toarray(im, "<u2")

    x.shape = im.size[1], im.size[0]
    return x.astype("=u2")


def load_file(fn, roi=None, bad_pixels=[], zip_file=None):
    """
    Load an image file
    """
    if os.path.exists(fn):
        im = load_timepix.load(fn)

    elif zip_file is not None:
        raise NotImplementedError

        # loading from a zip file is just about as fast (when not running in
        # parallel)
        f = zip_file.open(fn)
        stream = StringIO.StringIO()
        stream.write(f.read())
        f.close()

        stream.seek(0)
        im = plt.imread(stream, format="tif")
    else:
        raise Exception("File not found: %s" % fn)

    if bad_pixels is not None:
        for x, y in bad_pixels:
            im[x, y] = 0

    if roi is not None:
        x1, y1, x2, y2 = roi
        im = im[x1 : x2 + 1, y1 : y2 + 1]

    xline = np.sum(im, axis=1)
    yline = np.sum(im, axis=0)

    fx = np.fft.fftshift(np.fft.ifft(xline))
    fy = np.fft.fftshift(np.fft.ifft(yline))
    return im, fx, fy


def xj_test(filename, i, j, roi=None, bad_pixels=[], **kwargs):
    try:
        im, fx, fy = load_file(filename, zip_file=zip_file, roi=roi, bad_pixels=bad_pixels)
    except Exception as ex:
        print("Failed to load file %s: %s" % (filename, ex))
        return 0.0, 0.0, 0.0

    wx, wy = im.shape
    gx = np.sum(im[: wx / 2, :]) - np.sum(im[wx / 2 :, :])
    gy = np.sum(im[:, : wy / 2]) - np.sum(im[:, wy / 2 :])
    return 0, gx, gy


def run_dpc(
    filename,
    i,
    j,
    ref_fx=None,
    ref_fy=None,
    start_point=[1, 0],
    pixel_size=55,
    focus_to_det=1.46,
    dx=0.1,
    dy=0.1,
    energy=19.5,
    zip_file=None,
    roi=None,
    bad_pixels=[],
    max_iters=1000,
    solver="Nelder-Mead",
    invers=False,
):
    """
    All units in micron

    pixel_size
    focus_to_det: focus to detector distance
    dx: scan step size x
    dy: scan step size y
    energy: in keV
    """
    try:
        img, fx, fy = load_file(filename, zip_file=zip_file, roi=roi, bad_pixels=bad_pixels)
    except Exception as ex:
        print("Failed to load file %s: %s" % (filename, ex))
        return 0.0, 0.0, 0.0

    # vx = fmin(rss, start_point, args=(ref_fx, fx, get_beta(ref_fx)),
    #          maxiter=max_iters, maxfun=max_iters, disp=0)
    res = minimize(
        rss,
        start_point,
        args=(ref_fx, fx, get_beta(ref_fx)),
        method=solver,
        tol=1e-4,
        options=dict(maxiter=max_iters),
    )

    vx = res.x
    a = vx[0]
    if invers:
        gx = -vx[1]
    else:
        gx = vx[1]

    # vy = fmin(rss, start_point, args=(ref_fy, fy, get_beta(ref_fy)),
    #          maxiter=max_iters, maxfun=max_iters, disp=0)
    res = minimize(
        rss,
        start_point,
        args=(ref_fy, fy, get_beta(ref_fy)),
        method=solver,
        tol=1e-6,
        options=dict(maxiter=max_iters),
    )

    vy = res.x
    gy = vy[1]

    # print(i, j, vx[0], vx[1], vy[1])
    return a, gx, gy


def recon(gx, gy, dx=0.1, dy=0.1, pad=1, w=1.0):
    """
    Reconstruct the final phase image
    Parameters
    ----------
    gx : 2-D numpy array
        phase gradient along x direction

    gy : 2-D numpy array
        phase gradient along y direction

    dx : float
        scanning step size in x direction (in micro-meter)

    dy : float
        scanning step size in y direction (in micro-meter)

    pad : float
        padding parameter
        default value, pad = 1 --> no padding
                    p p p
        pad = 3 --> p v p
                    p p p

    w : float
        weighting parameter for the phase gradient along x and y direction when
        constructing the final phase image

    Returns
    ----------
    phi : 2-D numpy array
        final phase image

    References
    ----------
    [1] Yan, Hanfei, Yong S. Chu, Jorg Maser, Evgeny Nazaretski, Jungdae Kim,
    Hyon Chol Kang, Jeffrey J. Lombardo, and Wilson KS Chiu, "Quantitative
    x-ray phase imaging at the nanoscale by multilayer Laue lenses," Scientific
    reports 3 (2013).

    """

    rows, cols = gx.shape

    gx_padding = np.zeros((pad * rows, pad * cols), dtype="d")
    gy_padding = np.zeros((pad * rows, pad * cols), dtype="d")

    gx_padding[(pad // 2) * rows : (pad // 2 + 1) * rows, (pad // 2) * cols : (pad // 2 + 1) * cols] = gx
    gy_padding[(pad // 2) * rows : (pad // 2 + 1) * rows, (pad // 2) * cols : (pad // 2 + 1) * cols] = gy

    tx = np.fft.fftshift(np.fft.fft2(gx_padding))
    ty = np.fft.fftshift(np.fft.fft2(gy_padding))

    c = np.zeros((pad * rows, pad * cols), dtype=complex)

    mid_col = pad * cols // 2.0 + 1
    mid_row = pad * rows // 2.0 + 1

    ax = 2 * np.pi * (np.arange(pad * cols) + 1 - mid_col) / (pad * cols * dx)
    ay = 2 * np.pi * (np.arange(pad * rows) + 1 - mid_row) / (pad * rows * dy)

    kappax, kappay = np.meshgrid(ax, ay)

    c = -1j * (kappax * tx + w * kappay * ty)

    c = np.ma.masked_values(c, 0)
    c /= kappax ** 2 + w * kappay ** 2
    c = np.ma.filled(c, 0)

    c = np.fft.ifftshift(c)
    phi_padding = np.fft.ifft2(c)
    phi_padding = -phi_padding.real

    phi = phi_padding[(pad // 2) * rows : (pad // 2 + 1) * rows, (pad // 2) * cols : (pad // 2 + 1) * cols]

    return phi


def main(
    file_format="SOFC/SOFC_%05d.tif",
    dx=0.1,
    dy=0.1,
    ref_image=1,
    zip_file=None,
    rows=121,
    cols=121,
    start_point=[1, 0],
    pixel_size=55,
    focus_to_det=1.46e6,
    energy=19.5,
    pool=None,
    first_image=1,
    x1=None,
    x2=None,
    y1=None,
    y2=None,
    bad_pixels=[],
    solver="Nelder-Mead",
    display_fcn=None,
    invers=False,
):

    print("DPC")
    print("---")
    print("\tFile format: %s" % file_format)
    print("\tdx: %s" % dx)
    print("\tdy: %s" % dy)
    print("\trows: %s" % rows)
    print("\tcols: %s" % cols)
    print("\tstart point: %s" % start_point)
    print("\tpixel size: %s" % pixel_size)
    print("\tfocus to det: %s" % (focus_to_det / 1e6))
    print("\tenergy: %s" % energy)
    print("\tfirst image: %s" % first_image)
    print("\treference image: %s" % ref_image)
    print("\tsolver: %s" % solver)
    print("\tROI: (%s, %s)-(%s, %s)" % (x1, y1, x2, y2))

    t0 = time.time()

    roi = None
    if x1 is not None and x2 is not None:
        if y1 is not None and y2 is not None:
            roi = (x1, y1, x2, y2)

    # read the reference image: only one reference image
    reference, ref_fx, ref_fy = load_file(
        file_format % ref_image, zip_file=zip_file, roi=roi, bad_pixels=bad_pixels
    )

    a = np.zeros((rows, cols), dtype="d")
    gx = np.zeros((rows, cols), dtype="d")
    gy = np.zeros((rows, cols), dtype="d")

    dpc_settings = dict(
        start_point=start_point,
        pixel_size=pixel_size,
        focus_to_det=focus_to_det,
        dx=dx,
        dy=dy,
        energy=energy,
        zip_file=zip_file,
        ref_fx=ref_fx,
        ref_fy=ref_fy,
        roi=roi,
        bad_pixels=bad_pixels,
        solver=solver,
        invers=invers,
    )

    def get_filename(i, j):
        frame_num = first_image + i * cols + j

        # scan 1   9669
        #          12261 images
        #        = 21930
        # scan 2   21950
        # if frame_num >= 21930:
        #    frame_num += 20
        return file_format % frame_num

    # Wavelength in micron
    lambda_ = 12.4e-4 / energy
    if pool is None:
        for i in range(rows):
            trow = time.time()
            print("Row %d" % i, end="")
            rss_iters = 0
            for j in range(cols):
                _a, _gx, _gy = run_dpc(get_filename(i, j), i, j, **dpc_settings)
                a[i, j] = _a
                gx[i, j] = _gx
                gy[i, j] = _gy

            row_elapsed = 1000 * (time.time() - trow)
            print(" elapsed %.3fms" % row_elapsed, end=" ")
            print(" (per frame %.3fms, rss iters %d)" % (row_elapsed / cols, rss_iters))

    else:
        args = [(get_filename(i, j), i, j) for i in range(rows) for j in range(cols)]

        _t0 = time.time()
        try:
            if 1:
                fcn = run_dpc
            else:
                fcn = xj_test

            if display_fcn is not None:
                np.random.shuffle(args)

            results = [pool.apply_async(fcn, arg, kwds=dpc_settings) for arg in args]

            if display_fcn is not None:
                total_results = len(results)
                k = 0
                while k < total_results:
                    k = 0
                    for arg, result in zip(args, results):
                        if result.ready():
                            _a, _gx, _gy = result.get()
                            fn, i, j = arg

                            a[i, j] = _a
                            gx[i, j] = _gx
                            gy[i, j] = _gy
                            k += 1

                    try:
                        gx *= len(ref_fx) * pixel_size / (lambda_ * focus_to_det * 1e6)
                        gy *= len(ref_fy) * pixel_size / (lambda_ * focus_to_det * 1e6)
                        display_fcn(a, gx, gy, None)
                    except Exception as ex:
                        print("Failed to update display: (%s) %s" % (ex.__class__.__name__, ex))

                    time.sleep(1.0)

            pool.close()
            pool.join()
        except KeyboardInterrupt:
            print("Cancelled")
            return

        for arg, result in zip(args, results):
            fn, i, j = arg
            _a, _gx, _gy = result.get()
            a[i, j] = _a
            gx[i, j] = _gx
            gy[i, j] = _gy
            k += 1

        _t1 = time.time()
        elapsed = _t1 - _t0
        print(
            "Multiprocess elapsed=%.3f frames=%d (per frame %.3fms)"
            % (elapsed, rows * cols, 1000 * elapsed / (rows * cols))
        )

    gx *= len(ref_fx) * pixel_size / (lambda_ * focus_to_det * 1e6)
    gy *= len(ref_fy) * pixel_size / (lambda_ * focus_to_det * 1e6)

    dim = len(np.squeeze(gx).shape)

    if dim != 1:
        imsave("a.jpg", a)
        np.savetxt("a.txt", a)
        imsave("gx.jpg", gx)
        np.savetxt("gx.txt", gx)
        imsave("gy.jpg", gy)
        np.savetxt("gy.txt", gy)

        # -------------reconstruct the final phase image using gx and gy--------------------#
        phi = recon(gx, gy, dx, dy)
        imsave("phi.jpg", phi)
        np.savetxt("phi.txt", phi)

        t1 = time.time()
        print("Elapsed", t1 - t0)

        return a, gx, gy, phi

    else:
        """
        #plt.hold(False)
        plt.plot(np.squeeze(a), '-*')
        plt.savefig('a.jpg')
        np.savetxt('a.txt', a)

        plt.plot(np.squeeze(gx), '-*')
        plt.savefig('gx.jpg')
        np.savetxt('gx.txt', gx)

        plt.plot(np.squeeze(gy), '-*')
        plt.savefig('gy.jpg')
        np.savetxt('gy.txt', gy)

        t1 = time.time()
        print('Elapsed', t1 - t0)
        """

        phi = None
        return a, gx, gy, phi

    # plt.imshow(phi, cmap=cm.Greys_r)
    # plt.show()


if __name__ == "__main__":
    zip_file = None  # zipfile.ZipFile('SOFC.zip')
    main(zip_file=zip_file, processes=0, rows=121, cols=121)
