#!/usr/bin python
# -*- coding: utf-8 -*-
# FIXME: -par_plot option fails:
# ValueError: range parameter must be finite
# ('vis_range', [nan, nan])
# ('ticks', [nan, nan])

from __future__ import (print_function)
import os
import sys
import argparse
import matplotlib
matplotlib.use('Agg')
label_size = 6
matplotlib.rcParams['xtick.labelsize'] = label_size
matplotlib.rcParams['ytick.labelsize'] = label_size
path = os.path.normpath(os.path.join(os.path.dirname(sys.argv[0]), '..'))
sys.path.insert(0, path)
import glob
import numpy as np
from vlbi_errors.uv_data import UVData
from vlbi_errors.spydiff import (import_difmap_model, modelfit_difmap)
from vlbi_errors.bootstrap import CleanBootstrap
from vlbi_errors.model import Model
from vlbi_errors.utils import hdi_of_mcmc

try:
    import corner as triangle
except ImportError:
    triangle = None

outname = 'boot_uv'
errors_fname = 'bootstrap_errors.txt'

if 'DIFMAP_LOGIN' in os.environ:
    del os.environ['DIFMAP_LOGIN']

def xy_2_rtheta(params):
    flux, x, y = params[:3]
    r = np.sqrt(x ** 2 + y ** 2)
    theta = np.rad2deg(np.arctan(x / y))
    result = [flux, r, theta]
    try:
        result.extend(params[3:])
    except IndexError:
        pass
    return result

def analyze_bootstrap_samples(dfm_model_fname, booted_mdl_paths,
                              dfm_model_dir=None, plot_comps=None,
                              plot_file=None, txt_file=None, cred_mass=0.68,
                              coordinates='xy'):
    """
    Plot bootstrap distribution of model component parameters.

    :param dfm_model_fname:
        File name of original difmap model.
    :param booted_mdl_paths:
        Iterable of paths to bootstrapped difmap models.
    :param dfm_model_dir: (optional)
        Directory with original difmap model. If ``None`` then CWD. (default:
        ``None``)
    :param plot_comps: (optional)
        Iterable of components number to plot on same plot. If ``None`` then
        plot parameter distributions of all components.
    :param plot_file: (optional)
        File to save picture. If ``None`` then don't save picture. (default:
        ``None``)
    :param txt_file: (optional)
        File to save credible intervals for parameters. If ``None`` then don't
        save credible intervals. (default: ``None``)
    :param cred_mass: (optional)
        Value of credible interval mass. Float in range (0., 1.). (default:
        ``0.68``)
    :param coordinates: (optional)
        Type of coordinates to use. ``xy`` or ``rtheta``. (default: ``xy``)
    """
    n_boot = len(booted_mdl_paths)
    # Get params of initial model used for bootstrap
    comps_orig = import_difmap_model(dfm_model_fname, dfm_model_dir)
    comps_params0 = {i: [] for i in range(len(comps_orig))}
    for i, comp in enumerate(comps_orig):
        # FIXME: Move (x, y) <-> (r, theta) mapping to ``Component``
        if coordinates == 'xy':
            params = comp.p
        elif coordinates == 'rtheta':
            params = xy_2_rtheta(comp.p)
        else:
            raise Exception
        comps_params0[i].extend(list(params))

    # Load bootstrap models
    comps_params = {i: [] for i in range(len(comps_orig))}
    for booted_mdl_path in booted_mdl_paths:
        path, booted_mdl_file = os.path.split(booted_mdl_path)
        comps = import_difmap_model(booted_mdl_file, path)
        for i, comp in enumerate(comps):
            # FIXME: Move (x, y) <-> (r, theta) mapping to ``Component``
            if coordinates == 'xy':
                params = comp.p
            elif coordinates == 'rtheta':
                params = xy_2_rtheta(comp.p)
            else:
                raise Exception
            comps_params[i].extend(list(params))

    comps_to_plot = [comps_orig[k] for k in plot_comps]
    # (#boot, #parameters)
    boot_data = np.hstack(np.array(comps_params[i]).reshape((n_boot,
                                                             comps_orig[i].size)) for
                          i in plot_comps)

    # Optionally plot
    if plot_file:
        if triangle:
            lens = list(np.cumsum([comp.size for comp in comps_orig]))
            lens.insert(0, 0)

            labels = list()
            for comp in comps_to_plot:
                for lab in np.array(comp._parnames)[~comp._fixed]:
                    # FIXME: Move (x, y) <-> (r, theta) mapping to ``Component``
                    if coordinates == 'rtheta':
                        if lab == 'x':
                            lab = 'r'
                        if lab == 'y':
                            lab = 'theta'
                    elif coordinates == 'xy':
                        pass
                    else:
                        raise Exception
                    labels.append(lab)

            try:
                n = sum([c.size for c in comps_to_plot])
                figure, axes = matplotlib.pyplot.subplots(nrows=n, ncols=n)
                figure.set_size_inches(19.5, 19.5)
                triangle.corner(boot_data, labels=labels, plot_contours=False,
                                truths=np.hstack([comps_params0[i] for i in
                                                  plot_comps]),
                                title_kwargs={"fontsize": 6},
                                label_kwargs={"fontsize": 6},
                                quantiles=[0.16, 0.5, 0.84], fig=figure,
                                use_math_text=True, show_titles=True,
                                title_fmt=".3f")
                figure.gca().annotate("Components {}".format(plot_comps),
                                      xy=(0.5, 1.0),
                                      xycoords="figure fraction",
                                      xytext=(0, -5),
                                      textcoords="offset points", ha="center",
                                      va="top")
                figure.savefig(plot_file, bbox_inches='tight', dpi=300)
            except ValueError:
                print("Failed to plot... ValueError")
        else:
            print("Install ``corner`` for corner-plots")

    if txt_file:
        # Print credible intervals
        fn = open(txt_file, 'w')
        fn.write("# parameter original.value low.boot high.boot mean.boot"
                 " median.boot (mean-low).boot (high-mean).boot\n")
        recorded = 0
        for i in plot_comps:
            comp = comps_orig[i]
            for j in range(comp.size):
                low, high, mean, median = hdi_of_mcmc(boot_data[:, recorded+j],
                                                      cred_mass=cred_mass,
                                                      return_mean_median=True)
                # FIXME: Move (x, y) <-> (r, theta) mapping to ``Component``
                parnames = comp._parnames
                if coordinates == 'xy':
                    params = comp.p
                elif coordinates == 'rtheta':
                    params = xy_2_rtheta(comp.p)
                    parnames[1] = 'r'
                    parnames[2] = 'theta'
                else:
                    raise Exception
                fn.write("{:<4} {:.4f} {:.4f} {:.4f} {:.4f} {:.4f} {:.4f}"
                         " {:.4f}".format(parnames[j], params[j], low,
                                          high, mean, median, abs(median - low),
                                          abs(high - median)))
                fn.write("\n")
            recorded += (j + 1)
        fn.close()

if __name__ == "__main__":
    parser = \
        argparse.ArgumentParser(description='Bootstrap Difmap models.\n'
                                            ' Required '
                                            'arguments are ``dfm_model_path`` &'
                                            ' one of the ``-uv_fits_path`` or'
                                            ' ``-booted_mdl_path``. If'
                                            ' ``-booted_mdl_path`` is used then'
                                            ' options ``-parametric``,'
                                            ' ``-n_boot``, ``-n_iter``,'
                                            ' ``-res_plot`` are not used.')

    parser.add_argument('dfm_model_path', type=str, metavar='dfm_model_path',
                        help='Path to Difmap-format file with model.')
    parser.add_argument('-uv_fits_path', action='store', nargs='?', type=str,
                        metavar='PATH TO UV-DATA FITS FILE', default=None,
                        help='Path to FITS-file with self-calibrated UV-data.')
    parser.add_argument('-booted_mdl_card', action='store', nargs='?', type=str,
                        default=None,
                        help='Wildcard to find bootstrapped model'
                             ' files', metavar='WILCARD WITH FULL PATH')
    parser.add_argument('-n_boot', action='store', nargs='?', default=100,
                        type=int, help='Number of bootstrap realizations.'
                                       ' Default value = 100',
                        metavar='INT')
    parser.add_argument('-n_iter', action='store', nargs='?', default=50,
                        type=int, help='Number of iterations in difmap internal'
                                       ' fitting. Default is 50.',
                        metavar='INT')
    parser.add_argument('-cred_value', action='store', nargs='?', default=0.68,
                        type=float, help='Credible interval specification.'
                                         ' Float from (0, 1) interval. Default'
                                         ' is 0.68.',
                        metavar='FLOAT FROM (0, 1)')
    parser.add_argument('-out_dir', action='store', nargs='?',
                        default=os.getcwd(), type=str, help='Directory to store'
                                                     ' bootstrap files, models'
                                                     ' & results.',
                        metavar='DIRECTORY')
    parser.add_argument('-errors_file', action='store', nargs='?',
                        default='bootstrap_errors.txt', type=str,
                        help='File name to store bootstrap errors. Default is'
                             '`bootstrap_errors.txt`.',
                        metavar='FILE NAME')
    parser.add_argument('-res_plot', action='store', nargs='?', default=None,
                        type=str, help='File name to store IF-averages'
                                       ' residuals of Stokes I real & imag part'
                                       ' plot in output directory.',
                        metavar='FILE NAME')
    parser.add_argument('-amplitude_scale_sigma', action='store', nargs='?',
                        default=None, type=float, help='Sigma of amplitude'
                                                       ' scale. In fractions of'
                                                       ' amplitude.',
                        metavar='FLOAT nearly 0.05-0.1')
    parser.add_argument('-res_plot_full', action='store', nargs='?', default=None,
                        type=str, help='File name to store residuals of Stokes '
                                       'RR & LL real & imag part'
                                       ' plot in output directory.',
                        metavar='FILE NAME')
    parser.add_argument('-par_plot', action='store', nargs='?', default=None,
                        type=str, help='File name to store parameters plot in'
                                       ' output directory.',
                        metavar='FILE NAME')
    parser.add_argument('-plot_comps', action='store', nargs='+', default=None,
                        type=str, help='Components numbers to plot.',
                        metavar='COMPONENT #')
    parser.add_argument('-txt_comps', action='store', nargs='*', default=None,
                        type=str, help='Components numbers to output'
                                       ' parameters in a text file.',
                        metavar='COMPONENT #')
    parser.add_argument('-parametric', action='store_true', dest='parametric',
                        default=False,
                        help='Use parametric bootstrap instead of'
                             ' nonparametric (nonparametric is the default).')
    parser.add_argument('-recenter', action='store_true', dest='recenter',
                        default=False,
                        help='Recenter residuals on each baseline.')
    parser.add_argument('-clean_after', action='store_true', dest='clean_after',
                        default=False,
                        help='Remove bootstrapped data & model files in the'
                             ' end.')
    parser.add_argument('-bic', action='store_true', dest='bic',
                        default=False,
                        help='Calculate BIC criterion value for original model'
                             ' and bootstrapped samples.')
    parser.add_argument('-rtheta', action='store_true', dest='use_rtheta',
                        default=False,
                        help='Use `r-theta` coordinates instead of `xy`.')
    parser.add_argument('-split_scans', action='store_true', dest='split_scans',
                        default=False, help='Resample each scan individually?')

    args = parser.parse_args()

    data_dir = args.out_dir
    if not os.path.exists(data_dir):
        os.makedirs(data_dir)
    print("Data directory: {}".format(data_dir))

    cred_value = args.cred_value
    uv_fits_path = args.uv_fits_path
    booted_mdl_card = args.booted_mdl_card
    dfm_model_path = args.dfm_model_path
    n_boot = args.n_boot
    niter = args.n_iter
    nonparametric = not args.parametric
    errors_fname = args.errors_file
    par_plot = args.par_plot
    recenter = args.recenter
    plot_comps = args.plot_comps
    txt_comps = args.txt_comps
    split_scans = args.split_scans
    amplitude_scale_sigma = args.amplitude_scale_sigma

    bic = args.bic
    if args.use_rtheta:
        coordinates = 'rtheta'
    else:
        coordinates = 'xy'

    if par_plot and not plot_comps:
        raise Exception("Use -plot_comps argument to specify # of components"
                        " to plot")

    dfm_model_dir, dfm_model_fname = os.path.split(dfm_model_path)
    try:
        comps = import_difmap_model(dfm_model_fname, dfm_model_dir)
    except ValueError:
        print("Problem importing difmap model...")
        sys.exit(1)

    # Check that component numbers in input are among model components
    if plot_comps:
        for c in plot_comps:
            if int(c) not in range(len(comps)):
                raise Exception("No such component {} in current"
                                " model!".format(c))
    if not txt_comps:
        txt_comps = range(len(comps))
    else:
        txt_comps = [int(k) for k in txt_comps]
    for c in txt_comps:
        if int(c) not in range(len(comps)):
            raise Exception("No such component {} in current model!".format(c))

    if uv_fits_path:
        print("Bootstrapping uv-data")
        uv_fits_dir, uv_fits_fname = os.path.split(uv_fits_path)
        boot_type_dict = {True: "non-parametric", False: "parametric"}
        print("==================================")
        print("Bootstrap uv-data: {}".format(uv_fits_fname))
        print("With model: {}".format(dfm_model_fname))
        print("Using {} bootstrap".format(boot_type_dict[nonparametric]))
        if not nonparametric:
            if recenter:
                print("Recentering KDE-fitted residuals")
            else:
                print("Using fitted KDE Model to generate resamples")
        print("Using {} bootstrap replications".format(n_boot))
        if amplitude_scale_sigma:
            print("Using amplitude sigma for scale"
                  " = {}".format(amplitude_scale_sigma))
        print("Using {} fitting iterations".format(niter))
        print("Finding {}-confidence regions".format(cred_value))
        print("Using directory {} for storing output".format(data_dir))
        txt_save_dict = {'None': "all"}
        try:
            print("Saving errors of {} components to file"
                  " {}".format(txt_save_dict[str(txt_comps)], errors_fname))
        except KeyError:
            print("Saving errors of {} components to file"
                  " {}".format(txt_comps, errors_fname))

        if split_scans:
            print("Resampling each scan individually")
        if par_plot:
            print("Saving components {} parameters distributions plot to file"
                  " {}".format(plot_comps, par_plot))
        if args.res_plot:
            print("Saving residuals I plot to file {}".format(args.res_plot))
        if args.res_plot_full:
            print("Saving residulas RR & LL plots to files"
                  " {}*".format(args.res_plot_full))
        print("==================================")

        uvdata = UVData(uv_fits_path)
        model = Model(stokes='I')
        model.add_components(*comps)

        if bic:
            bic_orig = model.bic(uvdata)
            bic_booted = list()

        try:
            boot = CleanBootstrap([model], uvdata,
                                  sigma_ampl_scale=amplitude_scale_sigma)
        # If uv-data contains only one Stokes parameter (e.g. `0838+133`)
        except IndexError:
            print("Problem in bootstrapping data...")
            sys.exit(1)
        # FIXME: Broken - ValueError
        if args.res_plot:
            print("Plotting histograms of I residuals...")
            boot.plot_residuals(args.res_plot)

        curdir = os.getcwd()
        os.chdir(data_dir)
        boot.run(n=n_boot, nonparametric=nonparametric, outname=[outname,
                                                                 '.fits'],
                 recenter=recenter, use_kde=True, use_v=False)
        if args.res_plot_full:
            print("Plotting histograms of RR & LL residuals...")
            boot.plot_residuals_trio(args.res_plot_full, split_scans,
                                     stokes=['RR', 'LL'])

        os.chdir(curdir)

        booted_uv_paths = sorted(glob.glob(os.path.join(data_dir,
                                                        outname + "*")))
        booted_mdl_paths = list()
        # Modelfit bootstrapped uvdata
        for booted_uv_path in booted_uv_paths:
            path, booted_uv_file = os.path.split(booted_uv_path)
            i = booted_uv_file.split('_')[-1].split('.')[0]
            out_fname = dfm_model_fname + '_' + i
            modelfit_difmap(booted_uv_file, dfm_model_fname, out_fname,
                            path=path, mdl_path=dfm_model_dir,
                            out_path=data_dir, niter=niter)
            booted_mdl_paths.append(os.path.join(data_dir, out_fname))
            if bic:
                uvdata_ = UVData(booted_uv_path)
                model_ = Model(stokes='I')
                comps_ = import_difmap_model(out_fname, data_dir)
                model_.add_components(comps_)
                bic_booted.append(model_.bic(uvdata_))

        if bic:
            low_, high_, mean_, median_ = hdi_of_mcmc(bic_booted,
                                                      cred_mass=0.68,
                                                      return_mean_median=True)
            print("Model BIC with bootstrapped 68% interval = {:.2f}"
                  " -{:.2f} +{:.2f}".format(bic_orig, abs(mean_ - low_),
                                            abs(high_ - mean_)))

    elif booted_mdl_card:
        print("Using already bootstrapped uv-data")
        booted_mdl_paths = glob.glob(booted_mdl_card)
        n_boot = len(booted_mdl_paths)
        print("==================================")
        print("With {} bootstrap replications".format(n_boot))
        print("Finding {}-confidence regions".format(cred_value))
        print("Using directory {} for storing output".format(data_dir))
        txt_save_dict = {'None': "all"}
        try:
            print("Saving errors of {} components to file"
                  " {}".format(txt_save_dict[str(txt_comps)], errors_fname))
        except KeyError:
            print("Saving errors of {} components to file"
                  " {}".format(txt_comps, errors_fname))
        if par_plot:
            print("Saving components {} parameters distributions plot to file"
                  " {}".format(plot_comps, par_plot))
        print("==================================")

    else:
        raise Exception("Use -uv_fits_path or -booted_mdl_card to create/get"
                        " bootstrapped models.")

    # Optionally plot component parameters
    if par_plot:
        plot_comps = [int(k) for k in plot_comps]
        analyze_bootstrap_samples(dfm_model_fname, booted_mdl_paths,
                                  dfm_model_dir=dfm_model_dir,
                                  plot_comps=plot_comps,
                                  plot_file=os.path.join(data_dir, par_plot),
                                  cred_mass=cred_value,
                                  coordinates=coordinates)
    analyze_bootstrap_samples(dfm_model_fname, booted_mdl_paths,
                              dfm_model_dir=dfm_model_dir,
                              plot_comps=txt_comps,
                              txt_file=os.path.join(data_dir, errors_fname),
                              cred_mass=cred_value, coordinates=coordinates)

    if args.clean_after:
        for rmfile in booted_uv_paths:
            os.unlink(rmfile)
        for rmfile in booted_mdl_paths:
            os.unlink(rmfile)


"""
@author: reubendo
"""

# -*- coding: utf-8 -*-
from __future__ import absolute_import, print_function, division

import numpy as np
import tensorflow as tf

from niftynet.layer.base_layer import TrainableLayer
from niftynet.layer.convolution import ConvolutionalLayer
from niftynet.layer.deconvolution import DeconvolutionalLayer
from niftynet.layer.fully_connected import FullyConnectedLayer
from niftynet.layer.activation import ActiLayer
from niftynet.layer.bn import InstanceNormLayer
from niftynet.layer.convolution import ConvLayer
from niftynet.layer.elementwise import ElementwiseLayer
from niftynet.layer.downsample import DownSampleLayer as Pooling
from niftynet.layer.linear_resize import LinearResizeLayer

MODALITIES = ['T1', 'T1c', 'T2', 'Flair', 'seg']

HIDDEN_SPACE = 512

NB_CONV = 8
tf.set_random_seed(1)
class U_HeMIS(TrainableLayer):
    """
    Implementation of U-HeMIS introduced [1] mixing HeMIS [2] and a U-Net architecture [3]
    [1] Dorent, et al. "Hetero-Modal Variational Encoder-Decoder for
        Joint Modality Completion and Segmentation".
        MICCAI 2019.
    [2] Havaei, et al. "HeMIS: Hetero-Modal Image Segmentation".
        MICCAI 2016. https://arxiv.org/abs/1607.05194
    [3] Ronneberger, et al. "U-Net: Convolutional Networks for Biomedical Image Segmentation".
        MICCAI 2015. https://arxiv.org/abs/1505.04597
    """

    def __init__(self,
                num_classes,
                w_initializer=None,
                w_regularizer=None,
                b_initializer=None,
                b_regularizer=None,
                acti_func='leakyrelu',
                name='VAE'):

        super(U_HeMIS, self).__init__(name=name)

        self.initializers = {'w': w_initializer, 'b': b_initializer}
        self.regularizers = {'w': w_regularizer, 'b': b_regularizer}
        self.num_classes = num_classes

    def layer_op(self, images, choices, is_training=True, is_inference=False, **unused_kwargs):

        encoder = ConvEncoder(w_initializer=self.initializers['w'], w_regularizer=self.regularizers['w'], b_initializer=self.initializers['b'],b_regularizer=self.regularizers['b'])

        abstraction_op = HeMISAbstractionBlock()

        img_decoder = ConvDecoderImg(w_initializer=self.initializers['w'], w_regularizer=self.regularizers['w'], b_initializer=self.initializers['b'],b_regularizer=self.regularizers['b'])

        mod_img = MODALITIES[:4]

        # Encode the input
        list_skips = encoder(images)

        # Sample from the posterior distribution P(latent variables|input)
        skip_flow = []
        for k in range(len(list_skips)):
            sample = abstraction_op(list_skips[k],  choices,  is_inference)
            skip_flow.append(sample)

        img_output = img_decoder(skip_flow)

        return img_output

class ConvEncoder(TrainableLayer):
    """
        Each modality are encoded indepedently.
    """

    def __init__(self,
                 w_initializer=None,
                 w_regularizer=None,
                 b_initializer=None,
                 b_regularizer=None,
                 name='ConvEncoder'):

        super(ConvEncoder, self).__init__(name=name)

        self.initializers = {'w': w_initializer, 'b': b_initializer}
        self.regularizers = {'w': w_regularizer, 'b': b_regularizer}

        self.ini_f = NB_CONV
        self.layers = [
            {'name': 'conv_0', 'n_features': self.ini_f, 'kernel_size': (1,1,1)},
            {'name': 'block_1', 'n_features': self.ini_f, 'kernels': ((3,3,3), (3,3,3)), 'downsampling':True},
            {'name': 'block_2', 'n_features': 2*self.ini_f, 'kernels': ((3,3,3), (3,3,3)), 'downsampling':True},
            {'name': 'block_3', 'n_features': 4*self.ini_f, 'kernels': ((3,3,3), (3,3,3)), 'downsampling':True},
            {'name': 'block_4', 'n_features': 8*self.ini_f, 'kernels': ((3,3,3), (3,3,3)), 'downsampling':False}]

        self.skip_ind = [1, 3, 5, 7]
        self.hidden = [self.layers[k]['n_features'] for k in range(1,len(self.layers))]
        self.hidden = [int(k/2) for k in self.hidden]

    def layer_op(self, images):
        # Define the encoding convolutional layers
        def clip(input):
            # This is for clipping logvars,
            # so that variances = exp(logvars) behaves well
            output = tf.maximum(input, -50)
            output = tf.minimum(output, 50)
            return output

        layer_instances = [] #list layers
        means = dict()
        logvars = dict()

        pooling_params = {'kernel_size': 2, 'stride': 2}

        list_skip_flow = [[] for k in range(len(self.skip_ind))]

        layer_fc_mod = dict()
        layer_cnn_mod = dict()

        for mod in MODALITIES[:4]:
            layer_cnn_mod[mod] = []
            layer_fc_mod[mod] = []

            params = self.layers[0]
            first_conv_layer = ConvolutionalLayer(
                n_output_chns=params['n_features'],
                kernel_size=params['kernel_size'],
                acti_func='leakyrelu',
                with_bn=False,
                w_initializer=self.initializers['w'],
                w_regularizer=self.regularizers['w'],
                name='%s_%s' % (params['name'],mod))

            layer_instances.append(first_conv_layer)
            layer_cnn_mod[mod].append(first_conv_layer)

            for i in range(1,len(self.layers)):

                params = self.layers[i]
                res_block = ResBlock(
                    n_output_chns=params['n_features'],
                    kernels=params['kernels'],
                    acti_func='leakyrelu',
                    encoding=True,
                    w_initializer=self.initializers['w'],
                    w_regularizer=self.regularizers['w'],
                    name='%s_%s' % (params['name'],mod))
                layer_instances.append(res_block)
                layer_cnn_mod[mod].append(res_block)

                if params['downsampling']:
                    downsampler = Pooling(func='MAX', kernel_size=2, stride=2,)

                    layer_instances.append(downsampler)
                    layer_cnn_mod[mod].append(downsampler)

        for mod in MODALITIES[:4]:
            flow_mod = images[mod]
            print(flow_mod)

            for ind, cnn_mod in enumerate(layer_cnn_mod[mod]):

                flow_mod = cnn_mod(flow_mod)
                layer_cnn_mod[mod][ind] = cnn_mod
                if ind in self.skip_ind:
                    pos = self.skip_ind.index(ind)
                    list_skip_flow[pos].append(flow_mod)
                print(flow_mod)
            print('list_flow')

        print(list_skip_flow)

        output = list_skip_flow

        if True:
            self._print(layer_instances)
            return output
        return output

    def _print(self, list_of_layers):
        for op in list_of_layers:
            print(op)

class HeMISAbstractionBlock(TrainableLayer):
    def __init__(self,
                 pooling_type='average',
                 name='HeMISAbstractionBlock'):

        super(HeMISAbstractionBlock, self).__init__(name=name)

        self.pooling_type = pooling_type
        self.name = name

    def layer_op(self, input_tensor, choices, is_training):
        """
        Written by Thomas Varsavsky.

        Function will drop all zero columns and compute E[C] and Var[C]
        :param backend_output: backend_output
        :return: 1xC tensor where C is the number of features.
        """
        # Omit zero columns from average
        # intermediate_tensor = tf.reduce_sum(tf.abs(input_tensor), 0)
        # zero_vector = tf.zeros(shape=(1, 1), dtype=tf.float32)
        # bool_mask = tf.not_equal(intermediate_tensor, zero_vector)
        # omit_zero_columns = tf.boolean_mask(input_tensor, bool_mask)
        # Compute E[C]
        input_tensor = tf.boolean_mask(input_tensor, choices)
        average_over_modalities, variance_between_modalities = tf.nn.moments(input_tensor, axes=[0])
        abstraction_output = tf.concat([average_over_modalities, variance_between_modalities], axis=-1)
        return abstraction_output

class ConvDecoderImg(TrainableLayer):
    """
    Each modality are then decoded using the average of the skip-connections across
    the available modalities.
    """

    def __init__(self,
                 w_initializer=None,
                 w_regularizer=None,
                 b_initializer=None,
                 b_regularizer=None,
                 name='ConvDecoderImg'):

        super(ConvDecoderImg, self).__init__(name=name)

        self.initializers = {'w': w_initializer, 'b': b_initializer}
        self.regularizers = {'w': w_regularizer, 'b': b_regularizer}

        self.ini_f = NB_CONV
        self.layers = [
            {'name': 'block_1', 'n_features': 4*self.ini_f, 'kernels': ((3,3,3), (3,3,3))},
            {'name': 'block_2', 'n_features': 2*self.ini_f, 'kernels': ((3,3,3), (3,3,3))},
            {'name': 'block_3', 'n_features': self.ini_f, 'kernels': ((3,3,3), (3,3,3))}]

    def layer_op(self, list_skips):

        # Define the decoding convolutional layers
        layer_instances = [] #list layers
        layer_mod = dict()
        decoders_fc = dict()
        flow = dict()

        list_skips = list_skips[::-1]

        for mod in ['seg']:
            layer_mod[mod] = []

            flow_mod = list_skips[0]

            if mod =='seg':
                double = True
                n_output = 4
            else:
                double = True
                n_output = 1

            for i in range(len(self.layers)):

                params = self.layers[i]

                flow_mod = LinearResizeLayer(list_skips[i+1].shape.as_list()[1:-1])(flow_mod)

                print(mod)
                print(flow_mod)
                print('added with ')
                print(list_skips[i+1])
                flow_mod = ElementwiseLayer('CONCAT')(flow_mod, list_skips[i+1])
                print(flow_mod)

                res_block = ResBlock(
                    n_output_chns=params['n_features'],
                    kernels=params['kernels'],
                    acti_func='leakyrelu',
                    encoding=False,
                    double_n = double,
                    w_initializer=self.initializers['w'],
                    w_regularizer=self.regularizers['w'],
                    name='%s_%s' % (params['name'],mod))
                layer_instances.append(res_block)
                layer_mod[mod].append(res_block)
                flow_mod = res_block(flow_mod)

            last_conv = ConvolutionalLayer(
                n_output_chns=n_output,
                kernel_size=(1,1,1),
                with_bn=False,
                acti_func=None,
                w_initializer=self.initializers['w'],
                w_regularizer=self.regularizers['w'],
                name='final_conv_seg')
            layer_instances.append(last_conv)
            flow_mod = last_conv(flow_mod)
            flow[mod] = flow_mod

        if True:
            self._print(layer_instances)
            return flow['seg']
        return flow['seg']

    def _print(self, list_of_layers):
        for op in list_of_layers:
            print(op)

class ResBlock(TrainableLayer):
    """
    This class define a high-resolution block with residual connections
    kernels

        - specify kernel sizes of each convolutional layer
        - e.g.: kernels=(5, 5, 5) indicate three conv layers of kernel_size 5

    with_res

        - whether to add residual connections to bypass the conv layers
    """

    def __init__(self,
                 n_output_chns,
                 kernels=((3,3,3), (3,3,3)),
                 acti_func='leakyrelu',
                 encoding=False,
                 double_n = True,
                 w_initializer=None,
                 w_regularizer=None,
                 with_res=True,
                 name='ResBlock',
                 stride=1):

        super(ResBlock, self).__init__(name=name)

        self.n_output_chns = n_output_chns
        if hasattr(kernels, "__iter__"):  # a list of layer kernel_sizes
            self.kernels = kernels
        else:  # is a single number (indicating single layer)
            self.kernels = [kernels]
        self.acti_func = acti_func
        self.with_res = with_res

        self.initializers = {'w': w_initializer}
        self.regularizers = {'w': w_regularizer}
        self.stride = stride
        self.encoding = encoding
        self.double_n = double_n
        self.kernels = self.kernels if double_n else [self.kernels[0]]

    def layer_op(self, input_tensor):
        output_tensor = input_tensor
        for (i, k) in enumerate(self.kernels):
            # create parameterised layers
            if self.encoding:
                if i==0:
                    nb_channels = self.n_output_chns
                elif i==1:
                    nb_channels = int(self.n_output_chns/2)
            else:
                if self.double_n:
                    if i==0:
                        nb_channels = self.n_output_chns
                    elif i==1:
                        nb_channels = int(self.n_output_chns/2)
                else:
                    nb_channels = int(self.n_output_chns/2)

            in_op = InstanceNormLayer(name='in_{}'.format(i))
            acti_op = ActiLayer(func=self.acti_func,
                                regularizer=self.regularizers['w'],
                                name='acti_{}'.format(i))
            conv_op = ConvLayer(n_output_chns=nb_channels,
                                kernel_size=k,
                                stride=self.stride,
                                w_initializer=self.initializers['w'],
                                w_regularizer=self.regularizers['w'],
                                name='conv_{}'.format(i))
            # connect layers
            output_tensor = in_op(output_tensor)
            output_tensor = acti_op(output_tensor)
            output_tensor = conv_op(output_tensor)
        # make residual connections
        # if self.with_res:
        #     output_tensor = ElementwiseLayer('SUM')(output_tensor, input_tensor)
        return output_tensor


from collections import defaultdict
from collections.abc import MutableSequence, Iterable
import io

import numpy as np
from numpy.polynomial import Polynomial
import pandas as pd

from .data import NEUTRON_MASS
from .endf import get_head_record, get_cont_record, get_tab1_record, get_list_record
try:
    from .reconstruct import wave_number, penetration_shift, reconstruct_mlbw, \
        reconstruct_slbw, reconstruct_rm
    _reconstruct = True
except ImportError:
    _reconstruct = False
import openmc.checkvalue as cv

class Resonances(object):
    """Resolved and unresolved resonance data

    Parameters
    ----------
    ranges : list of openmc.data.ResonanceRange
        Distinct energy ranges for resonance data

    Attributes
    ----------
    ranges : list of openmc.data.ResonanceRange
        Distinct energy ranges for resonance data
    resolved : openmc.data.ResonanceRange or None
        Resolved resonance range
    unresolved : openmc.data.Unresolved or None
        Unresolved resonance range

    """

    def __init__(self, ranges):
        self.ranges = ranges

    def __iter__(self):
        for r in self.ranges:
            yield r

    @property
    def ranges(self):
        return self._ranges

    @property
    def resolved(self):
        resolved_ranges = [r for r in self.ranges
                           if not isinstance(r, Unresolved)]
        if len(resolved_ranges) > 1:
            raise ValueError('More than one resolved range present')
        elif len(resolved_ranges) == 0:
            return None
        else:
            return resolved_ranges[0]

    @property
    def unresolved(self):
        for r in self.ranges:
            if isinstance(r, Unresolved):
                return r
        else:
            return None

    @ranges.setter
    def ranges(self, ranges):
        cv.check_type('resonance ranges', ranges, MutableSequence)
        self._ranges = cv.CheckedList(ResonanceRange, 'resonance ranges',
                                      ranges)

    @classmethod
    def from_endf(cls, ev):
        """Generate resonance data from an ENDF evaluation.

        Parameters
        ----------
        ev : openmc.data.endf.Evaluation
            ENDF evaluation

        Returns
        -------
        openmc.data.Resonances
            Resonance data

        """
        file_obj = io.StringIO(ev.section[2, 151])

        # Determine whether discrete or continuous representation
        items = get_head_record(file_obj)
        n_isotope = items[4]  # Number of isotopes

        ranges = []
        for iso in range(n_isotope):
            items = get_cont_record(file_obj)
            abundance = items[1]
            fission_widths = (items[3] == 1)  # fission widths are given?
            n_ranges = items[4]  # number of resonance energy ranges

            for j in range(n_ranges):
                items = get_cont_record(file_obj)
                resonance_flag = items[2]  # flag for resolved (1)/unresolved (2)
                formalism = items[3]  # resonance formalism

                if resonance_flag in (0, 1):
                    # resolved resonance region
                    erange = _FORMALISMS[formalism].from_endf(ev, file_obj, items)

                elif resonance_flag == 2:
                    # unresolved resonance region
                    erange = Unresolved.from_endf(file_obj, items, fission_widths)

                # erange.material = self
                ranges.append(erange)

        return cls(ranges)

class ResonanceRange(object):
    """Resolved resonance range

    Parameters
    ----------
    target_spin : float
        Intrinsic spin, :math:`I`, of the target nuclide
    energy_min : float
        Minimum energy of the resolved resonance range in eV
    energy_max : float
        Maximum energy of the resolved resonance range in eV
    channel : dict
        Dictionary whose keys are l-values and values are channel radii as a
        function of energy
    scattering : dict
        Dictionary whose keys are l-values and values are scattering radii as a
        function of energy

    Attributes
    ----------
    channel_radius : dict
        Dictionary whose keys are l-values and values are channel radii as a
        function of energy
    energy_max : float
        Maximum energy of the resolved resonance range in eV
    energy_min : float
        Minimum energy of the resolved resonance range in eV
    scattering_radius : dict
        Dictionary whose keys are l-values and values are scattering radii as a
        function of energ
    target_spin : float
        Intrinsic spin, :math:`I`, of the target nuclide

    """
    def __init__(self, target_spin, energy_min, energy_max, channel, scattering):
        self.target_spin = target_spin
        self.energy_min = energy_min
        self.energy_max = energy_max
        self.channel_radius = channel
        self.scattering_radius = scattering

        self._prepared = False
        self._parameter_matrix = {}

    def __copy__(self):
        cls = type(self)
        new_copy = cls.__new__(cls)
        new_copy.__dict__.update(self.__dict__)
        new_copy._prepared = False
        return new_copy

    @classmethod
    def from_endf(cls, ev, file_obj, items):
        """Create resonance range from an ENDF evaluation.

        This factory method is only used when LRU=0, indicating that only a
        scattering radius appears in MF=2, MT=151. All subclasses of
        ResonanceRange override this method with their own.

        Parameters
        ----------
        ev : openmc.data.endf.Evaluation
            ENDF evaluation
        file_obj : file-like object
            ENDF file positioned at the second record of a resonance range
            subsection in MF=2, MT=151
        items : list
            Items from the CONT record at the start of the resonance range
            subsection

        Returns
        -------
        openmc.data.ResonanceRange
            Resonance range data

        """
        energy_min, energy_max = items[0:2]

        # For scattering radius-only, NRO must be zero
        assert items[4] == 0

        # Get energy-independent scattering radius
        items = get_cont_record(file_obj)
        target_spin = items[0]
        ap = Polynomial((items[1],))

        # Calculate channel radius from ENDF-102 equation D.14
        a = Polynomial((0.123 * (NEUTRON_MASS*ev.target['mass'])**(1./3.) + 0.08,))

        return cls(target_spin, energy_min, energy_max, {0: a}, {0: ap})

    def reconstruct(self, energies):
        """Evaluate cross section at specified energies.

        Parameters
        ----------
        energies : float or Iterable of float
            Energies at which the cross section should be evaluated

        Returns
        -------
        3-tuple of float or numpy.ndarray
            Elastic, capture, and fission cross sections at the specified
            energies

        """
        if not _reconstruct:
            raise RuntimeError("Resonance reconstruction not available.")

        # Pre-calculate penetrations and shifts for resonances
        if not self._prepared:
            self._prepare_resonances()

        if isinstance(energies, Iterable):
            elastic = np.zeros_like(energies)
            capture = np.zeros_like(energies)
            fission = np.zeros_like(energies)

            for i, E in enumerate(energies):
                xse, xsg, xsf = self._reconstruct(self, E)
                elastic[i] = xse
                capture[i] = xsg
                fission[i] = xsf
        else:
            elastic, capture, fission = self._reconstruct(self, energies)

        return {2: elastic, 102: capture, 18: fission}

class MultiLevelBreitWigner(ResonanceRange):
    """Multi-level Breit-Wigner resolved resonance formalism data.

    Multi-level Breit-Wigner resolved resonance data is identified by LRF=2 in
    the ENDF-6 format.

    Parameters
    ----------
    target_spin : float
        Intrinsic spin, :math:`I`, of the target nuclide
    energy_min : float
        Minimum energy of the resolved resonance range in eV
    energy_max : float
        Maximum energy of the resolved resonance range in eV
    channel : dict
        Dictionary whose keys are l-values and values are channel radii as a
        function of energy
    scattering : dict
        Dictionary whose keys are l-values and values are scattering radii as a
        function of energy

    Attributes
    ----------
    atomic_weight_ratio : float
        Atomic weight ratio of the target nuclide given as a function of
        l-value. Note that this may be different than the value for the
        evaluation as a whole.
    channel_radius : dict
        Dictionary whose keys are l-values and values are channel radii as a
        function of energy
    energy_max : float
        Maximum energy of the resolved resonance range in eV
    energy_min : float
        Minimum energy of the resolved resonance range in eV
    parameters : pandas.DataFrame
        Energies, spins, and resonances widths for each resonance
    q_value : dict
        Q-value to be added to incident particle's center-of-mass energy to
        determine the channel energy for use in the penetrability factor. The
        keys of the dictionary are l-values.
    scattering_radius : dict
        Dictionary whose keys are l-values and values are scattering radii as a
        function of energy
    target_spin : float
        Intrinsic spin, :math:`I`, of the target nuclide

    """

    def __init__(self, target_spin, energy_min, energy_max, channel, scattering):
        super().__init__(target_spin, energy_min, energy_max, channel,
                         scattering)
        self.parameters = None
        self.q_value = {}
        self.atomic_weight_ratio = None

        # Set resonance reconstruction function
        if _reconstruct:
            self._reconstruct = reconstruct_mlbw
        else:
            self._reconstruct = None

    @classmethod
    def from_endf(cls, ev, file_obj, items):
        """Create MLBW data from an ENDF evaluation.

        Parameters
        ----------
        ev : openmc.data.endf.Evaluation
            ENDF evaluation
        file_obj : file-like object
            ENDF file positioned at the second record of a resonance range
            subsection in MF=2, MT=151
        items : list
            Items from the CONT record at the start of the resonance range
            subsection

        Returns
        -------
        openmc.data.MultiLevelBreitWigner
            Multi-level Breit-Wigner resonance parameters

        """

        # Read energy-dependent scattering radius if present
        energy_min, energy_max = items[0:2]
        nro, naps = items[4:6]
        if nro != 0:
            params, ape = get_tab1_record(file_obj)

        # Other scatter radius parameters
        items = get_cont_record(file_obj)
        target_spin = items[0]
        ap = Polynomial((items[1],))  # energy-independent scattering-radius
        NLS = items[4]  # number of l-values

        # Read resonance widths, J values, etc
        channel_radius = {}
        scattering_radius = {}
        q_value = {}
        records = []
        for l in range(NLS):
            items, values = get_list_record(file_obj)
            l_value = items[2]
            awri = items[0]
            q_value[l_value] = items[1]
            competitive = items[3]

            # Calculate channel radius from ENDF-102 equation D.14
            a = Polynomial((0.123 * (NEUTRON_MASS*awri)**(1./3.) + 0.08,))

            # Construct scattering and channel radius
            if nro == 0:
                scattering_radius[l_value] = ap
                if naps == 0:
                    channel_radius[l_value] = a
                elif naps == 1:
                    channel_radius[l_value] = ap
            elif nro == 1:
                scattering_radius[l_value] = ape
                if naps == 0:
                    channel_radius[l_value] = a
                elif naps == 1:
                    channel_radius[l_value] = ape
                elif naps == 2:
                    channel_radius[l_value] = ap

            energy = values[0::6]
            spin = values[1::6]
            gt = np.asarray(values[2::6])
            gn = np.asarray(values[3::6])
            gg = np.asarray(values[4::6])
            gf = np.asarray(values[5::6])
            if competitive > 0:
                gx = gt - (gn + gg + gf)
            else:
                gx = np.zeros_like(gt)

            for i, E in enumerate(energy):
                records.append([energy[i], l_value, spin[i], gt[i], gn[i],
                                gg[i], gf[i], gx[i]])

        columns = ['energy', 'L', 'J', 'totalWidth', 'neutronWidth',
                   'captureWidth', 'fissionWidth', 'competitiveWidth']
        parameters = pd.DataFrame.from_records(records, columns=columns)

        # Create instance of class
        mlbw = cls(target_spin, energy_min, energy_max,
                   channel_radius, scattering_radius)
        mlbw.q_value = q_value
        mlbw.atomic_weight_ratio = awri
        mlbw.parameters = parameters

        return mlbw

    def _prepare_resonances(self):
        df = self.parameters.copy()

        # Penetration and shift factors
        p = np.zeros(len(df))
        s = np.zeros(len(df))

        # Penetration and shift factors for competitive reaction
        px = np.zeros(len(df))
        sx = np.zeros(len(df))

        l_values = []
        competitive = []

        A = self.atomic_weight_ratio
        for i, E, l, J, gt, gn, gg, gf, gx in df.itertuples():
            if l not in l_values:
                l_values.append(l)
                competitive.append(gx > 0)

            # Determine penetration and shift corresponding to resonance energy
            k = wave_number(A, E)
            rho = k*self.channel_radius[l](E)
            rhohat = k*self.scattering_radius[l](E)
            p[i], s[i] = penetration_shift(l, rho)

            # Determine penetration at modified energy for competitive reaction
            if gx > 0:
                Ex = E + self.q_value[l]*(A + 1)/A
                rho = k*self.channel_radius[l](Ex)
                rhohat = k*self.scattering_radius[l](Ex)
                px[i], sx[i] = penetration_shift(l, rho)
            else:
                px[i] = sx[i] = 0.0

        df['p'] = p
        df['s'] = s
        df['px'] = px
        df['sx'] = sx

        self._l_values = np.array(l_values)
        self._competitive = np.array(competitive)
        for l in l_values:
            self._parameter_matrix[l] = df[df.L == l].values

        self._prepared = True

class SingleLevelBreitWigner(MultiLevelBreitWigner):
    """Single-level Breit-Wigner resolved resonance formalism data.

    Single-level Breit-Wigner resolved resonance data is is identified by LRF=1
    in the ENDF-6 format.

    Parameters
    ----------
    target_spin : float
        Intrinsic spin, :math:`I`, of the target nuclide
    energy_min : float
        Minimum energy of the resolved resonance range in eV
    energy_max : float
        Maximum energy of the resolved resonance range in eV
    channel : dict
        Dictionary whose keys are l-values and values are channel radii as a
        function of energy
    scattering : dict
        Dictionary whose keys are l-values and values are scattering radii as a
        function of energy

    Attributes
    ----------
    atomic_weight_ratio : float
        Atomic weight ratio of the target nuclide given as a function of
        l-value. Note that this may be different than the value for the
        evaluation as a whole.
    channel_radius : dict
        Dictionary whose keys are l-values and values are channel radii as a
        function of energy
    energy_max : float
        Maximum energy of the resolved resonance range in eV
    energy_min : float
        Minimum energy of the resolved resonance range in eV
    parameters : pandas.DataFrame
        Energies, spins, and resonances widths for each resonance
    q_value : dict
        Q-value to be added to incident particle's center-of-mass energy to
        determine the channel energy for use in the penetrability factor. The
        keys of the dictionary are l-values.
    scattering_radius : dict
        Dictionary whose keys are l-values and values are scattering radii as a
        function of energy
    target_spin : float
        Intrinsic spin, :math:`I`, of the target nuclide

    """

    def __init__(self, target_spin, energy_min, energy_max, channel, scattering):
        super().__init__(target_spin, energy_min, energy_max, channel,
                         scattering)

        # Set resonance reconstruction function
        if _reconstruct:
            self._reconstruct = reconstruct_slbw
        else:
            self._reconstruct = None

class ReichMoore(ResonanceRange):
    """Reich-Moore resolved resonance formalism data.

    Reich-Moore resolved resonance data is identified by LRF=3 in the ENDF-6
    format.

    Parameters
    ----------
    target_spin : float
        Intrinsic spin, :math:`I`, of the target nuclide
    energy_min : float
        Minimum energy of the resolved resonance range in eV
    energy_max : float
        Maximum energy of the resolved resonance range in eV
    channel : dict
        Dictionary whose keys are l-values and values are channel radii as a
        function of energy
    scattering : dict
        Dictionary whose keys are l-values and values are scattering radii as a
        function of energy

    Attributes
    ----------
    angle_distribution : bool
        Indicate whether parameters can be used to compute angular distributions
    atomic_weight_ratio : float
        Atomic weight ratio of the target nuclide given as a function of
        l-value. Note that this may be different than the value for the
        evaluation as a whole.
    channel_radius : dict
        Dictionary whose keys are l-values and values are channel radii as a
        function of energy
    energy_max : float
        Maximum energy of the resolved resonance range in eV
    energy_min : float
        Minimum energy of the resolved resonance range in eV
    num_l_convergence : int
        Number of l-values which must be used to converge the calculation
    scattering_radius : dict
        Dictionary whose keys are l-values and values are scattering radii as a
        function of energy
    parameters : pandas.DataFrame
        Energies, spins, and resonances widths for each resonance
    target_spin : float
        Intrinsic spin, :math:`I`, of the target nuclide

    """

    def __init__(self, target_spin, energy_min, energy_max, channel, scattering):
        super().__init__(target_spin, energy_min, energy_max, channel,
                         scattering)
        self.parameters = None
        self.angle_distribution = False
        self.num_l_convergence = 0

        # Set resonance reconstruction function
        if _reconstruct:
            self._reconstruct = reconstruct_rm
        else:
            self._reconstruct = None

    @classmethod
    def from_endf(cls, ev, file_obj, items):
        """Create Reich-Moore resonance data from an ENDF evaluation.

        Parameters
        ----------
        ev : openmc.data.endf.Evaluation
            ENDF evaluation
        file_obj : file-like object
            ENDF file positioned at the second record of a resonance range
            subsection in MF=2, MT=151
        items : list
            Items from the CONT record at the start of the resonance range
            subsection

        Returns
        -------
        openmc.data.ReichMoore
            Reich-Moore resonance parameters

        """
        # Read energy-dependent scattering radius if present
        energy_min, energy_max = items[0:2]
        nro, naps = items[4:6]
        if nro != 0:
            params, ape = get_tab1_record(file_obj)

        # Other scatter radius parameters
        items = get_cont_record(file_obj)
        target_spin = items[0]
        ap = Polynomial((items[1],))
        angle_distribution = (items[3] == 1)  # Flag for angular distribution
        NLS = items[4]  # Number of l-values
        num_l_convergence = items[5]  # Number of l-values for convergence

        # Read resonance widths, J values, etc
        channel_radius = {}
        scattering_radius = {}
        records = []
        for i in range(NLS):
            items, values = get_list_record(file_obj)
            apl = Polynomial((items[1],)) if items[1] != 0.0 else ap
            l_value = items[2]
            awri = items[0]

            # Calculate channel radius from ENDF-102 equation D.14
            a = Polynomial((0.123 * (NEUTRON_MASS*awri)**(1./3.) + 0.08,))

            # Construct scattering and channel radius
            if nro == 0:
                scattering_radius[l_value] = apl
                if naps == 0:
                    channel_radius[l_value] = a
                elif naps == 1:
                    channel_radius[l_value] = apl
            elif nro == 1:
                if naps == 0:
                    channel_radius[l_value] = a
                    scattering_radius[l_value] = ape
                elif naps == 1:
                    channel_radius[l_value] = scattering_radius[l_value] = ape
                elif naps == 2:
                    channel_radius[l_value] = apl
                    scattering_radius[l_value] = ape

            energy = values[0::6]
            spin = values[1::6]
            gn = values[2::6]
            gg = values[3::6]
            gfa = values[4::6]
            gfb = values[5::6]

            for i, E in enumerate(energy):
                records.append([energy[i], l_value, spin[i], gn[i], gg[i],
                                gfa[i], gfb[i]])

        # Create pandas DataFrame with resonance data
        columns = ['energy', 'L', 'J', 'neutronWidth', 'captureWidth',
                   'fissionWidthA', 'fissionWidthB']
        parameters = pd.DataFrame.from_records(records, columns=columns)

        # Create instance of ReichMoore
        rm = cls(target_spin, energy_min, energy_max,
                 channel_radius, scattering_radius)
        rm.parameters = parameters
        rm.angle_distribution = angle_distribution
        rm.num_l_convergence = num_l_convergence
        rm.atomic_weight_ratio = awri

        return rm

    def _prepare_resonances(self):
        df = self.parameters.copy()

        # Penetration and shift factors
        p = np.zeros(len(df))
        s = np.zeros(len(df))

        l_values = []
        lj_values = []

        A = self.atomic_weight_ratio
        for i, E, l, J, gn, gg, gfa, gfb in df.itertuples():
            if l not in l_values:
                l_values.append(l)
            if (l, abs(J)) not in lj_values:
                lj_values.append((l, abs(J)))

            # Determine penetration and shift corresponding to resonance energy
            k = wave_number(A, E)
            rho = k*self.channel_radius[l](E)
            rhohat = k*self.scattering_radius[l](E)
            p[i], s[i] = penetration_shift(l, rho)

        df['p'] = p
        df['s'] = s

        self._l_values = np.array(l_values)
        for (l, J) in lj_values:
            self._parameter_matrix[l, J] = df[(df.L == l) &
                                              (abs(df.J) == J)].values

        self._prepared = True

class RMatrixLimited(ResonanceRange):
    """R-matrix limited resolved resonance formalism data.

    R-matrix limited resolved resonance data is identified by LRF=7 in the
    ENDF-6 format.

    Parameters
    ----------
    energy_min : float
        Minimum energy of the resolved resonance range in eV
    energy_max : float
        Maximum energy of the resolved resonance range in eV
    particle_pairs : list of dict
        List of particle pairs. Each particle pair is represented by a
        dictionary that contains the mass, atomic number, spin, and parity of
        each particle as well as other characteristics.
    spin_groups : list of dict
        List of spin groups. Each spin group is characterized by channels,
        resonance energies, and resonance widths.

    Attributes
    ----------
    reduced_width : bool
        Flag indicating whether channel widths in eV or reduced-width amplitudes
        in eV^1/2 are given
    formalism : int
        Flag to specify which formulae for the R-matrix are to be used
    particle_pairs : list of dict
        List of particle pairs. Each particle pair is represented by a
        dictionary that contains the mass, atomic number, spin, and parity of
        each particle as well as other characteristics.
    spin_groups : list of dict
        List of spin groups. Each spin group is characterized by channels,
        resonance energies, and resonance widths.

    """

    def __init__(self, energy_min, energy_max, particle_pairs, spin_groups):
        super().__init__(0.0, energy_min, energy_max, None, None)
        self.reduced_width = False
        self.formalism = 3
        self.particle_pairs = particle_pairs
        self.spin_groups = spin_groups

    @classmethod
    def from_endf(cls, ev, file_obj, items):
        """Read R-Matrix limited resonance data from an ENDF evaluation.

        Parameters
        ----------
        ev : openmc.data.endf.Evaluation
            ENDF evaluation
        file_obj : file-like object
            ENDF file positioned at the second record of a resonance range
            subsection in MF=2, MT=151
        items : list
            Items from the CONT record at the start of the resonance range
            subsection

        Returns
        -------
        openmc.data.RMatrixLimited
            R-matrix limited resonance parameters

        """
        energy_min, energy_max = items[0:2]

        items = get_cont_record(file_obj)
        reduced_width = (items[2] == 1)  # reduced width amplitude?
        formalism = items[3]  # Specify which formulae are used
        n_spin_groups = items[4]  # Number of Jpi values (NJS)

        particle_pairs = []
        spin_groups = []

        items, values = get_list_record(file_obj)
        n_pairs = items[5]//2  # Number of particle pairs (NPP)
        for i in range(n_pairs):
            first = {'mass': values[12*i],
                     'z': int(values[12*i + 2]),
                     'spin': values[12*i + 4],
                     'parity': values[12*i + 10]}
            second = {'mass': values[12*i + 1],
                      'z': int(values[12*i + 3]),
                      'spin': values[12*i + 5],
                      'parity': values[12*i + 11]}

            q_value = values[12*i + 6]
            penetrability = values[12*i + 7]
            shift = values[12*i + 8]
            mt = int(values[12*i + 9])

            particle_pairs.append(ParticlePair(
                first, second, q_value, penetrability, shift, mt))

        # loop over spin groups
        for i in range(n_spin_groups):
            items, values = get_list_record(file_obj)
            J = items[0]
            if J == 0.0:
                parity = '+' if items[1] == 1.0 else '-'
            else:
                parity = '+' if J > 0. else '-'
                J = abs(J)
            kbk = items[2]
            kps = items[3]
            n_channels = items[5]
            channels = []
            for j in range(n_channels):
                channel = {}
                channel['particle_pair'] = particle_pairs[
                    int(values[6*j]) - 1]
                channel['l'] = values[6*j + 1]
                channel['spin'] = values[6*j + 2]
                channel['boundary'] = values[6*j + 3]
                channel['effective_radius'] = values[6*j + 4]
                channel['true_radius'] = values[6*j + 5]
                channels.append(channel)

            # Read resonance energies and widths
            items, values = get_list_record(file_obj)
            n_resonances = items[3]
            records = []
            m = n_channels//6 + 1
            for j in range(n_resonances):
                energy = values[6*m*j]
                records.append([energy] + [values[6*m*j + k + 1]
                                           for k in range(n_channels)])

            # Determine column names
            columns = ['energy']
            for channel in channels:
                mt = channel['particle_pair'].mt
                if mt == 2:
                    columns.append('neutronWidth')
                elif mt == 18:
                    columns.append('fissionWidth')
                elif mt == 102:
                    columns.append('captureWidth')
                else:
                    columns.append('width (MT={})'.format(mt))

            # Create Pandas dataframe with resonance parameters
            parameters = pd.DataFrame.from_records(records, columns=columns)

            # Construct SpinGroup instance and add to list
            sg = SpinGroup(J, parity, channels, parameters)
            spin_groups.append(sg)

            # Optional extension (Background R-Matrix)
            if kbk > 0:
                items, values = get_list_record(file_obj)
                lbk = items[4]
                if lbk == 1:
                    params, rbr = get_tab1_record(file_obj)
                    params, rbi = get_tab1_record(file_obj)

            # Optional extension (Tabulated phase shifts)
            if kps > 0:
                items, values = get_list_record(file_obj)
                lps = items[4]
                if lps == 1:
                    params, psr = get_tab1_record(file_obj)
                    params, psi = get_tab1_record(file_obj)

        rml = cls(energy_min, energy_max, particle_pairs, spin_groups)
        rml.reduced_width = reduced_width
        rml.formalism = formalism

        return rml

class ParticlePair(object):
    def __init__(self, first, second, q_value, penetrability,
                 shift, mt):
        self.first = first
        self.second = second
        self.q_value = q_value
        self.penetrability = penetrability
        self.shift = shift
        self.mt = mt

class SpinGroup(object):
    """Resonance spin group

    Attributes
    ----------
    spin : float
        Total angular momentum (nuclear spin)
    parity : {'+', '-'}
        Even (+) or odd(-) parity
    channels : list of openmc.data.Channel
        Available channels
    parameters : pandas.DataFrame
        Energies/widths for each resonance/channel

    """

    def __init__(self, spin, parity, channels, parameters):
        self.spin = spin
        self.parity = parity
        self.channels = channels
        self.parameters = parameters

    def __repr__(self):
        return '<SpinGroup: Jpi={}{}>'.format(self.spin, self.parity)

class Unresolved(ResonanceRange):
    """Unresolved resonance parameters as identified by LRU=2 in MF=2.

    Parameters
    ----------
    target_spin : float
        Intrinsic spin, :math:`I`, of the target nuclide
    energy_min : float
        Minimum energy of the unresolved resonance range in eV
    energy_max : float
        Maximum energy of the unresolved resonance range in eV
    channel : openmc.data.Function1D
        Channel radii as a function of energy
    scattering : openmc.data.Function1D
        Scattering radii as a function of energy

    Attributes
    ----------
    add_to_background : bool
        If True, file 3 contains partial cross sections to be added to the
        average unresolved cross sections calculated from parameters.
    atomic_weight_ratio : float
        Atomic weight ratio of the target nuclide
    channel_radius : openmc.data.Function1D
        Channel radii as a function of energy
    energies : Iterable of float
        Energies at which parameters are tabulated
    energy_max : float
        Maximum energy of the unresolved resonance range in eV
    energy_min : float
        Minimum energy of the unresolved resonance range in eV
    parameters : list of pandas.DataFrame
        Average resonance parameters at each energy
    scattering_radius : openmc.data.Function1D
        Scattering radii as a function of energy
    target_spin : float
        Intrinsic spin, :math:`I`, of the target nuclide

    """

    def __init__(self, target_spin, energy_min, energy_max, channel, scattering):
        super().__init__(target_spin, energy_min, energy_max, channel,
                         scattering)
        self.energies = None
        self.parameters = None
        self.add_to_background = False
        self.atomic_weight_ratio = None

    @classmethod
    def from_endf(cls, file_obj, items, fission_widths):
        """Read unresolved resonance data from an ENDF evaluation.

        Parameters
        ----------
        file_obj : file-like object
            ENDF file positioned at the second record of a resonance range
            subsection in MF=2, MT=151
        items : list
            Items from the CONT record at the start of the resonance range
            subsection
        fission_widths : bool
            Whether fission widths are given

        Returns
        -------
        openmc.data.Unresolved
            Unresolved resonance region parameters

        """
        # Read energy-dependent scattering radius if present
        energy_min, energy_max = items[0:2]
        nro, naps = items[4:6]
        if nro != 0:
            params, ape = get_tab1_record(file_obj)

        # Get SPI, AP, and LSSF
        formalism = items[3]
        if not (fission_widths and formalism == 1):
            items = get_cont_record(file_obj)
            target_spin = items[0]
            if nro == 0:
                ap = Polynomial((items[1],))
            add_to_background = (items[2] == 0)

        if not fission_widths and formalism == 1:
            # Case A -- fission widths not given, all parameters are
            # energy-independent
            NLS = items[4]
            columns = ['L', 'J', 'd', 'amun', 'gn0', 'gg']
            records = []
            for ls in range(NLS):
                items, values = get_list_record(file_obj)
                awri = items[0]
                l = items[2]
                NJS = items[5]
                for j in range(NJS):
                    d, j, amun, gn0, gg = values[6*j:6*j + 5]
                    records.append([l, j, d, amun, gn0, gg])
            parameters = pd.DataFrame.from_records(records, columns=columns)
            energies = None

        elif fission_widths and formalism == 1:
            # Case B -- fission widths given, only fission widths are
            # energy-dependent
            items, energies = get_list_record(file_obj)
            target_spin = items[0]
            if nro == 0:
                ap = Polynomial((items[1],))
            add_to_background = (items[2] == 0)
            NE, NLS = items[4:6]
            records = []
            columns = ['L', 'J', 'E', 'd', 'amun', 'amuf', 'gn0', 'gg', 'gf']
            for ls in range(NLS):
                items = get_cont_record(file_obj)
                awri = items[0]
                l = items[2]
                NJS = items[4]
                for j in range(NJS):
                    items, values = get_list_record(file_obj)
                    muf = items[3]
                    d = values[0]
                    j = values[1]
                    amun = values[2]
                    gn0 = values[3]
                    gg = values[4]
                    gfs = values[6:]
                    for E, gf in zip(energies, gfs):
                        records.append([l, j, E, d, amun, muf, gn0, gg, gf])
            parameters = pd.DataFrame.from_records(records, columns=columns)

        elif formalism == 2:
            # Case C -- all parameters are energy-dependent
            NLS = items[4]
            columns = ['L', 'J', 'E', 'd', 'amux', 'amun', 'amuf', 'gx', 'gn0',
                       'gg', 'gf']
            records = []
            for ls in range(NLS):
                items = get_cont_record(file_obj)
                awri = items[0]
                l = items[2]
                NJS = items[4]
                for j in range(NJS):
                    items, values = get_list_record(file_obj)
                    ne = items[5]
                    j = items[0]
                    amux = values[2]
                    amun = values[3]
                    amuf = values[5]
                    energies = []
                    for k in range(1, ne + 1):
                        E = values[6*k]
                        d = values[6*k + 1]
                        gx = values[6*k + 2]
                        gn0 = values[6*k + 3]
                        gg = values[6*k + 4]
                        gf = values[6*k + 5]
                        energies.append(E)
                        records.append([l, j, E, d, amux, amun, amuf, gx, gn0,
                                        gg, gf])
            parameters = pd.DataFrame.from_records(records, columns=columns)

        # Calculate channel radius from ENDF-102 equation D.14
        a = Polynomial((0.123 * (NEUTRON_MASS*awri)**(1./3.) + 0.08,))

        # Determine scattering and channel radius
        if nro == 0:
            scattering_radius = ap
            if naps == 0:
                channel_radius = a
            elif naps == 1:
                channel_radius = ap
        elif nro == 1:
            scattering_radius = ape
            if naps == 0:
                channel_radius = a
            elif naps == 1:
                channel_radius = ape
            elif naps == 2:
                channel_radius = ap

        urr = cls(target_spin, energy_min, energy_max, channel_radius,
                  scattering_radius)
        urr.parameters = parameters
        urr.add_to_background = add_to_background
        urr.atomic_weight_ratio = awri
        urr.energies = energies

        return urr

_FORMALISMS = {0: ResonanceRange,
               1: SingleLevelBreitWigner,
               2: MultiLevelBreitWigner,
               3: ReichMoore,
               7: RMatrixLimited}

_RESOLVED = (SingleLevelBreitWigner, MultiLevelBreitWigner,
             ReichMoore, RMatrixLimited)


# coding: utf-8
# Copyright (c) Max-Planck-Institut für Eisenforschung GmbH - Computational Materials Design (CM) Department
# Distributed under the terms of "New BSD License", see the LICENSE file.

import numpy as np
from pyiron.atomistics.job.interactivewrapper import InteractiveWrapper, ReferenceJobOutput
from pyiron_base import InputList, Settings

__author__ = "Osamu Waseda"
__copyright__ = "Copyright 2020, Max-Planck-Institut für Eisenforschung GmbH " \
                "- Computational Materials Design (CM) Department"
__version__ = "1.0"
__maintainer__ = "Osamu Waseda"
__email__ = "waseda@mpie.de"
__status__ = "development"
__date__ = "Sep 1, 2018"

s = Settings()

class ARTInteractive(object):
    def __init__(self, art_id, direction, gamma=0.1, fix_layer=False, non_art_id=None):
        if int(art_id) != art_id or art_id < 0:
            raise ValueError('art_id must be a posive integer')
        if len(direction) != 3 or np.isclose(np.linalg.norm(direction), 0):
            raise ValueError('direction must be a finite 3d vector')
        if gamma < 0:
            raise ValueError('gamma must be a positive float')
        if fix_layer and non_art_id is not None:
            raise ValueError('fix_layer and non_art_id cannot be set at the same time')
        self.art_id = art_id
        self.direction = direction
        self.gamma = gamma
        self.non_art_id = non_art_id
        if non_art_id is not None:
            self.non_art_id = np.array([non_art_id]).flatten()
        self.fix_layer = fix_layer

    @property
    def _R(self):
        value = np.array(self.direction)
        value = value / np.linalg.norm(value)
        return np.outer(value, value)

    def get_forces(self, f_in):
        f = np.array(f_in)
        if len(f.shape) == 2:
            f = np.array([f])
        if self.non_art_id is None:
            self.non_art_id = np.arange(len(f[0])) != self.art_id
        f_art = (1.0+self.gamma)*np.einsum('ij,nj->ni', self._R, f[:, self.art_id])
        if self.fix_layer:
            f[:, self.non_art_id] = np.einsum('nmj,ij->nmi', f[:, self.non_art_id], np.identity(3)-self._R)
        else:
            f[:, self.non_art_id] += f_art[:, np.newaxis, :] / np.sum(self.non_art_id != False)
        f[:, self.art_id] -= f_art
        return f.reshape(np.array(f_in).shape)

class ART(InteractiveWrapper):
    """
    Apply an artificial force according to the Activation Relaxation Technique (ART)
    DOI:https://doi.org/10.1103/PhysRevE.57.2419

    The applied force f_art is calculated from the original force f by:

    f_art = f-(1+gamma)*np.dot(n,f)*n

    where gamma is a parameter to be determined in the input (default: 0.1) and n is
    the direction along which the force is reversed (3d-vector). In order to homogenize
    the total force in the system, f_art is distributed among atoms specified by
    non_art_id (default: all atoms), or if fix_layer is defined, the forces acting on
    all the other atoms along the direction n are cancelled. Note: Since the energy
    is not compatible with the forces anymore, structure optimization methods which
    rely on the energy variation (such as conjugate gradient) cannot/should not be used.

    Input:
        - art_id (int): atom id on which ART force is applied
        - direction (list/numpy.ndarray): direction along which force is reversed
        - gamma (float): prefactor for force inversion. v.s.
        - non_art_id (list/numpy.ndarray): list of atoms to be used for the force cancellation
        - fix_layer (bool): whether or not to fix all other atoms on the layer perpendicular
            to the direction along which force is reversed.

    Example:

        # Structure creation
        >>> vacancy_position = structure.positions[0]
        >>> del structure[0]
        >>> neighbors = structure.get_neighborhood(vacancy_position)
        >>> direction = neighbors.vecs[0]
        >>> art_id = neighbors.indices[0]
        >>> structure.positions[art_id] -= 0.5*direction

        # Job creation
        >>> some_atomistic_job.structure = structure
        >>> art = ART(job_name='art')
        >>> art.ref_job = some_atomistic_job
        >>> art.input.art_id = art_id
        >>> art.input.direction = direction
        >>> some_minimizer.ref_job = art
        >>> some_minimizer.run()

    """
    def __init__(self, project, job_name):
        super(ART, self).__init__(project, job_name)
        self.__name__ = "ART"
        self.input = InputList(table_name='custom_dict')
        self.input.gamma = 0.1
        self.input.fix_layer = False
        self.input.non_art_id = None
        self.input.art_id = None
        self.input.direction = None
        self.output = ARTIntOutput(job=self)
        self.server.run_mode.interactive = True
        self._interactive_interface = None
        self._art = None

    def set_input_to_read_only(self):
        """
        This function enforces read-only mode for the input classes, but it has to be implement in the individual
        classes.
        """
        self.input.read_only = True

    @property
    def art(self):
        if self._art is None:
            self._art = ARTInteractive(
                art_id=self.input.art_id,
                direction=self.input.direction,
                gamma=self.input.gamma,
                fix_layer=self.input.fix_layer,
                non_art_id=self.input.non_art_id
            )
        return self._art

    def run_if_interactive(self):
        self._logger.debug('art status: ' + str(self.status))
        if not self.status.running:
            self.ref_job_initialize()
        self.status.running = True
        if self.ref_job.server.run_mode.interactive:
            self.ref_job.run()
        else:
            self.ref_job.run(run_again=True)
        self._logger.debug('art status: ' + str(self.status))

    def interactive_forces_getter(self):
        return self.art.get_forces(self.ref_job.output.forces[-1])

    def validate_ready_to_run(self):
        """
            check whether art_id and direction are set
        """
        if self.input.art_id is None or self.input.direction is None:
            raise AssertionError('art_id and/or direction not set')

    def interactive_close(self):
        self.status.collect = True
        if self.ref_job.server.run_mode.interactive:
            self.ref_job.interactive_close()
        self.status.finished = True

class ARTIntOutput(ReferenceJobOutput):
    def __init__(self, job):
        super(ARTIntOutput, self).__init__(job=job)

    @property
    def forces(self):
        return self._job.art.get_forces(self._job.ref_job.output.forces)