"""The dct module provide helpers functions to work with experimental diffraction contrast tomography data.
"""
import os
import h5py
import numpy as np
from scipy import ndimage
from matplotlib import pyplot as plt, cm
from pymicro.xray.experiment import ForwardSimulation
from pymicro.crystal.lattice import HklPlane
from pymicro.xray.xray_utils import lambda_keV_to_nm, radiograph, radiographs
from pymicro.crystal.microstructure import Grain, Orientation


class DctForwardSimulation(ForwardSimulation):
    """Class to represent a Forward Simulation."""

    def __init__(self, verbose=False):
        super(DctForwardSimulation, self).__init__('dct', verbose=verbose)
        self.hkl_planes = []
        self.check = 1  # grain id to display infos in verbose mode
        self.omegas = None
        self.reflections = []

    def set_hkl_planes(self, hkl_planes):
        self.hkl_planes = hkl_planes

    def set_diffracting_famillies(self, hkl_list):
        """Set the list of diffracting hk planes using a set of families."""
        symmetry = self.exp.get_sample().get_material().get_symmetry()
        hkl_planes = []
        for hkl in hkl_list:
            # here we set include_friedel_pairs to False as we take it into account in the calculation
            planes = HklPlane.get_family(hkl, include_friedel_pairs=True, crystal_structure=symmetry)
            for plane in planes:  # fix the lattice
                plane.set_lattice(self.exp.get_sample().get_material())
            hkl_planes.extend(planes)
        self.set_hkl_planes(hkl_planes)


    def setup(self, omega_step, grain_ids=None):
        """Setup the forward simulation.

        :param float omega_step: the angular integration step (in degrees) use to compute the diffraction comditions.
        :param list grain_ids: a list of grain ids to restrict the forward simulation (use all grains by default).
        """
        assert self.exp.source.min_energy == self.exp.source.max_energy  # monochromatic case
        lambda_keV = self.exp.source.max_energy
        self.omegas = np.linspace(0.0, 360.0, num=int(360.0 / omega_step), endpoint=False)
        self.reflections = []
        for omega in self.omegas:
            self.reflections.append([])
        if grain_ids:
            # make a list of the grains selected for the forward simulation
            grains = [self.exp.sample.microstructure.get_grain(gid) for gid in grain_ids]
        else:
            grains = self.exp.sample.microstructure.grains
        for g in grains:
            for plane in self.hkl_planes: