#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""FDTD cell phantom with tilted and rolled axis of rotation

The *in silico* data set was created with the
:abbr:`FDTD (Finite Difference Time Domain)` software `meep`_. The data
are 2D projections of a 3D refractive index phantom that is rotated
about an axis which is tilted by 0.2 rad (11.5 degrees) with respect to
the imaging plane and rolled by -.42 rad (-24.1 degrees) within the
imaging plane. The data are the same as were used in the previous
example. A brief description of this algorithm is given in
:cite:`Mueller2015tilted`.

The figure shows the 3D reconstruction from
:abbr:`FDTD (Finite Difference Time Domain)` data created by `meep`_
simulations. The known tilted axis of rotation is used in the
reconstruction process.

This example requires Python3 because the data are lzma-compressed.

.. _`meep`: http://ab-initio.mit.edu/wiki/index.php/Meep
"""
from __future__ import division, print_function

import matplotlib.pylab as plt
import numpy as np
from scipy.ndimage import rotate

import odtbrain as odt

from example_helper import get_file, load_tar_lzma_data

lzmafile = get_file("fdtd_3d_sino_A220_R6.500_tiltyz0.2.tar.lzma")

sino, angles, phantom, cfg = load_tar_lzma_data(lzmafile)

# Perform titlt by -.42 rad in detector plane
rotang = -0.42
rotkwargs = {"mode": "constant",
             "order": 2,
             "reshape": False,
             }
for ii in range(len(sino)):
    sino[ii].real = rotate(
        sino[ii].real, np.rad2deg(rotang), cval=1, **rotkwargs)
    sino[ii].imag = rotate(
        sino[ii].imag, np.rad2deg(rotang), cval=0, **rotkwargs)

A = angles.shape[0]

print("Example: Backpropagation from 3d FDTD simulations")
print("Refractive index of medium:", cfg["nm"])
print("Measurement position from object center:", cfg["lD"])
print("Wavelength sampling:", cfg["res"])
print("Axis tilt in y-z direction:", cfg["tilt_yz"])
print("Number of projections:", A)

# Apply the Rytov approximation
sinoRytov = odt.sinogram_as_rytov(sino)

# Determine tilted axis
tilted_axis = [0, np.cos(cfg["tilt_yz"]), np.sin(cfg["tilt_yz"])]
rotmat = np.array([
    [np.cos(rotang), -np.sin(rotang), 0],
    [np.sin(rotang), np.cos(rotang), 0],
    [0, 0, 1],
    ])
tilted_axis = np.dot(rotmat, tilted_axis)

print("Performing tilted backpropagation.")
# Perform tilted backpropagation
f_tilt = odt.backpropagate_3d_tilted(uSin=sinoRytov,
                                     angles=angles,
                                     res=cfg["res"],
                                     nm=cfg["nm"],
                                     lD=cfg["lD"],
                                     tilted_axis=tilted_axis,
                                     )

# compute refractive index n from object function
n_tilt = odt.odt_to_ri(f_tilt, res=cfg["res"], nm=cfg["nm"])

sx, sy, sz = n_tilt.shape
px, py, pz = phantom.shape

sino_phase = np.angle(sino)

# compare phantom and reconstruction in plot
fig, axes = plt.subplots(1, 3, figsize=(8, 2.4))
kwri = {"vmin": n_tilt.real.min(), "vmax": n_tilt.real.max()}
kwph = {"vmin": sino_phase.min(), "vmax": sino_phase.max(),
        "cmap": plt.cm.coolwarm}  # @UndefinedVariable

# Sinogram
axes[0].set_title("phase projection")
phmap = axes[0].imshow(sino_phase[A // 2, :, :], **kwph)
axes[0].set_xlabel("detector x")
axes[0].set_ylabel("detector y")

axes[1].set_title("sinogram slice")
axes[1].imshow(sino_phase[:, :, sino.shape[2] // 2],
               aspect=sino.shape[1] / sino.shape[0], **kwph)
axes[1].set_xlabel("detector y")
axes[1].set_ylabel("angle [rad]")
# set y ticks for sinogram
labels = np.linspace(0, 2 * np.pi, len(axes[1].get_yticks()))
labels = ["{:.2f}".format(i) for i in labels]
axes[1].set_yticks(np.linspace(0, len(angles), len(labels)))
axes[1].set_yticklabels(labels)

axes[2].set_title("tilt correction (nucleolus)")
rimap = axes[2].imshow(n_tilt[int(sx / 2 + 2 * cfg["res"])].real, **kwri)
axes[2].set_xlabel("x")
axes[2].set_ylabel("y")

# color bars
cbkwargs = {"fraction": 0.045,
            "format": "%.3f"}
plt.colorbar(phmap, ax=axes[0], **cbkwargs)
plt.colorbar(phmap, ax=axes[1], **cbkwargs)
plt.colorbar(rimap, ax=axes[2], **cbkwargs)

plt.tight_layout()
plt.show()
