function Base.:*(A::XRTArray, B::XRTArray)
    ddots = DimNums((1,), (0,), (), ())
    HloDot(ddots)(A, B)
end
function Base.:*(A::XRTVector, B::XRTArray)
    ddots = DimNums((1,), (0,), (), ())
    HloDot(ddots)(HloReshape((length(A), 1))(A), B)
end
Base.:*(A::XRTArray{<:Any, (), 0}, B::XRTVector) = broadcast(*, A, B)

using LinearAlgebra

make_hlo_transpose(A::XRTArray) = HloTranspose((1,0))(A)
make_hlo_transpose(A::XRTVector) = make_hlo_transpose(HloReshape((length(A), 1))(A))

function Base.:*(A::Adjoint{<:Any, <:XRTArray}, B::XRTArray)
    ddots = DimNums((1,), (0,), (), ())
    HloDot(ddots)(make_hlo_transpose(A.parent), B)
end

function Base.:*(A::Adjoint{<:Any, <:XRTVector}, B::XRTVector)
    ddots = DimNums((0,), (0,), (), ())
    HloDot(ddots)(A.parent, B)
end

function Base.:*(A::Transpose{<:Any, <:XRTArray}, B::XRTArray)
    ddots = DimNums((1,), (0,), (), ())
    HloDot(ddots)(make_hlo_transpose(A.parent), B)
end

function Base.:*(A::XRTArray{T}, B::Transpose{S}) where {T, S}
    ddots = DimNums((1,), (0,), (), ())
    HloDot(ddots)(convert(XRTArray{promote_type(T, S)}, A), make_hlo_transpose(convert(XRTArray{promote_type(T, S)}, B.parent)))
end

function Base.:*(A::XRTVector{T}, B::Transpose{S, <:XRTVector}) where {T, S}
    ddots = DimNums((1,), (0,), (), ())
    HloDot(ddots)(reshape(convert(XRTArray{promote_type(T, S)}, A), (size(A, 1), 1)),
        make_hlo_transpose(convert(XRTArray{promote_type(T, S)}, B.parent)))
end

function Base.:*(A::Transpose{<:Any, <:XRTVector}, B::XRTVector)
    ddots = DimNums((0,), (0,), (), ())
    HloDot(ddots)(A.parent, B)
end

Base.convert(::Type{XRTArray}, x::XLAScalar) = XRTArray(x)

const XRTScalar{T} = XRTArray{T, (), 0}
@noinline Base.convert(::Type{Bool}, A::XRTScalar{Bool}) = convert(Array, A)[]
function Base.promote_rule(A::Type{XRTScalar{T}}, B::Type{XRTScalar{S}}) where {T<:XLAScalar, S<:XLAScalar}
    XRTArray{promote_type(T, S), (), 0}
end

function Base.promote_rule(A::Type{T}, B::Type{XRTScalar{S}}) where {T<:XLAScalar, S<:XLAScalar}
    XRTArray{promote_type(T, S), (), 0}
end

# Scalar conversions - Happens via HloConvert. In the future, we may
# want to make sure to more closely match julia semantics.
Base.convert(::Type{XRTScalar{T}}, x::XRTScalar{S}) where {T,S} = GenericHloOp{:convert}(T, ())(x)
Base.convert(::Type{XRTScalar{T}}, x::T) where T = XRTScalar{T}(x)
Base.convert(::Type{XRTScalar{T}}, x::XRTScalar{T}) where T = x
Base.convert(::Type{XRTScalar{T}}, x::XLAScalar) where {T<:XLAScalar} = XRTScalar{T}(convert(T, x))

# A couple of scalar embeddings (could use casette in the future)
using Base.Meta
for (binop, xlaop) in (
        (:+, :add),
        (:-, :subtract),
        (:/, :divide),
        (:*, :multiply),
        (:%, :remainder),
        (:&, :and),
        (:<<, Symbol("shift-left")),
        (:>>>, Symbol("shift-right-logical")),
        (:^, :power),
    )
    @eval Base.$(binop)(A::XRTScalar{T},
                        B::XRTScalar{T}) where {T<:XLAScalar} =
        GenericHloOp{$(quot(xlaop))}(T, ())(A, B)
    @eval Base.$(binop)(A::XRTScalar, B::XRTScalar, args::XRTScalar{<:XLAScalar}...) =
        $(binop)(promote(A, B, args...)...)
    @eval Base.$(binop)(A::XRTScalar, B::Number) = $(binop)(promote(A, B)...)
    @eval Base.$(binop)(A::Number, B::XRTScalar) = $(binop)(promote(A, B)...)
    @eval Base.$(binop)(A::XRTScalar{T}, B::XRTScalar{T}) where {T<:XRTArray} =
        Base.$(binop)(convert(T, A), convert(T, B))
end
Base.:>>(A::XRTScalar{<:Unsigned}, B::XRTScalar) = A >>> B
Base.:>>(A::XRTScalar{T}, B::XRTScalar) where {T <: Signed} = GenericHloOp{Symbol("shift-right-arithmetic")}(T, ())(A, B)

Base.:^(A::XRTScalar{T}, B::Integer) where {T} = A ^ XRTArray(Float32(B))

Base.literal_pow(::typeof(^), x::XRTScalar, ::Val{0}) = one(x)
Base.literal_pow(::typeof(^), x::XRTScalar, ::Val{1}) = x
Base.literal_pow(::typeof(^), x::XRTScalar, ::Val{2}) = x*x
Base.literal_pow(::typeof(^), x::XRTScalar, ::Val{3}) = x*x*x

# XRT Scalars are not numbers so import some functions manually
import Base: /, \, *
for f in (:/, :\, :*)
    if f != :/
        @eval ($f)(A::XRTArray{<:Any, (), 0}, B::XRTArray) = broadcast($f, A, B)
    end
    if f != :\
        @eval ($f)(A::XRTArray, B::XRTArray{<:Any, (), 0}) = broadcast($f, A, B)
    end
end

embed(x::XLAScalar) = XRTArray(x)
embed(x) = x

concat_dims(::Type{<:XRTArray{T, Sz1}}, Sz2::Tuple) where {T, Sz1} = (Sz1..., Sz2...)
concat_dims(::Type, Sz2::Tuple) where {T, Sz1} = Sz2

@noinline make_zeros_dims_tuple(::Type{T}) where {T} = ntuple(i->i-1, ndims(T))

function Base.zeros(::Type{XRTArray{T, Sz, N}}) where {T, Sz, N}
    # TODO: This assignment being necessary is probably a compiler bug
    TT = T
    HloBroadcast(make_zeros_dims_tuple(TT), Sz)(XRTArray((zero(TT),)))::XRTArray{T, Sz, N}
end
Base.zero(X::Type{XRTArray{T, Sz, N}}) where {T, Sz, N} = zeros(X)

Base.zero(A::Type{XRTArray{T, (), 0}}) where T = XRTArray(zero(T))
Base.zero(A::XRTArray) = zero(typeof(A))
Base.one(A::Type{XRTArray{T, (), 0}}) where T = XRTArray(one(T))
Base.one(A::XRTArray{<:Any, (), 0}) = one(typeof(A))
Base.max(A::XRTArray{T, (), 0}, B::XRTArray{T, (), 0}) where {T<:XLAScalar} =
    GenericHloOp{:maximum}(T, ())(A, B)

for (fname, hlo) in (
        (:exp, :exponential),
        (:sin, :sine),
        (:cos, :cosine),
        (:tanh, :tanh),
        (:expm1, Symbol("exponential-minus-one")),
        (:log, :log),
        (:log1p, Symbol("log-plus-one"))
    )
    @eval Base.$(fname)(x::XRTArray{T, (), 0}) where {T} = GenericHloOp{$(quot(hlo))}(T, ())(x)
end

Base.sqrt(A::XRTArray{T, (), 0}) where {T} = GenericHloOp{:power}(T, ())(A, XRTArray(convert(T, 0.5)))

@eval Base.isless(A::XRTArray{<:Any, (), 0}, B::XRTArray{<:Any, (), 0}) =
    convert(Bool, GenericHloOp{$(QuoteNode(Symbol("less-than")))}(Bool, ())(A, B))
@eval Base.:(<=)(A::XRTArray{<:Any, (), 0}, B::XRTArray{<:Any, (), 0}) =
    GenericHloOp{$(QuoteNode(Symbol("less-than-or-equal-to")))}(Bool, ())(A, B)
@eval Base.:(==)(A::XRTArray{<:Any, (), 0}, B::XRTArray{<:Any, (), 0}) =
    GenericHloOp{$(QuoteNode(Symbol("equal-to")))}(Bool, ())(A, B)
@eval Base.:(!=)(A::XRTArray{<:Any, (), 0}, B::XRTArray{<:Any, (), 0}) =
    GenericHloOp{$(QuoteNode(Symbol("not-equal-to")))}(Bool, ())(A, B)
Base.inv(x::XRTScalar{T}) where {T} = XRTArray(one(T))/x

Base.:-(x::XRTScalar{T}) where {T} = -(XRTArray(zero(T)), x)

# XRTArrays are immutable. Copy is identity.
Base.copy(A::XRTArray) = A
#Base.transpose(A::XRTArray) = HloTranspose((1,0))(A)
#Base.adjoint(A::XRTArray{<:Real}) = HloTranspose((1,0))(A)
Base.permutedims(A::XRTArray, perm) = HloTranspose(map(x->x-1, perm))(A)

# Recursive version of promote shape
# TODO: This shouldn't be necessary if inference did better
# purity/nothrow modeling.
check_shape_match(a::Tuple{}, b::Tuple{}) = true
function check_shape_match(a::Tuple, b::Tuple{})
    first(a) == 1:1 || return false
    check_shape_match(Base.tail(a), ())
end
function check_shape_match(a::Tuple, b::Tuple)
    first(a) == first(b) || return false
    check_shape_match(Base.tail(a), Base.tail(b))
end
function promote_shape_rec(a, b)
    if length(a) < length(b)
        return promote_shape(b, a)
    end
    check_shape_match(a, b) || error("Dimensions must match")
    a
end

function Base.promote_shape(a::XRTArray, b::XRTArray)
    promote_shape_rec(axes(a), axes(b))
end

import Base.Broadcast

Base.similar(x::XRTArray) = error()

struct XRTArrayStyle{N} <: Broadcast.AbstractArrayStyle{N} end
(::Type{<:XRTArrayStyle})(::Val{N}) where {N} = XRTArrayStyle{N}()
Broadcast.BroadcastStyle(::Type{<:XRTArray{<:Any,Dims,N}}) where {Dims, N} =
    XRTArrayStyle{N}()

@noinline _ccdims(s) = tuple(findall(==(1), s)...)
@noinline _ncdims(s) = tuple(findall(x->x!=(1), s)...)

@Base.pure ccdims(s) = _ccdims(s)
@Base.pure ncdims(s) = _ncdims(s)

@Base.pure getindex_tuple(s::Tuple, k::Tuple) = s[collect(k)]

@inline function broadcast_to_size(arg, rsize)
    if size(arg) != rsize
        collapse_dims = ccdims(size(arg))
        non_collapse_dims = ncdims(size(arg))
        if !isa(arg, XRTArray)
            arg = convert(XRTArray, arg)
        end
        if collapse_dims != ()
            arg = HloReshape(
                getindex_tuple(size(arg), non_collapse_dims)
            )(arg)
        end
        return HloBroadcast(map(x->x - 1, non_collapse_dims), rsize)(
            arg
        )
    else
        return arg
    end
end

# TODO: This should probably assert that the sessions are equivalent
any_sess(a::XRTArray, args::AnyXLA...) = a.storage.sess
any_sess(a::XLAScalar, args::AnyXLA...) = any_sess(args...)

function Broadcast.copy(bc::Broadcast.Broadcasted{<:XRTArrayStyle})
    ElType = Broadcast.combine_eltypes(bc.f, bc.args)
    bc′ = Broadcast.flatten(bc)
    if Base.isconcretetype(ElType)
        args = bc′.args
        # This could be axes(bc′) if we had better constant prop
        rsize = map(length, Broadcast.combine_axes(bc′.args...))
        args = map(arg->convert(XRTArray, broadcast_to_size(arg, rsize)), bc′.args)
        return HloMap{Core.Typeof(bc′.f)}()(bc′.f, args...)
    end
    # TODO: Pull back CPU, do this there
    error("No hope $(bc.f) $(typeof(bc.args)) $(ElType)")
end

using NNlib

@Base.pure function conv_windows(sz_, pad_, stride_, dilation_)
    ntuple(length(pad_)) do i
        (sz, p, s, d) = sz_[i], pad_[i], stride_[i], dilation_[i]
        WindowDims(sz, s, p, p, d, 1, true)
    end
end

function (::NNlib._conv{pad, stride, dilation})(input::XRTArray{T}, kernel::XRTArray{T}) where {T, pad, stride, dilation}
    pad_, stride_ = NNlib.padtuple(input, pad), NNlib.padtuple(input, stride)
    dilation_ = NNlib.padtuple(kernel, dilation)
    sz_ = size(kernel)
    windows = conv_windows(sz_, pad_, stride_, dilation_)
    convdims = ConvDimNums(
        3, 2, (0, 1),
        2, 3, (0, 1),
        3, 2, (0, 1)
    )
    HloConv(windows, convdims)(input, kernel)
end

function (::NNlib._∇conv_data{pad, stride, dilation, 0})(dy::XRTArray{T}, input::XRTArray{T}, kernel::XRTArray) where {T, pad, stride, dilation}
    pad_, stride_ = NNlib.padtuple(input, pad), NNlib.padtuple(input, stride)
    dilation_ = NNlib.padtuple(kernel, dilation)
    mirrored = kernel #permutedims(kernel, (2, 1, 3, 4)) #HloRev((0,1))(kernel)
    sz_ = size(kernel); isz = size(input); osz = size(dy)
    windows = ntuple(length(pad_)) do i
        (sz, p, s, d) = sz_[i], pad_[i], stride_[i], dilation_[i]
        padded_out_size = isz[i] + sz - 1
        pad_before = sz - 1 - p
        expanded_osz = (osz[i]-1)*s + 1
        pad_after = padded_out_size - expanded_osz - pad_before
        # N.B.: The window reversal flag is flipped here
        WindowDims(sz, 1, pad_before, pad_after, d, s, false)
    end
    convdims = ConvDimNums(
        3, 2, (0, 1),
        # N.B. The input and output dimensions are exchanged here from the
        # standard notion of convolution
        3, 2, (0, 1),
        3, 2, (0, 1)
    )
    r = HloConv(windows, convdims)(dy, mirrored)
    @assert size(r) == size(input)
    r
end

function (::NNlib._∇conv_filter{pad, stride, dilation, 0})(dy::XRTArray{T}, input::XRTArray{T}, kernel::XRTArray) where {T, pad, stride, dilation}
    # TODO: Validate that this is correct. Unfortunately, the NNlib
    # implementation itself is broken, so we need to fix that first
    pad_, stride_ = NNlib.padtuple(input, pad), NNlib.padtuple(input, stride)
    dilation_ = NNlib.padtuple(kernel, dilation)
    sz_ = size(kernel); isz = size(input); osz = size(dy)
    windows = ntuple(length(pad_)) do i
        (sz, p, s, d) = sz_[i], pad_[i], stride_[i], dilation_[i]
        expanded_osz = (osz[i]-1)*s + 1
        padded_in_size = expanded_osz + (sz - 1) * d
        pad_total = padded_in_size - isz[i]
        pad_before = max(ceil(Int, pad_total/2), 0)
        pad_after = pad_total - pad_before
        WindowDims(osz[i], d, pad_before, pad_after, s, 1, false)
    end
    convdims = ConvDimNums(
        # N.B. The input and output dimensions are exchanged here from the
        # standard notion of convolution
        2, 3, (0, 1),
        3, 2, (0, 1),
        2, 3, (0, 1)
    )
    r = HloConv(windows, convdims)(input, dy)
    r = HloRev((0,1))(r)
    @assert size(r) == size(kernel)
    r
end

function make_pooling_windows(x, k, pad, stride)
    k_, pad_, stride_ = NNlib.padtuple(x, k),
                        NNlib.padtuple(x, pad),
                        NNlib.padtuple(x, stride)
    windows = ntuple(ndims(x)) do i
        (sz, p, s) = i <= length(k) ? (k[i], pad[i], stride[i]) : (1, 0, 1)
        WindowDims(sz, s, p, p, 1, 1, false)
    end
end

function NNlib.maxpool(x::XRTArray, k; pad = map(_->0,k), stride = k)
    HloReduceWindow{typeof(max)}(
        make_pooling_windows(x, k, pad, stride)
    )(max, x, XRTArray(typemin(eltype(x))))
end

function NNlib.meanpool(x::XRTArray, k; pad = map(_->0,k), stride = k)
    s = HloReduceWindow{typeof(+)}(
        make_pooling_windows(x, k, pad, stride)
    )(+, x, XRTArray(zero(eltype(x))))
    wsize = convert(eltype(x), prod(k))
    s ./ XRTArray(wsize)
end

function NNlib.∇maxpool(dy::XRTArray, y::XRTArray, x::XRTArray, k; pad = map(_->0,k), stride = k)
    HloSelectAndScatter{typeof(>=), typeof(+)}(
        make_pooling_windows(x, k, pad, stride)
    )(>=, +, x, dy, XRTArray(zero(eltype(x))))
end

function make_pad_config(x, k, pad, stride)
    k_, pad_, stride_ = NNlib.padtuple(x, k),
                        NNlib.padtuple(x, pad),
                        NNlib.padtuple(x, stride)
    windows = ntuple(ndims(x)) do i
        (sz, p, s) = i <= length(k) ? (k[i], pad[i], stride[i]) : (1, 0, 1)
        #@assert p == 0 # TODO: What does this do? Probably reduce the edge padding
        PaddingConfigDim(sz - 1, sz - 1, s - 1)
    end
end

function NNlib.∇meanpool(dy::XRTArray, y::XRTArray, x::XRTArray, k; pad = map(_->0,k), stride = k)
    # General algorithm for ∇meanpool.
    # The idea here is relatively simple. Without loss of generality, we can ignore the
    # division by a constant - i.e. we're looking at the gradient of a (potentially,
    # padded and strided sum of a window). These windows can be overlapping, so in general,
    # we once again need to do a summed reduction of all the gradients. We do so as follows:
    # Suppose we have a 3x3 array:
    #
    #   123
    #   456
    #   789
    #
    # And we sum in a 2x2 window, i.e.
    # our gradients look like:
    #
    #   ab
    #   cd
    #
    # with a=1+2+4+5, c=4+5+7+8, etc. To compute, the gradient, we pad up to
    #
    #  0000
    #  0ab0
    #  0cd0
    #  0000
    #
    # and once again sum with a 2x2 window. It is easy to check that this procedure gives
    # the correct answer. If we have a stride, we need to add interior padding in addition.
    # Another way to think of this and arrive at the same result is to see
    # the summation as a convolution with an all-ones kernel.
    #
    # TODO: Does it make sense to special case to the case of non-overlapping windows
    # or does XLA take care of that?
    wsize = convert(eltype(x), prod(k))
    dy = dy ./ XRTArray(wsize)
    pdy = HloPad(make_pad_config(x, k, pad, stride))(dy, XRTArray(zero(eltype(dy))))
    HloReduceWindow{typeof(+)}(
        make_pooling_windows(x, k, pad, map(_->1, k))
    )(+, pdy, XRTArray(zero(eltype(x))))
end

@Base.pure dims_tuple(n) = tuple((0:n-1)...)
@Base.pure rev_dims_tuple(n) = tuple((n-1:-1:0)...)
dims_tuple(A, ::Colon) = dims_tuple(ndims(A))
dims_tuple(A, t::Tuple) = map(x->x-1, t)
dims_tuple(A, n::Int) = (n-1,)

@inline function Base.reshape(A::XRTArray, dims::Tuple{Vararg{Union{Int,Colon}}})
    reshape(A, Base._reshape_uncolon(A, dims))
end

@Base.pure rev_tuple(n) = reverse(n)

@inline function Base.reshape(A::XRTArray, dims::Tuple{Vararg{Int}})
    prod(dims) == prod(size(A)) || Base._throw_dmrsa(dims, prod(size(A)))
    HloTranspose(rev_dims_tuple(length(dims)))(
       HloReshape(rev_tuple(dims))(
            # HLO reshape semantics collapse the opposite way
            HloTranspose(rev_dims_tuple(ndims(A)))(A)
       )
    )
end

@Base.pure reduced_dimensions_collapes(sz, dims) = ntuple(i->i in dims ? 1 : sz[i], length(sz))
function Base.mapreduce(f, op, A::XRTArray; dims=:)
    dt = dims_tuple(A, dims)
    res = HloReduce{Core.Typeof(op)}(dt)(op,
        HloMap{Core.Typeof(f)}()(f, A),
        XRTArray(zero(eltype(A)))
    )
    if dims != (:)
        # Put back the dimensions that HloReduce dropped;
        # Julia semantics require this.
        res = HloReshape(reduced_dimensions_collapes(size(A), dims))(res)
    end
    return res
end

using Flux

function Flux.rand_similar(x::XRTArray{T, Shape}) where {T, Shape}
    HloRng(T, Shape, xla.RandomDistribution.RNG_UNIFORM)(
        XRTArray(zero(T)),
        XRTArray(one(T)))
end

using LinearAlgebra
# TODO: Base scales by maxabs if necessary. Do we need that?
function LinearAlgebra.dot(A::XRTVector, B::XRTVector)
    ddots = DimNums((0,), (0,), (), ())
    HloDot(ddots)(A, B)
end
LinearAlgebra.norm2(A::XRTArray, p) = sqrt(dot(A, A))
function LinearAlgebra.norm(A::XRTArray, p::Real)
    if p == 2
        return LinearAlgebra.norm2(A, p)
    end
    error("Not implemented")
end

function Base.setindex(A::XRTVector{T}, v::T, i::XRTArray{<:Integer, (), 0}) where {T}
    Base.setindex(A, XRTArray{T, (), 0}(v), i)
end
function Base.setindex(A::XRTVector{T}, v::XRTArray{T, (), 0}, i::XRTArray{<:Integer, (), 0}) where {T}
    # TODO: It would be better to handle the dim mapping internally,
    # but that requires access to the julia element type, which we don't currently
    # expose
    vb = HloBroadcast(ntuple(i->i-1, ndims(T)), ntuple(_->1, ndims(A)))(v)
    inds = HloBroadcast((), (1,))(i - 1)
    if ndims(T) != 0
        ones = HloBroadcast((), (ndims(T),))(XRTArray(1))
        inds = HloConcatenate(0)(ones, inds)
    end
    HloDynamicUpdateSlice()(A, vb, inds)
end
Base.setindex(A::XRTVector, v, i::Int64) = Base.setindex(A, v, XRTArray{Int64, (), 0}(i))

# TODO: Support trailing singleton
@Base.pure start_idxs_tuple(rs) = ntuple(i->isa(rs[i], Colon) ? 0 : first(rs[i]) - 1, Val(length(rs)))
function Base.setindex(A::XRTArray{T,<:Any,N}, up::XRTArray{T,<:Any,N}, rs::Vararg{Union{UnitRange{<:Integer}, Colon}, N}) where {T, N}
    inds = map(x->HloBroadcast((), (1,))(XRTArray(x)), start_idxs_tuple(rs))
    HloDynamicUpdateSlice()(A, up, HloConcatenate(0)(inds...))
end

@Base.pure function compute_result_size(tA, tInds)
    map(i->size(tA, i),
        tuple(Iterators.filter(
            i->tInds.parameters[i] <: Colon, 1:length(tInds.parameters))...))
end

function Base.getindex(A::XRTArray{T, <:Any, N}, inds::Vararg{Union{Colon, XRTArray{<:Integer, (), 0}}, N}) where {T, N}
    # TODO: It would be better to handle the dim mapping internally,
    # but that requires access to the julia element type, which we don't currently
    # expose
    ind_vector = HloConcatenate(0)(ntuple(Val(length(inds))) do i
        if isa(inds[i], Colon)
            HloBroadcast((), (1,))(XRTArray(0))
        elseif isa(inds[i], XRTArray)
            HloBroadcast((), (1,))(inds[i] - 1)
        end
    end...)
    sizes = ntuple(Val(length(inds))) do i
        isa(inds[i], Colon) ? size(A, i) : 1
    end
    if ndims(T) != 0
        ones = HloBroadcast((), (ndims(T),))(XRTArray(1))
        ind_vector = HloConcatenate(0)(ones, ind_vector)
        sizes = (size(T)..., sizes...)
    end
    result_sizes = compute_result_size(typeof(A), typeof(inds))
    ret = HloDynamicSlice(sizes)(A, ind_vector)
    ret = HloReshape(tuple((ndims(T) == 0 ? () : size(T))..., result_sizes...))(ret)
    if result_sizes === () && !(T <: XLAScalar)
        return convert(T, ret)
    end
    return ret
end
Base.getindex(A::XRTVector, i::Int64) = Base.getindex(A, XRTArray{Int64, (), 0}(i))
# Override for scalars
Base.getindex(x::XRTArray{T,(),0}) where {T} = x

@Base.pure function compute_gather_tuples(A, idxs)
    array_idxs = tuple(Iterators.filter(i->!(idxs.parameters[i] <: Colon), 1:length(idxs.parameters))...)
    offset_dims = tuple(Iterators.filter(i->idxs.parameters[i+1] <: Colon, 0:length(idxs.parameters)-1)...)
    start_index_map = tuple(Iterators.filter(i->idxs.parameters[i+1] <: XRTArray, 0:length(idxs.parameters)-1)...)
    slice_sizes = map(i->size(A, i), tuple(filter(i->idxs.parameters[i] <: Colon, 1:length(idxs.parameters))...))
    slice_sizes = tuple(slice_sizes..., (1 for _ in 1:length(array_idxs))...)
    (array_idxs, offset_dims, start_index_map, slice_sizes)
end

function Base.getindex(A::XRTArray{<:Any, <:Any, N}, idxs::Vararg{Union{Colon, XRTArray{<:Integer, <:Any, 1}}, N}) where {N}
    (array_idxs_idxs, offset_dims, start_index_map, slice_sizes) = compute_gather_tuples(typeof(A), typeof(idxs))
    array_idxs = ntuple(i->idxs[array_idxs_idxs[i]], Val{length(array_idxs_idxs)}())
    M = length(array_idxs)
    if length(array_idxs) == 1
        gather_idxs = array_idxs[1] .- XRTArray(one(eltype(array_idxs[1])))
    else
        gather_shape = map(length, array_idxs)
        gather_idxs = HloConcatenate(M)(ntuple(Val{M}) do i
            HloBroadcast((i,), (gather_shape..., 1))(array_idxs[i] .- XRTArray(one(eltype(array_idxs[i]))))
        end...)
    end
    dims = GatherDims(
        offset_dims,
        start_index_map, #= collapsed_slice_dims =#
        start_index_map,
        M #= index_vector_dim =#
    )
    HloGather(dims, slice_sizes)(A, gather_idxs)
end

@noinline function compute_slice_dim(rs, sz, i)
    r = rs[i]
    isa(r, UnitRange) ?
        SliceDim(first(r)-1, last(r), step(r)) :
        SliceDim(0, sz[i], 1)
end

@Base.pure slices_tuple(rs, sz) = ntuple(i->compute_slice_dim(rs, sz, i), Val(length(sz)))

function Base.getindex(A::XRTArray{<:Any, <:Any, N}, rs::Vararg{Union{UnitRange{<:Integer}, Colon}, N}) where {N}
    HloSlice(slices_tuple(rs, size(typeof(A))))(A)
end

using DiffRules
DiffRules.select(p::XRTArray{Bool, (), 0}, x::XRTArray{T, (), 0}, y::XRTArray{T, (), 0}) where {T} =
    GenericHloOp{:select}(T, ())(p, x, y)
DiffRules.select(p::Bool, x::XRTArray{T, (), 0}, y::XRTArray{T, (), 0}) where {T} =
    GenericHloOp{:select}(T, ())(XRTArray{Bool, (), 0}(p), x, y)

function XRTArray(a::UnitRange{T}) where {T}
    HloIota(T, (length(a),), 0)() .+ first(a)
end

function XRTArray(a::Transpose)
    make_hlo_transpose(XRTArray(a.parent))
end

function change_eltype(T::Type, x::XRTArray)
    GenericHloOp{:convert}(T, size(x))(x)
end
change_eltype(::Type{T}, x::XRTArray{T}) where {T} = x
Base.convert(::Type{XRTArray{T}}, x::XRTArray{S}) where {T,S} = change_eltype(T, x)
Base.convert(::Type{XRTArray{T}}, x::XRTArray{T}) where {T} = x
Base.convert(::Type{XRTArray}, x::XRTArray) = x

for T in Base.uniontypes(XLA.XLAScalar)
    @eval (::Type{$T})(x::XRTScalar{S}) where {S<:XLAScalar} = convert(XRTArray{$T}, x)
    @eval (::Type{$T})(x::XRTScalar{$T}) where {S<:XLAScalar} = x
end

@Base.pure @noinline function padding_tuple(aiT::Type{<:XRTArray}, dims::Int)
    ntuple(_->1, Val(dims))
end

struct ntupler{dims, T}
    a::T
end
function (nt::ntupler{dims, T})(i::Int) where {dims, T}
    ai = nt.a[i]
    if ndims(ai) < dims
        pt = ntuple(_->1, dims-ndims(ai))
        HloBroadcast(ntuple(i->i-1, ndims(ai)), (size(ai)...,
            pt...))(ai)
    else
        ai
    end
end

function _cat(dims::Int, a::XRTArray{T}...) where {T}
    b = ntuple(ntupler{dims, typeof(a)}(a), length(a))
    HloConcatenate(dims-1)(b...)
end
function _cat(dims::Int, a::XRTArray...)
    let eT = Base.promote_eltype(a...)
        _cat(dims, map(a->change_eltype(eT, a), a)...)
    end
end
Base.hcat(a::XRTArray...) = _cat(2, a...)
Base.vcat(a::XRTArray...) = _cat(1, a...)
Base.cat(a::XRTArray...; dims) = _cat(dims, a...)

using Statistics
@Base.pure @noinline count_summands(T::Type{<:XRTArray}, dims) = prod(size(T)[[dims...]])
@Base.pure @noinline count_summands(T::Type{<:XRTArray}, dims::Colon) = prod(size(T))
dimsum(A, dims) = sum(A; dims=dims)
function Statistics.mean(A::XRTArray; dims=:)
    summed = dimsum(A, dims)
    x = convert(eltype(A), count_summands(typeof(A), dims))
    nsummands = XRTArray(x)
    summed ./ nsummands
end

Base.@pure @noinline function compute_dropdims_d(A, dims)
    for i in 1:length(dims)
        1 <= dims[i] <= ndims(A) || throw(ArgumentError("dropped dims must be in range 1:ndims(A)"))
        length(axes(A, dims[i])) == 1 || throw(ArgumentError("dropped dims must all be size 1"))
        for j = 1:i-1
            dims[j] == dims[i] && throw(ArgumentError("dropped dims must be unique"))
        end
    end
    d = ()
    for i = 1:ndims(A)
        if !in(i, dims)
            d = tuple(d..., axes(A, i))
        end
    end
    return d
end
function Base._dropdims(A::XRTArray, dims::Base.Dims)
    reshape(A, compute_dropdims_d(typeof(A), dims))
end

@Base.pure function make_repeat_dims(sz, rep_dims)
    padded_sz = ntuple(i->i <= length(sz) ? sz[i] : 1, length(rep_dims))
    dim_mapping = tuple((2*i+1 for i = 0:length(sz)-1)...)
    result_szs = tuple(Iterators.flatten(zip(rep_dims, padded_sz))...)
    final_sz = tuple((a*b for (a,b) in zip(rep_dims, padded_sz))...)
    dim_mapping, result_szs, final_sz
end

function Base.repeat(A::XRTArray, dims::Integer...)
    dims_mapping, result_szs, final_sz = make_repeat_dims(size(A), dims)
    XLA.HloReshape(final_sz)(XLA.HloBroadcast(dims_mapping, result_szs)(A))
end


"""
    SplitLayer(xs...)

A layer that calls a number of sublayers/mappings with the same input and
returns a tuple of their outputs. Can be used in a regular Flux model:

```julia-repl
julia> m = Chain(Dense(2,3), SplitLayer(Dense(3,2), x->x .* 2))
julia> length(params(m)) == 4
julia> (x,y) = m(rand(2))
(Float32[-1.0541434, 1.1694773], Float32[-3.1472511, -0.86115724, -0.39665926])
```
Comes with a convenient constructor for a SplitLayer built from Dense layers
with given activation(s):
```julia-repl
julia> m = Chain(Dense(2,3), SplitLayer(3, [2,5], σ))
julia> (x,y) = m(rand(2))
(Float32[0.3069554, 0.3362006], Float32[0.437131, 0.4982477, 0.6465078, 0.4523438, 0.5068563])
```

You can also provide just a vector / scalar that should be trained but have the
same value for all inputs (like a lonely bias vector). This functionality is
provided by the `TrainableVector`/`TrainableScalar` types. For vector inputs
they simply return the array they are wrapping. For matrix (i.e. batch) inputs
they return appropriately repeated arrays:
```julia-repl
julia> m = SplitLayer(Dense(2,3), ones(Float32,3))
julia> length(params(m)) == 3
julia> (x,y) = m(rand(2,5))
julia> size(y) == (3,5)
julia> y
3×3 Array{Float32,2}:
 1.0  1.0  1.0
 1.0  1.0  1.0
 1.0  1.0  1.0
```
"""
struct SplitLayer{T<:Tuple}
    layers::T
end

SplitLayer(layers...) = SplitLayer(map(maybe_trainable, layers))

function (m::SplitLayer)(x)
    Tuple(layer(x) for layer in m.layers)
end

@functor SplitLayer

# for use as e.g. shared variance
struct TrainableVector{T<:AbstractArray}
    v::T
end
(v::TrainableVector)(x::AbstractVector) = v.v
(v::TrainableVector)(x::AbstractMatrix) = v.v .* reshape(fillsimilar(v.v,size(x,ndims(x)),1),1,:)
(v::TrainableVector)() = v.v
@functor TrainableVector

# for use as e.g. shared variance
struct TrainableScalar{T<:Real}
    s::AbstractVector{T}
    TrainableScalar{T}(x::T) where T<:Real = new{T}([x])
end
TrainableScalar(x::T) where T<:Real = TrainableScalar{T}(x)
(s::TrainableScalar)(x::AbstractVector) = s.s[1]
(s::TrainableScalar)(x::AbstractMatrix) = fillsimilar(x,size(x,ndims(x)),1) .* s.s
@functor TrainableScalar

maybe_trainable(x) = x
maybe_trainable(x::AbstractArray) = TrainableVector(x)
maybe_trainable(x::Real) = TrainableScalar(x)

fillsimilar(x::AbstractArray, s::Tuple, value::Real) = fill!(similar(x, s...), value)
fillsimilar(x::AbstractArray, s, value::Real) = fill!(similar(x, s), value)
@non_differentiable fillsimilar(::Any, ::Any, ::Any)


# Run model fitting and crossvalidation
# for real or simulated data

# Main runner for model fitting and model comparison
# NOTE: These procedures take a long time to run.
# The settings are different from the ones used in the paper (DEBUG settings).
# Comment/Uncomment the corresponding lines (around 383 and 445)
# to perform the fitting and comparison procedures as in the paper.
function runner_crossval_multipletimes_fit()

    # -------- Fit (once)
    i=""
    result_path_base, batch_fit_models_spec = model_list_main(filestr = string(i))
    runner_fmri_fit(result_path_base, batch_fit_models_spec)

    # -------- Crossvalidation (multiple runs)
    nruns = 5
    for i in 1:nruns
        result_path_base, batch_fit_models_spec = model_list_main(filestr = string(i))
        runner_fmri_3foldscrossval(result_path_base,
                                        batch_fit_models_spec)
    end
end

# Main runner for model and parameter recovery
# NOTE: You need to EDIT which data (real or simulated)
# should be used in
# runner_fmri_fit and runner_fmri_3foldscrossval!!!!
# Right now the real participants' data are used.
# See lines 367 and 428
function runner_multipletimes_fit_recovery()

    # -------- Fit for parameter recovery
    # ---- We run fit multiple times to get some error bars
    nruns = 5
    for i in 1:nruns
        result_path_base, batch_fit_models_spec = model_list_paramrecovery(filestr = string(i))
        runner_fmri_fit(result_path_base, batch_fit_models_spec)

    end

    # -------- Crossvalidation (multiple runs) for model recovery
    nruns = 5
    for i in 1:nruns
        result_path_base, batch_fit_models_spec = model_list_main(filestr = string(i))
        runner_fmri_3foldscrossval(result_path_base,
                                        batch_fit_models_spec)
    end
end

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# ------- Model list -----------------------------------------------------------
# ------------------------------------------------------------------------------
# List and specifications of algorithms for Fig. 3 of paper
function model_list_main(; filestr = "_")

    result_path_base = makehomesavepath("test_main"*filestr)

    stochmax = 20.

    # # Biased Random Walk
    # --------------------
    biasedRandomWalk_prior = [
        RL_Fit.paramBias(0.0, 1.)
    ]
    biasedRandomWalk_prior_std = [.15]
    biasedRandomWalk_spec = RL_Fit.Model_spec_per_subj("Biased Random Walk (0-Model)",
                                RL_Fit.getLL_BiasedRandomWalk_V,
                                biasedRandomWalk_prior, biasedRandomWalk_prior_std,
                                string(result_path_base, "biased_random_walk/"))
    # Sarsa-λ
    # -------
    sarsaL_prior = [
        RL_Fit.paramAlpha(0.0, 1.0),
        RL_Fit.paramGamma(0.0, 1.0),
        RL_Fit.paramLambda(0.0, 1.0),
        RL_Fit.paramTemp(0.01, 10.)
        ]
    sarsaL_prior_std = [.15, .15, .15, 1.5]
    sarsaL_spec = RL_Fit.Model_spec_per_subj("Sarsa Lambda",
                    RL_Fit.getLL_SarsaLambda_V,
                    sarsaL_prior, sarsaL_prior_std,
                    string(result_path_base, "sarsa_lambda/"))
    # Prioritized Sweeping with Particle Filtering
    # --------------------------------------------
    priosweep_particle_prior = [
        RL_Fit.paramGamma(0.0, 1.0),
        RL_Fit.paramTemp(0.01, 10.),
        RL_Fit.paramSurpriseProbability(0.0, 1.),
        RL_Fit.paramStochasticity(0.00001, stochmax),
        RL_Fit.paramNrUpdateCycles(0, 9),
        RL_Fit.paramNParticles(1, 20)
        ]
    priosweep_particle_prior_std = [.15, 1.5, .15, 3., 1.35, 3.]
    priosweep_particle_spec = RL_Fit.Model_spec_per_subj("Prioritized Sweeping Particle Filter",
                                RL_Fit.getLL_Prio_Sweep_TParticleFilterJump_V,
                                priosweep_particle_prior,
                                priosweep_particle_prior_std,
                                string(result_path_base, "priosweep_particle/"))
    # Hybrid-λ-PS-PF
    # --------------
    hybrid_pspf_sarsaL_prior = [
        RL_Fit.paramGamma(0.0, 1.0),
        RL_Fit.paramSurpriseProbability(0.0, 1.),
        RL_Fit.paramStochasticity(0.00001, stochmax),
        RL_Fit.paramNrUpdateCycles(0, 9),
        RL_Fit.paramNParticles(1, 20),
        RL_Fit.paramAlpha(0.0, 1.),
        RL_Fit.paramLambda(0.0, 1.0),
        RL_Fit.paramWeight(0.0, 20.0; name="w_MB"),
        RL_Fit.paramWeight(0.0, 20.0; name="w_MF")
        ]
    hybrid_pspf_sarsaL_prior_std = [.15, .15, 3., 1.35, 3., .15, .15, 3., 3.]
    hybrid_pspf_sarsaL_spec = RL_Fit.Model_spec_per_subj("Hybrid: PS with Particle Filtering + Sarsa Lambda",
                                RL_Fit.getLL_Hybrid_Prio_Sweep_TParticleFilterJump_V,
                                hybrid_pspf_sarsaL_prior, hybrid_pspf_sarsaL_prior_std,
                                string(result_path_base, "hybrid_pspf_sarsaL/"))
    # Actor-Critic
    # ------------
    actorCriticTD_prior = [
        RL_Fit.paramAlpha(0.0, 1.; name="learning rate actor"),
        RL_Fit.paramEta(0.0, 1.; name="learning rate critic"),
        RL_Fit.paramGamma(0.0, 1.),
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace actor"), # actor
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace critic"), # critic
        RL_Fit.paramTemp(0.01, 10.)
    ]
    actorCriticTD_prior_std = [.15, .15, .15, .15, .15, 1.5]
    actorCriticTD_spec = RL_Fit.Model_spec_per_subj("Actor-Critic (TD-Lambda)",
                                RL_Fit.getLL_ActorCritic_TDLambda_V,
                                actorCriticTD_prior, actorCriticTD_prior_std,
                                string(result_path_base, "actor_critic_TD/"))
    # REINFORCE
    # ---------
    REINFORCE_prior = [
        RL_Fit.paramAlpha(0.0, 1.),
        RL_Fit.paramGamma(0.0, 1.),
        RL_Fit.paramTemp(0.01, 10.)
    ]
    REINFORCE_prior_std = [.15, .15, 1.5]
    REINFORCE_spec = RL_Fit.Model_spec_per_subj("REINFORCE (Williams 1992)",
                            RL_Fit.getLL_REINFORCE_V,
                            REINFORCE_prior, REINFORCE_prior_std,
                            string(result_path_base, "REINFORCE/"))
    # Hybrid Actor-Critic
    # -------------------
    hybrid_actorCritic_FWD_prior = [
        RL_Fit.paramAlpha(0.0, 1.; name="learning rate actor"),
        RL_Fit.paramAlpha(0.0, 1.; name="learning rate critic"),
        RL_Fit.paramGamma(0.0, 1.),
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace actor"), # actor
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace critic"), # critic
        RL_Fit.paramEta(0., 1.; name="learning rate fwd"),
        RL_Fit.paramWeight(0.0, 20.0; name="w_MB"),
        RL_Fit.paramWeight(0.0, 20.0; name="w_MF")
    ]
    hybrid_actorCritic_FWD_prior_std = [.15, .15, .15, .15, .15, .15, 3., 3.]
    hybrid_actorCritic_FWD_spec = RL_Fit.Model_spec_per_subj("Hybrid Actor-Critic with FWD",
                                RL_Fit.getLL_Hybrid_AC_FWD_V,
                                hybrid_actorCritic_FWD_prior, hybrid_actorCritic_FWD_prior_std,
                                string(result_path_base, "hybrid_actorcritic_FWD/"))

    # Forward Learner (MB only learner)
    # ---------------------------------
    forwardLearner_prior = [
        RL_Fit.paramEta(0., 1.),
        RL_Fit.paramGamma(0., 1.),
        RL_Fit.paramTemp(0.01, 10.)
    ]
    forwardLearner_prior_std = [.15, .15, 1.5]
    forwardLearner_spec = RL_Fit.Model_spec_per_subj("Forward Learner",
                                RL_Fit.getLL_Forward_Learner_V,
                                forwardLearner_prior, forwardLearner_prior_std,
                                string(result_path_base, "forward_learner/"))
    # Hybrid-0 (Glaescher et al. 2010)
    # --------------------------------
    glaescherHybrid_prior = [
        RL_Fit.paramAlpha(0., 1.),
        RL_Fit.paramGamma(0., 1.),
        RL_Fit.paramEta(0., 1.),
        RL_Fit.paramExpDecayOffset(0., 1.),
        RL_Fit.paramExpDecaySlope(-.1, +.8),
        RL_Fit.paramTemp(0.01, 10.)
    ]
    glaescherHybrid_prior_std = [.15, .15, .15, .15, .15, 1.5]
    glaescherHybrid_spec = RL_Fit.Model_spec_per_subj("Glaescher-Hybrid",
                                RL_Fit.getLL_GlaescherHybrid_V,
                                glaescherHybrid_prior, glaescherHybrid_prior_std,
                                string(result_path_base, "glaescher_hybrid/"))
    # Hybrid-λ (Daw et al. 2011)
    # --------------------------
    hybridSarsaL_prior = [
        RL_Fit.paramAlpha(0.0, 1.),
        RL_Fit.paramGamma(0., 1.),
        RL_Fit.paramLambda(0.0, 1.),
        RL_Fit.paramEta(0., 1.),
        RL_Fit.paramExpDecayOffset(0., 1.),
        RL_Fit.paramExpDecaySlope(0., 1.),
        RL_Fit.paramTemp(0.01, 10.)
    ]
    hybridSarsaL_prior_std = [.15, .15, .15, .15, .15, .15, 1.5]
    hybridSarsaL_prior_spec = RL_Fit.Model_spec_per_subj("Hybrid-SarsaLambda",
                                RL_Fit.getLL_HybridSarsaLambda_V,
                                hybridSarsaL_prior, hybridSarsaL_prior_std,
                                string(result_path_base, "hybrid_FWD_SarsaLambda/"))

    # REINFORCE with baseline
    # -----------------------
    REINFORCE_baseline_prior = [
        RL_Fit.paramAlpha(0.0, 1., name="learning rate"),
        RL_Fit.paramGamma(0.0, 1.),
        RL_Fit.paramTemp(0.01, 10.),
        RL_Fit.paramAlpha(0.0, 1.; name="learning rate baseline")
    ]
    REINFORCE_baseline_prior_std = [.15, .15, 1.5, .15]
    REINFORCE_baseline_spec = RL_Fit.Model_spec_per_subj("REINFORCE (Williams 1992) baseline",
                            RL_Fit.getLL_REINFORCE_baseline_V,
                            REINFORCE_baseline_prior, REINFORCE_baseline_prior_std,
                            string(result_path_base, "REINFORCE_baseline/"))
    # Surprise Sarsa-λ
    # ----------------
    sarsaL_modulated_prior = [
        RL_Fit.paramGamma(0.0, 1.0),
        RL_Fit.paramTemp(0.01, 10.),
        RL_Fit.paramSurpriseProbability(0.0, 1.),
        RL_Fit.paramStochasticity(0.00001, stochmax),
        RL_Fit.paramNParticles(1, 20),
        RL_Fit.paramAlpha(0.0, 1.0),
        RL_Fit.paramLambda(0.0, 1.0)
        ]
    sarsaL_modulated_prior_std = [.15, 1.5, .15, 3., 3., .15, .15]
    sarsaL_modulated_spec = RL_Fit.Model_spec_per_subj("Sarsa Lambda modulated by Particle Filtering",
                                RL_Fit.getLL_sarsaL_modulated_TParticleFilterJump_V,
                                sarsaL_modulated_prior,
                                sarsaL_modulated_prior_std,
                                string(result_path_base, "sarsaL_modulated/"))

    # Surprise Actor-Critic
    # ---------------------
    actorCriticTD_modulatedActor_prior = [
        RL_Fit.paramAlpha(0.0, 1.; name="learning rate actor"),
        RL_Fit.paramSurpriseProbability(0.0, 1.),
        RL_Fit.paramStochasticity(0.00001, stochmax),
        RL_Fit.paramEta(0.0, 1.; name="learning rate critic"),
        RL_Fit.paramGamma(0.0, 1.),
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace actor"), # actor
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace critic"), # critic
        RL_Fit.paramTemp(0.01, 10.),
        RL_Fit.paramNParticles(1, 20)
    ]
    actorCriticTD_modulatedActor_prior_std = [.15, .15, 3., .15, .15, .15, .15, 1.5, 3.]
    actorCriticTD_modulatedActor_spec = RL_Fit.Model_spec_per_subj("Actor-Critic (TD-Lambda) modulated actor (1 degree)",
                                RL_Fit.getLL_ActorCritic_TDLambda_modulatedActor_V,
                                actorCriticTD_modulatedActor_prior, actorCriticTD_modulatedActor_prior_std,
                                string(result_path_base, "actor_critic_TD_modulatedActor_1degree/"))
    # Surprise REINFORCE
    # ------------------
    REINFORCE_modulated_prior = [
        RL_Fit.paramAlpha(0.0, 1.),
        RL_Fit.paramSurpriseProbability(0.0, 1.),
        RL_Fit.paramStochasticity(0.00001, stochmax),
        RL_Fit.paramGamma(0.0, 1.),
        RL_Fit.paramTemp(0.01, 10.),
        RL_Fit.paramNParticles(1, 20)
    ]
    REINFORCE_modulated_prior_std = [.15, .15, 3., .15, 1.35, 3.]
    REINFORCE_modulated_spec = RL_Fit.Model_spec_per_subj("REINFORCE (Williams 1992) modulated",
                            RL_Fit.getLL_REINFORCE_modulated_V,
                            REINFORCE_modulated_prior, REINFORCE_modulated_prior_std,
                            string(result_path_base, "REINFORCE_modulated/"))

    # Hybrid Actor-Critic PF
    # ----------------------
    hybrid_actorCritic_FWD_BF_prior = [
        RL_Fit.paramAlpha(0.0, 1.; name="learning rate actor"),
        RL_Fit.paramAlpha(0.0, 1.; name="learning rate critic"),
        RL_Fit.paramGamma(0.0, 1.),
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace actor"), # actor
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace critic"), # critic
        RL_Fit.paramSurpriseProbability(0.0, 1.),
        RL_Fit.paramStochasticity(0.00001, stochmax),
        RL_Fit.paramNParticles(1, 20),
        RL_Fit.paramWeight(0.0, 20.0; name="w_MB"),
        RL_Fit.paramWeight(0.0, 20.0; name="w_MF")
    ]
    hybrid_actorCritic_FWD_BF_prior_std = [.15, .15, .15, .15, .15, .15, 3., 3., 3., 3.,]
    hybrid_actorCritic_FWD_BF_spec = RL_Fit.Model_spec_per_subj("Hybrid Actor-Critic with FWD - BF",
                                RL_Fit.getLL_Hybrid_AC_FWD_BF_V,
                                hybrid_actorCritic_FWD_BF_prior, hybrid_actorCritic_FWD_BF_prior_std,
                                string(result_path_base, "hybrid_actorcritic_FWD_BF/"))
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # --- Run for the paper:
    batch_fit_models_spec = [
                            actorCriticTD_spec,
                            actorCriticTD_modulatedActor_spec,
                            hybrid_actorCritic_FWD_BF_spec,
                            hybrid_actorCritic_FWD_spec,
                            REINFORCE_spec,
                            REINFORCE_modulated_spec,
                            REINFORCE_baseline_spec,
                            sarsaL_spec,
                            sarsaL_modulated_spec,
                            forwardLearner_spec,
                            priosweep_particle_spec,
                            hybridSarsaL_prior_spec,
                            glaescherHybrid_spec,
                            hybrid_pspf_sarsaL_spec,
                            biasedRandomWalk_spec
                          ]
    result_path_base, batch_fit_models_spec
end

# List and specifications of algorithm for parameter recovery
function model_list_paramrecovery(; filestr = "_")

    result_path_base = makehomesavepath("test_main_recovery"*filestr)

    stochmax = 20.

    # Surprise Actor-Critic
    # ---------------------
    actorCriticTD_modulatedActor_prior = [
        RL_Fit.paramAlpha(0.0, 1.; name="learning rate actor"),
        RL_Fit.paramSurpriseProbability(0.0, 1.),
        RL_Fit.paramStochasticity(0.00001, stochmax),
        RL_Fit.paramEta(0.0, 1.; name="learning rate critic"),
        RL_Fit.paramGamma(0.0, 1.),
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace actor"), # actor
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace critic"), # critic
        RL_Fit.paramTemp(0.01, 10.),
        RL_Fit.paramNParticles(1, 20)
    ]
    actorCriticTD_modulatedActor_prior_std = [.15, .15, 3., .15, .15, .15, .15, 1.5, 3.]
    actorCriticTD_modulatedActor_spec = RL_Fit.Model_spec_per_subj("Actor-Critic (TD-Lambda) modulated actor (1 degree)",
                                RL_Fit.getLL_ActorCritic_TDLambda_modulatedActor_V,
                                actorCriticTD_modulatedActor_prior, actorCriticTD_modulatedActor_prior_std,
                                string(result_path_base, "actor_critic_TD_modulatedActor_1degree/"))

    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # --- Run for parameter recovery:
    batch_fit_models_spec = [
                            actorCriticTD_modulatedActor_spec,
                          ]

    result_path_base, batch_fit_models_spec
end

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# ------- MCMC settings and helpers --------------------------------------------
# ------------------------------------------------------------------------------

# Settings for data fitting
# Same settings as in runner_fmri_3foldscrossval (see below),
# but here: only fit, not crossvalidation
function runner_fmri_fit(result_path_base,
                                        batch_fit_models_spec
                                        # ; filestr = "_")
                                        )
    rl_module_path = "hi"
    # --------------------------------------------------------------------------
    # ------ Which data? Real or simulated? Edit!

    # Real participants' data:
    data_path = "./mcmc_rl_fit/projects/fmri/data/SARSPEIC_all_fMRI.csv"

    # Simulated data:
    # data_path = "./mcmc_rl_fit/projects/fmri/data_sim/SARSPEICZVG_all_fMRI_Sim_exporder_SurpAC_run2.csv"

    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    description = "Batch fit: Models / subjects"
    # ---- DEBUG settings
    max_nr_episodes = 50  # FIT to the first 10 episodes. For the main analysis, I always used 40
    max_nr_episodes = 2  # FIT to the first 10 episodes. For the main analysis, I always used 40
    nr_repetitions = 2
    top_n_fraction = 0.1
    nr_burnin = 3
    nr_samples = 3
    inter_collect_interval = 4
    # # end debug settings
    # ---------------------------
    # ---- Settings used for paper
    # max_nr_episodes = 50  # FIT to the first n episodes.
    # nr_repetitions = 50 #20 # one repetition generates nr_mcmc_samples. doing several repetitions, starting from different random initial positions helps finding isolated regions (e.g. bimodal distributions).
    # top_n_fraction = 0.001 # not important. I sometimes a the top 1percent of mcmc samples
    # nr_burnin = 1500 # important. mcmc needs to "burn in". The more params (larger space) the longer it takes to walk into the volume of interest.
    # nr_samples = 2000 # how many samples to collect from each independent repetition
    # inter_collect_interval = 10 # samples are correlated. therefore make N steps before collecting a sample. 7 to 10 seems a good value.
    # # end
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------

    fit_spec = RL_Fit.FitSpec_multiple_models(
        batch_fit_models_spec,
        max_nr_episodes,
        nr_burnin,
        nr_samples,
        inter_collect_interval,
        top_n_fraction,
        nr_repetitions,
        # false,  #do_per_subject_fit
        true,  #do_per_subject_fit
        result_path_base
    )
    # --------------------------------------------------------------------------
    # ----------------- RUN -----------------
    RL_Fit.run_fit_multiple_models(description, data_path, fit_spec, rl_module_path)

end

# Settings for crossvalidation
# Same settings as in runner_fmri_fit (see above),
# but here: ONLY crossvalidation, not fit
function runner_fmri_3foldscrossval(result_path_base,
                                        batch_fit_models_spec
                                        )
    rl_module_path = "hi"
    # --------------------------------------------------------------------------
    # ------ Which data? Real or simulated? Edit!

    # Real participants' data:
    data_path = "./mcmc_rl_fit/projects/fmri/data/SARSPEIC_all_fMRI.csv"

    # Simulated data:
    # data_path = "./mcmc_rl_fit/projects/fmri/data_sim/SARSPEICZVG_all_fMRI_Sim_exporder_SurpAC_run2.csv"

    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    description = "Batch fit: Models / subjects"
    # ---- DEBUG settings:
    max_nr_episodes = 50  # FIT to the first 10 episodes. For the main analysis, I always used 40
    max_nr_episodes = 2  # FIT to the first 10 episodes. For the main analysis, I always used 40
    nr_repetitions = 2
    top_n_fraction = 0.1
    nr_burnin = 3
    nr_samples = 3
    inter_collect_interval = 4
    # # end debug settings
    # ---------------------------
    # ---- Settings used for paper
    # max_nr_episodes = 50  # FIT to the first n episodes.
    # nr_repetitions = 50 #20 # one repetition generates nr_mcmc_samples. doing several repetitions, starting from different random initial positions helps finding isolated regions (e.g. bimodal distributions).
    # top_n_fraction = 0.001 # not important. I sometimes a the top 1percent of mcmc samples
    # nr_burnin = 1500 # important. mcmc needs to "burn in". The more params (larger space) the longer it takes to walk into the volume of interest.
    # nr_samples = 2000 # how many samples to collect from each independent repetition
    # inter_collect_interval = 10 # samples are correlated. therefore make N steps before collecting a sample. 7 to 10 seems a good value.
    # # end
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------

    fit_spec = RL_Fit.FitSpec_multiple_models(
        batch_fit_models_spec,
        max_nr_episodes,
        nr_burnin,
        nr_samples,
        inter_collect_interval,
        top_n_fraction,
        nr_repetitions,
        false,  #do_per_subject_fit
        # true,  #do_per_subject_fit
        result_path_base
    )
    # --------------------------------------------------------------------------
    # # ----------------- CROSSVALIDATE -----------------
    (dataFmriAll, all_data_stats) = RL_Fit.load_SARSPEI(data_path; max_episodes_per_subj = max_nr_episodes)
    # when running just an additional model, reuse the same (initially randomly) sampled
    # (make results comparable with previously fit algorithms)

    # I called kf_Spec = RL_Fit.get_k_folds_all_out(Set(all_data_stats.all_subject_ids), 3) once
    # and then fix the subject ID per fold
    kf_Spec = RL_Fit.get_k_folds( Set(all_data_stats.all_subject_ids),
    [
    Set([7, 2, 19, 5, 8, 12, 20])
    Set([18, 4, 13, 3, 15, 21, 6])
    Set([9, 14, 10, 16, 17, 11, 1])])
    # Otherwise shuffle(collect(1:21))... and split in 3
    CV_multi_model_spec = RL_Fit.CV_multiple_models_spec(
        batch_fit_models_spec,
        kf_Spec,
        nr_burnin,
        nr_samples,
        inter_collect_interval,
        top_n_fraction,
        nr_repetitions,
        result_path_base)
    cv_result = RL_Fit.run_cross_validation(description, dataFmriAll,
                                            CV_multi_model_spec, rl_module_path)
end


export AbstractBlock, PrimitiveBlock, CompositeBlock, AbstractContainer, TagBlock

"""
    AbstractBlock

Abstract type for quantum circuit blocks.
"""
abstract type AbstractBlock{N} end

"""
    PrimitiveBlock{N} <: AbstractBlock{N}

Abstract type that all primitive block will subtype from. A primitive block
is a concrete block who can not be decomposed into other blocks. All composite
block can be decomposed into several primitive blocks.

!!! note

    subtype for primitive block with parameter should implement `hash` and `==`
    method to enable key value cache.
"""
abstract type PrimitiveBlock{N} <: AbstractBlock{N} end

"""
    CompositeBlock{N} <: AbstractBlock{N}

Abstract supertype which composite blocks will inherit from. Composite blocks
are blocks composited from other [`AbstractBlock`](@ref)s, thus it is a `AbstractBlock`
as well.
"""
abstract type CompositeBlock{N} <: AbstractBlock{N} end

"""
    AbstractContainer{BT, N} <: CompositeBlock{N}

Abstract type for container block. Container blocks are blocks contain a single
block. Container block should have a
"""
abstract type AbstractContainer{BT<:AbstractBlock,N} <: CompositeBlock{N} end

"""
    TagBlock{BT, N} <: AbstractContainer{BT, N}

`TagBlock` is a special kind of Container block, it forwards most of the methods
but tag the block with some extra information.
"""
abstract type TagBlock{BT,N} <: AbstractContainer{BT,N} end

"""
    content(x)

Returns the content of `x`.
"""
@interface content

"""
    chcontent(x, blk)

Create a similar block of `x` and change its content to blk.
"""
@interface chcontent

"""
    apply!(register, block)

Apply a block (of quantum circuit) to a quantum register.
"""

@interface apply!

"""
    occupied_locs(x)

Return a tuple of occupied locations of `x`.
"""
@interface occupied_locs

"""
    subblocks(x)

Returns an iterator of the sub-blocks of a composite block. Default is empty.
"""
@interface subblocks

"""
    chsubblocks(composite_block, itr)

Change the sub-blocks of a [`CompositeBlock`](@ref) with given iterator `itr`.
"""
@interface chsubblocks

@interface print_block

"""
    mat([T=ComplexF64], blk)

Returns the matrix form of given block.
"""
@interface mat

# parameters
"""
    getiparams(block)

Returns the intrinsic parameters of node `block`, default is an empty tuple.
"""
@interface getiparams

"""
    setiparams!(block, itr)
    setiparams!(block, params...)

Set the parameters of `block`.
"""
@interface setiparams!

"""
    parameters(block)

Returns all the parameters contained in block tree with given root `block`.
"""
@interface parameters

"""
    nparameters(block) -> Int

Return number of parameters in `block`. See also [`niparams`](@ref).
"""
@interface nparameters

"""
    niparam(block) -> Int

Return number of intrinsic parameters in `block`. See also [`nparameters`](@ref).
"""
@interface niparams

"""
    iparams_eltype(block)

Return the element type of [`getiparams`](@ref).
"""
@interface iparams_eltype

"""
    parameters_eltype(x)

Return the element type of [`parameters`](@ref).
"""
@interface parameters_eltype

"""
    dispatch!(x::AbstractBlock, collection)

Dispatch parameters in collection to block tree `x`.

!!! note

    it will try to dispatch the parameters in collection first.
"""
@interface dispatch!

"""
    render_params(r::AbstractBlock, params)

`params` can be a number or a symbol like `:zero` and `:random`.
This function renders the input parameter to a consumable type to `r`.
"""
@interface render_params

"""
    expect(op::AbstractBlock, reg) -> Vector
    expect(op::AbstractBlock, reg => circuit) -> Vector
    expect(op::AbstractBlock, density_matrix) -> Vector

Get the expectation value of an operator, the second parameter can be a register `reg` or a pair of input register and circuit `reg => circuit`.

expect'(op::AbstractBlock, reg=>circuit) -> Pair
expect'(op::AbstractBlock, reg) -> AbstracRegister

Obtain the gradient with respect to registers and circuit parameters.
For pair input, the second return value is a pair of `gψ=>gparams`,
with `gψ` the gradient of input state and `gparams` the gradients of circuit parameters.
For register input, the return value is a register.

!!! note

    For batched register, `expect(op, reg=>circuit)` returns a vector of size number of batch as output. However, one can not differentiate over a vector loss, so `expect'(op, reg=>circuit)` accumulates the gradient over batch, rather than returning a batched gradient of parameters.
"""
@interface expect

"""
    operator_fidelity(b1::AbstractBlock, b2::AbstractBlock) -> Number

Operator fidelity defined as

```math
F^2 = \\frac{1}{d^2}\\left[{\\rm Tr}(b1^\\dagger b2)\\right]
```

Here, `d` is the size of the Hilbert space. Note this quantity is independant to global phase.
See arXiv: 0803.2940v2, Equation (2) for reference.
"""
@interface operator_fidelity

"""
    apply_back!((ψ, ∂L/∂ψ*), circuit::AbstractBlock, collector) -> AbstractRegister

back propagate and calculate the gradient ∂L/∂θ = 2*Re(∂L/∂ψ*⋅∂ψ*/∂θ), given ∂L/∂ψ*.
`ψ` is the output register, ∂L/∂ψ* should also be register type.

Note: gradients are stored in `Diff` blocks, it can be access by either `diffblock.grad` or `gradient(circuit)`.
Note2: now `apply_back!` returns the inversed gradient!
"""
@interface apply_back!

"""
    mat_back!(T, rb::AbstractBlock, adjy, collector)

Back propagate the matrix gradients.
"""
@interface mat_back!


"""
`module  JAC.PhotoExcitation`
    ... a submodel of JAC that contains all methods for computing photo-excitation properties between some initial
        and final-state multiplets.
"""
module PhotoExcitation

    using Printf, ..AngularMomentum, ..Basics,  ..Basics,  ..Defaults, ..ManyElectron, ..Radial, ..PhotoEmission, ..TableStrings

    """
    `struct  PhotoExcitation.Settings`  ... defines a type for the details and parameters of computing photo-excitation  lines.

        + multipoles              ::Array{EmMultipole,1}    ... Specifies the multipoles of the radiation field that are to be included.
        + gauges                  ::Array{UseGauge,1}       ... Specifies the gauges to be included into the computations.
        + calcForStokes           ::Bool                    ... True, if the excitation cross sections are to be calculated (and false otherwise)
                                                                for given Stokes parameter of the incident plane-wave photons.
        + calcPhotonDm            ::Bool                    ... True, if the photon density matrix of a subsequently emitted fluorescence photon
                                                                is to be calculated and false otherwise.
        + calcTensors             ::Bool                    ... True, if statistical tensors of the excited atom are to be calculated, false otherwise.
        + printBefore             ::Bool                    ... True, if all energies and lines are printed before their evaluation.
        + lineSelection           ::LineSelection           ... Specifies the selected levels, if any.
        + photonEnergyShift       ::Float64                 ... An overall energy shift for all photon energies.
        + mimimumPhotonEnergy     ::Float64                 ... minimum transition energy for which (photon) transitions are included into the
                                                                computation.
        + maximumPhotonEnergy     ::Float64                 ... maximum transition energy for which (photon) transitions are included.
        + stokes                  ::ExpStokes               ... Stokes parameters of the incident radiation.
    """
    struct Settings
        multipoles                ::Array{EmMultipole,1}
        gauges                    ::Array{UseGauge,1}
        calcForStokes             ::Bool
        calcPhotonDm              ::Bool
        calcTensors               ::Bool
        printBefore               ::Bool
        lineSelection             ::LineSelection
        photonEnergyShift         ::Float64
        mimimumPhotonEnergy       ::Float64
        maximumPhotonEnergy       ::Float64
        stokes                    ::ExpStokes
    end

    """
    `PhotoExcitation.Settings()`  ... 'empty' constructor for the default values of photo-excitation line computations
    """
    function Settings()
        Settings(EmMultipole[], UseGauge[], false, false, false, false, LineSelection(), 0., 0., 0., Basics.ExpStokes())
    end

    """
    `PhotoExcitation.Settings(set::PhotoExcitation.Settings;`

            multipoles=..,          gauges=..,                  calcForStokes=..,           calcPhotonDm=..,
            calcTensors=..,         printBefore=..,             lineSelection=..,
            photonEnergyShift=..,   mimimumPhotonEnergy=..,     maximumPhotonEnergy=..,     stokes=..)

        ... constructor for modifying the given PhotoExcitation.Settings by 'overwriting' the previously selected parameters.
    """
    function Settings(set::PhotoExcitation.Settings;
        multipoles::Union{Nothing,Array{EmMultipole,1}}=nothing,        gauges::Union{Nothing,Array{UseGauge,1}}=nothing,
        calcForStokes::Union{Nothing,Bool}=nothing,                     calcPhotonDm::Union{Nothing,Bool}=nothing,
        calcTensors::Union{Nothing,Bool}=nothing,                       printBefore::Union{Nothing,Bool}=nothing,
        lineSelection::Union{Nothing,LineSelection}=nothing,
        photonEnergyShift::Union{Nothing,Float64}=nothing,              mimimumPhotonEnergy::Union{Nothing,Float64}=nothing,
        maximumPhotonEnergy::Union{Nothing,Float64}=nothing,            stokes::Union{Nothing,ExpStokes}=nothing)

        if  multipoles          == nothing   multipolesx          = set.multipoles              else  multipolesx          = multipoles            end
        if  gauges              == nothing   gaugesx              = set.gauges                  else  gaugesx              = gauges                end
        if  calcForStokes       == nothing   calcForStokesx       = set.calcForStokes           else  calcForStokesx       = calcForStokes         end
        if  calcPhotonDm        == nothing   calcPhotonDmx        = set.calcPhotonDm            else  calcPhotonDmx        = calcPhotonDm          end
        if  calcTensors         == nothing   calcTensorsx         = set.calcTensors             else  calcTensorsx         = calcTensors           end
        if  printBefore         == nothing   printBeforex         = set.printBefore             else  printBeforex         = printBefore           end
        if  lineSelection       == nothing   lineSelectionx       = set.lineSelection           else  lineSelectionx       = lineSelection         end
        if  photonEnergyShift   == nothing   photonEnergyShiftx   = set.photonEnergyShift       else  photonEnergyShiftx   = photonEnergyShift     end
        if  mimimumPhotonEnergy == nothing   mimimumPhotonEnergyx = set.mimimumPhotonEnergy     else  mimimumPhotonEnergyx = mimimumPhotonEnergy   end
        if  maximumPhotonEnergy == nothing   maximumPhotonEnergyx = set.maximumPhotonEnergy     else  maximumPhotonEnergyx = maximumPhotonEnergy   end
        if  stokes              == nothing   stokesx              = set.stokes                  else  stokesx              = stokes                end

        Settings( multipolesx, gaugesx, calcForStokesx, calcPhotonDmx, calcTensorsx, printBeforex, lineSelectionx,
                  photonEnergyShiftx, mimimumPhotonEnergyx, maximumPhotonEnergyx, stokesx)
    end

    # `Base.show(io::IO, settings::PhotoExcitation.Settings)`  ... prepares a proper printout of the variable settings::PhotoExcitation.Settings.
    function Base.show(io::IO, settings::PhotoExcitation.Settings)
        println(io, "multipoles:               $(settings.multipoles)  ")
        println(io, "use-gauges:               $(settings.gauges)  ")
        println(io, "calcForStokes:            $(settings.calcForStokes)  ")
        println(io, "calcPhotonDm:             $(settings.calcPhotonDm)  ")
        println(io, "calcTensors:              $(settings.calcTensors)  ")
        println(io, "printBefore:              $(settings.printBefore)  ")
        println(io, "lineSelection:            $(settings.lineSelection)  ")
        println(io, "photonEnergyShift:        $(settings.photonEnergyShift)  ")
        println(io, "mimimumPhotonEnergy:      $(settings.mimimumPhotonEnergy)  ")
        println(io, "maximumPhotonEnergy:      $(settings.maximumPhotonEnergy)  ")
        println(io, "stokes:                   $(settings.stokes)  ")
    end

    """
    `struct  PhotoExcitation.Line`
        ... defines a type for a photo-excitation line that may include the definition of sublines and their
            corresponding amplitudes.

        + initialLevel   ::Level                       ... initial-(state) level
        + finalLevel     ::Level                       ... final-(state) level
        + omega          ::Float64                     ... Transition frequency of this line; can be shifted w.r.t. the level energies.
        + oscStrength    ::EmProperty                  ... Absorption oscillator strength
        + crossSection   ::EmProperty                  ... Total cross section of this line.
        + staTensor      ::Array{TensorComp,1}         ... Array of statistical tensor components rho_kq
        + hasSublines    ::Bool                        ... Determines whether the individual sublines are defined in terms of their
                                                           multipolarity, amplitude, or not; cf. PhotoEmission.Channel
        + channels       ::Array{PhotoEmission.Channel,1}  ... List of radiative (photon) channels
    """
    struct  Line
        initialLevel     ::Level
        finalLevel       ::Level
        omega            ::Float64
        oscStrength      ::EmProperty
        crossSection     ::EmProperty
        staTensor        ::Array{TensorComp,1}
        hasSublines      ::Bool
        channels         ::Array{PhotoEmission.Channel,1}
    end

    """
    `PhotoExcitation.Line(initialLevel::Level, finalLevel::Level, omega::Float64, crossSection::EmProperty)`
        ... constructor an photo-excitation line between a specified initial and final level.
    """
    function Line(initialLevel::Level, finalLevel::Level, omega::Float64, crossSection::EmProperty)
        Line(initialLevel, finalLevel, omega, EmProperty(0., 0.), crossSection, false, PhotoEmission.Channel[])
    end

    # `Base.show(io::IO, line::PhotoExcitation.Line)`  ... prepares a proper printout of the variable line::PhotoExcitation.Line.
    function Base.show(io::IO, line::PhotoExcitation.Line)
        println(io, "initialLevel:         $(line.initialLevel)  ")
        println(io, "finalLevel:           $(line.finalLevel)  ")
        println(io, "omega:                $(line.omega)  ")
        println(io, "oscStrength:          $(line.oscStrength)  ")
        println(io, "crossSection:         $(line.crossSection)  ")
        println(io, "staTensor:            $(line.staTensor)  ")
        println(io, "hasSublines:          $(line.hasSublines)  ")
        println(io, "channels:             $(line.channels)  ")
    end

    """
    `PhotoExcitation.computeAmplitudesProperties(line::PhotoExcitation.Line, grid::Radial.Grid, settings::PhotoExcitation.Settings; printout::Bool=true)`
        ... to compute all amplitudes and properties of the given line; a line::PhotoExcitation.Line is returned for which
            the amplitudes and properties have now been evaluated.
    """
    function  computeAmplitudesProperties(line::PhotoExcitation.Line, grid::Radial.Grid, settings::PhotoExcitation.Settings; printout::Bool=true)
        newChannels = PhotoEmission.Channel[];    Ji2 = Basics.twice(line.initialLevel.J);    Jf2 = Basics.twice(line.finalLevel.J)
        for  channel  in  line.channels
            amplitude = PhotoEmission.amplitude("absorption", channel.multipole, channel.gauge, line.omega,
                                                line.finalLevel, line.initialLevel, grid, printout=printout)
            push!( newChannels, PhotoEmission.Channel(channel.multipole, channel.gauge, amplitude) )
        end
        # Calculate the absorption oscillator strength
        oscCoulomb = oscBabushkin = 0.;   omega = line.omega;   alpha = Defaults.getDefaults("alpha")
        for  channel  in  newChannels
            ##x wa = line.omega / (Basics.twice(channel.multipole.L) + 1) * (alpha * line.omega)^(2*channel.multipole.L - 2)
            wa = 8pi * alpha * line.omega / (Ji2 + 1) * (Jf2 + 1)
            if      channel.gauge == Basics.Coulomb     oscCoulomb   = oscCoulomb    +  channel.amplitude * conj(channel.amplitude) * wa
            elseif  channel.gauge == Basics.Babushkin   oscBabushkin = oscBabushkin  +  channel.amplitude * conj(channel.amplitude) * wa
            elseif  channel.gauge == Basics.Magnetic    oscBabushkin = oscBabushkin  +  channel.amplitude * conj(channel.amplitude) * wa
                                                        oscCoulomb   = oscCoulomb    +  channel.amplitude * conj(channel.amplitude) * wa
            else    error("stop a")
            end
        end
        # Calculate the crossSection and stastistical tensors if requested
        csCoulomb = csBabushkin = 0.;
        for  channel  in  newChannels
            if      channel.gauge == Basics.Coulomb     csCoulomb   = csCoulomb    +  channel.amplitude * conj(channel.amplitude)
            elseif  channel.gauge == Basics.Babushkin   csBabushkin = csBabushkin  +  channel.amplitude * conj(channel.amplitude)
            elseif  channel.gauge == Basics.Magnetic    csBabushkin = csBabushkin  +  channel.amplitude * conj(channel.amplitude)
                                                        csCoulomb   = csCoulomb    +  channel.amplitude * conj(channel.amplitude)
            else    error("stop b")
            end
        end
        oscFactor    = (Ji2 + 1) / (Jf2 + 1) / 2. * Defaults.getDefaults("speed of light: c")^3 / line.omega^2
        csFactor     = 2 * pi^3 * Defaults.getDefaults("alpha") / (line.omega * (Ji2 + 1))
        oscStrength  = EmProperty(oscFactor*oscCoulomb, oscFactor*oscBabushkin)
        crossSection = EmProperty(csFactor*csCoulomb, csFactor*csBabushkin)
        staTensor    = TensorComp[]
        if  settings.calcTensors    push!(staTensor, TensorComp(0, 0, 1.))      end

        newLine = PhotoExcitation.Line( line.initialLevel, line.finalLevel, line.omega, oscStrength, crossSection, staTensor, true, newChannels)
        return( newLine )
    end

    """
    `PhotoExcitation.computeCrossSection(line::PhotoExcitation.Line, stokes::ExpStokes)`
        ... to compute the excitation cross section for the excitation of unpolarized atoms by plane-wave photons, whose polarization
            is described by the given (experimental) Stokes parameters. A cs::EmProperty is returned.
    """
    function  computeCrossSection(line::PhotoExcitation.Line, stokes::ExpStokes)
        csCoulomb = -2.;    csBabushkin = -2.
        cs = EmProperty( csCoulomb, csBabushkin )
        return( cs )
    end

    """
    `PhotoExcitation.computeLines(finalMultiplet::Multiplet, initialMultiplet::Multiplet, grid::Radial.Grid,
                                      settings::PhotoExcitation.Settings; output=true)`
        ... to compute the photo-excitation amplitudes and all properties as requested by the given settings. A list
            of lines::Array{PhotoExcitation.Lines} is returned.
    """
    function  computeLines(finalMultiplet::Multiplet, initialMultiplet::Multiplet, grid::Radial.Grid,
                           settings::PhotoExcitation.Settings; output=true)
        println("")
        printstyled("PhotoExcitation.computeLines(): The computation of the excitation cross sections, etc. starts now ... \n", color=:light_green)
        printstyled("----------------------------------------------------------------------------------------------------- \n", color=:light_green)
        println("")
        lines = PhotoExcitation.determineLines(finalMultiplet, initialMultiplet, settings)
        # Display all selected lines before the computations start
        if  settings.printBefore    PhotoExcitation.displayLines(lines)    end
        # Calculate all amplitudes and requested properties
        newLines = PhotoExcitation.Line[]
        for  line in lines
            newLine = PhotoExcitation.computeAmplitudesProperties(line, grid, settings)
            push!( newLines, newLine)
        end
        # Print all results to screen
        PhotoExcitation.displayCrossSections(stdout, newLines, settings)
        printSummary, iostream = Defaults.getDefaults("summary flag/stream")
        if  printSummary    PhotoExcitation.displayCrossSections(iostream, newLines, settings)   end
        #
        if    output    return( lines )
        else            return( nothing )
        end
    end

    """
    `PhotoExcitation.computeLinesCascade(finalMultiplet::Multiplet, initialMultiplet::Multiplet, grid::Radial.Grid,
                                         settings::PhotoExcitation.Settings; output::Bool=true, printout::Bool=true)`
        ... to compute the excitation (absorption) transition amplitudes and all properties as requested by the given settings.
            The computations and printout is adapted for larger cascade computations by including only lines with at least one channel
            and by sending all printout to a summary file only. A list of lines::Array{PhotoExcitation.Lines} is returned.
    """
    function  computeLinesCascade(finalMultiplet::Multiplet, initialMultiplet::Multiplet, grid::Radial.Grid,
                                         settings::PhotoExcitation.Settings; output::Bool=true, printout::Bool=true)
        # Define a common subshell list for both multiplets
        subshellList = Basics.generate("subshells: ordered list for two bases", finalMultiplet.levels[1].basis, initialMultiplet.levels[1].basis)
        Defaults.setDefaults("relativistic subshell list", subshellList; printout=false)
        lines = PhotoExcitation.determineLines(finalMultiplet, initialMultiplet, settings)
        # Display all selected lines before the computations start
        # if  settings.printBefore    PhotoExcitation.displayLines(lines)    end
        # Calculate all amplitudes and requested properties
        newLines = PhotoExcitation.Line[]
        for  (i,line)  in  enumerate(lines)
            if  rem(i,50) == 0    println("> Excitation line $i:")   end
            newLine = PhotoExcitation.computeAmplitudesProperties(line, grid, settings, printout=printout)
            #
            # Don't add this line if it does not contribute to the decay
            wa = 0.
            for  ch in newLine.channels   wa = wa + abs(ch.amplitude)^2    end
            if   wa == 0.    continue    end
            push!( newLines, newLine)
        end
        # Print all results to a summary file, if requested
        printSummary, iostream = Defaults.getDefaults("summary flag/stream")
        if  printSummary   PhotoExcitation.displayCrossSections(iostream, newLines, settings)    end
        #
        if    output    return( newLines )
        else            return( nothing )
        end
    end

    """
    `PhotoExcitation.computeStatisticalTensor(k::Int64, q::Int64, line::PhotoExcitation.Line, stokes::ExpStokes)`
        ... to compute the statistical tensor (component) rho_{k,q} of the final level for the excitation of unpolarized atoms by
            plane-wave photons, whose polarization is described by the given (experimental) Stokes parameters.
            A rho_kq::EmProperty is returned.
    """
    function  computeStatisticalTensor(k::Int64, q::Int64, line::PhotoExcitation.Line, stokes::ExpStokes)
        if  k == q  ||  k == -q   rhoCoulomb = -3.;    rhoBabushkin = -3.    else   rhoCoulomb = 0.;    rhoBabushkin = 0.    end
        rho = EmProperty( rhoCoulomb, rhoBabushkin )
        return( rho )
    end

    """
    `PhotoExcitation.determineChannels(finalLevel::Level, initialLevel::Level, settings::PhotoExcitation.Settings)`
        ... to determine a list of PhotoExcitation.Channel for a transitions from the initial to final level and by taking
            into account the particular settings of for this computation; an Array{PhotoExcitation.Channel,1} is returned.
    """
    function determineChannels(finalLevel::Level, initialLevel::Level, settings::PhotoExcitation.Settings)
        channels = PhotoEmission.Channel[];
        symi = LevelSymmetry(initialLevel.J, initialLevel.parity);    symf = LevelSymmetry(finalLevel.J, finalLevel.parity)
        for  mp in settings.multipoles
            if   AngularMomentum.isAllowedMultipole(symi, mp, symf)
                hasMagnetic = false
                for  gauge in settings.gauges
                    # Include further restrictions if appropriate
                    if     string(mp)[1] == 'E'  &&   gauge == Basics.UseCoulomb      push!(channels, PhotoEmission.Channel(mp, Basics.Coulomb,   0.) )
                    elseif string(mp)[1] == 'E'  &&   gauge == Basics.UseBabushkin    push!(channels, PhotoEmission.Channel(mp, Basics.Babushkin, 0.) )
                    elseif string(mp)[1] == 'M'  &&   !(hasMagnetic)               push!(channels, PhotoEmission.Channel(mp, Basics.Magnetic,  0.) );
                                                        hasMagnetic = true;
                    end
                end
            end
        end
        ##x println("PhotoExcitation.determineChannels-aa: channels = $channels ")
        return( channels )
    end

    """
    `PhotoExcitation.determineLines(finalMultiplet::Multiplet, initialMultiplet::Multiplet, settings::PhotoExcitation.Settings)`
        ... to determine a list of photo-excitation Line's for transitions between the levels from the given initial- and
            final-state multiplets and by taking into account the particular selections and settings for this computation;
            an Array{PhotoExcitation.Line,1} is returned. Apart from the level specification, all physical properties are set to
            zero during the initialization process.
    """
    function  determineLines(finalMultiplet::Multiplet, initialMultiplet::Multiplet, settings::PhotoExcitation.Settings)
        lines = PhotoExcitation.Line[]
        for  iLevel  in  initialMultiplet.levels
            for  fLevel  in  finalMultiplet.levels
                if  Basics.selectLevelPair(iLevel, fLevel, settings.lineSelection)
                    omega    = fLevel.energy - iLevel.energy + settings.photonEnergyShift
                    if   omega == 0.  ||  omega < settings.mimimumPhotonEnergy  ||  omega > settings.maximumPhotonEnergy    continue   end
                    channels = PhotoExcitation.determineChannels(fLevel, iLevel, settings)
                    if   length(channels) == 0   continue   end
                    push!( lines, PhotoExcitation.Line(iLevel, fLevel, omega, EmProperty(0., 0.), EmProperty(0., 0.), TensorComp[], true, channels) )
                end
            end
        end
        return( lines )
    end

    """
    `PhotoExcitation.displayLines(lines::Array{PhotoExcitation.Line,1})`
        ... to display a list of lines and channels that have been selected due to the prior settings. A neat table of all
            selected transitions and energies is printed but nothing is returned otherwise.
    """
    function  displayLines(lines::Array{PhotoExcitation.Line,1})
        nx = 120
        println(" ")
        println("  Selected photo-excitation lines:")
        println(" ")
        println("  ", TableStrings.hLine(nx))
        sa = "  ";   sb = "  "
        sa = sa * TableStrings.center(18, "i-level-f"; na=0);                         sb = sb * TableStrings.hBlank(18)
        sa = sa * TableStrings.center(18, "i--J^P--f"; na=4);                         sb = sb * TableStrings.hBlank(22)
        sa = sa * TableStrings.center(14, "Energy"; na=4);
        sb = sb * TableStrings.center(14, TableStrings.inUnits("energy"); na=4)
        sa = sa * TableStrings.flushleft(30, "List of multipoles"; na=4);             sb = sb * TableStrings.hBlank(34)
        println(sa);    println(sb);    println("  ", TableStrings.hLine(nx))
        #
        for  line in lines
            sa  = "";    isym = LevelSymmetry( line.initialLevel.J, line.initialLevel.parity)
                         fsym = LevelSymmetry( line.finalLevel.J,   line.finalLevel.parity)
            sa = sa * TableStrings.center(18, TableStrings.levels_if(line.initialLevel.index, line.finalLevel.index); na=2)
            sa = sa * TableStrings.center(18, TableStrings.symmetries_if(isym, fsym); na=4)
            sa = sa * @sprintf("%.8e", Defaults.convertUnits("energy: from atomic", line.omega)) * "    "
            mpGaugeList = Tuple{Basics.EmMultipole,Basics.EmGauge}[]
            for  i in 1:length(line.channels)
                push!( mpGaugeList, (line.channels[i].multipole, line.channels[i].gauge) )
            end
            ##x println("PhotoExcitation.diplayLines-ad: mpGaugeList = ", mpGaugeList)
            sa = sa * TableStrings.multipoleGaugeTupels(50, mpGaugeList)
            println( sa )
        end
        println("  ", TableStrings.hLine(nx))
        #
        return( nothing )
    end

    """
    `PhotoExcitation.displayCrossSections(stream::IO, lines::Array{PhotoExcitation.Line,1}, settings::PhotoExcitation.Settings)`
        ... to list all results, energies, cross sections, etc. of the selected lines. A neat table is printed but nothing
            is returned otherwise.
    """
    function  displayCrossSections(stream::IO, lines::Array{PhotoExcitation.Line,1}, settings::PhotoExcitation.Settings)
        #
        # Photoexcitation resonance strength as derived from oscillator strength
        nx = 122
        println(stream, " ")
        println(stream, "  Photoexcitation resonance strength as derived from oscillator strength:")
        println(stream, " ")
        println(stream, "  ", TableStrings.hLine(nx))
        sa = "  ";   sb = "  "
        sa = sa * TableStrings.center(18, "i-level-f"; na=0);                         sb = sb * TableStrings.hBlank(18)
        sa = sa * TableStrings.center(18, "i--J^P--f"; na=4);                         sb = sb * TableStrings.hBlank(22)
        sa = sa * TableStrings.center(14, "Energy"   ; na=3);
        sb = sb * TableStrings.center(14,TableStrings.inUnits("energy"); na=3)
        sa = sa * TableStrings.center(10, "Multipoles"; na=3);                        sb = sb * TableStrings.hBlank(13)
        sa = sa * TableStrings.center(24, "Cou -- f_ik -- Bab"; na=1);
        sb = sb * TableStrings.center(24, "     absorption   "; na=2)
        sa = sa * TableStrings.center(24, "Cou -- S -- Bab"; na=3);
        sb = sb * TableStrings.center(24, "   [Mb eV]     "; na=3)
        println(stream, sa);    println(stream, sb);    println(stream, "  ", TableStrings.hLine(nx))
        #
        for  line in lines
            sa  = "";    isym = LevelSymmetry( line.initialLevel.J, line.initialLevel.parity)
                         fsym = LevelSymmetry( line.finalLevel.J,   line.finalLevel.parity)
            sa = sa * TableStrings.center(18, TableStrings.levels_if(line.initialLevel.index, line.finalLevel.index); na=2)
            sa = sa * TableStrings.center(18, TableStrings.symmetries_if(isym, fsym); na=4)
            sa = sa * @sprintf("%.8e", Defaults.convertUnits("energy: from atomic", line.omega)) * "    "
            multipoles = EmMultipole[]
            for  ch in line.channels
                multipoles = push!( multipoles, ch.multipole)
            end
            multipoles = unique(multipoles);   mpString = TableStrings.multipoleList(multipoles) * "          "
            sa = sa * TableStrings.flushleft(11, mpString[1:10];  na=3)
            sa = sa * @sprintf("%.4e", line.oscStrength.Coulomb)   * "   "
            sa = sa * @sprintf("%.4e", line.oscStrength.Babushkin) * "    "
            sa = sa * @sprintf("%.4e", line.oscStrength.Coulomb   * 109.761)   * "   "
            sa = sa * @sprintf("%.4e", line.oscStrength.Babushkin * 109.761)   * "    "
            println(stream, sa)
        end
        println(stream, "  ", TableStrings.hLine(nx))
        #
        #
        # Photoexcitation cross sections from line strength for (completely) linearly-polarized plane-wave photons
        nx = 103
        println(stream, " ")
        println(stream, "  Photoexcitation cross sections for (completely) linearly-polarized plane-wave photons:  ... not yet correct !!")
        println(stream, " ")
        println(stream, "  ", TableStrings.hLine(nx))
        sa = "  ";   sb = "  "
        sa = sa * TableStrings.center(18, "i-level-f"; na=0);                         sb = sb * TableStrings.hBlank(18)
        sa = sa * TableStrings.center(18, "i--J^P--f"; na=4);                         sb = sb * TableStrings.hBlank(22)
        sa = sa * TableStrings.center(14, "Energy"   ; na=3);
        sb = sb * TableStrings.center(14,TableStrings.inUnits("energy"); na=3)
        sa = sa * TableStrings.center(10, "Multipoles"; na=2);                        sb = sb * TableStrings.hBlank(14)
        sa = sa * TableStrings.center(28, "Cou -- cross section -- Bab"; na=0);
        sb = sb * TableStrings.center(28, TableStrings.inUnits("cross section")*"       "*
                                          TableStrings.inUnits("cross section"); na=1)
        println(stream, sa);    println(stream, sb);    println(stream, "  ", TableStrings.hLine(nx))
        #
        for  line in lines
            sa  = "";    isym = LevelSymmetry( line.initialLevel.J, line.initialLevel.parity)
                         fsym = LevelSymmetry( line.finalLevel.J,   line.finalLevel.parity)
            sa = sa * TableStrings.center(18, TableStrings.levels_if(line.initialLevel.index, line.finalLevel.index); na=2)
            sa = sa * TableStrings.center(18, TableStrings.symmetries_if(isym, fsym); na=4)
            sa = sa * @sprintf("%.8e", Defaults.convertUnits("energy: from atomic", line.omega)) * "    "
            multipoles = EmMultipole[]
            for  ch in line.channels
                multipoles = push!( multipoles, ch.multipole)
            end
            multipoles = unique(multipoles);   mpString = TableStrings.multipoleList(multipoles) * "          "
            sa = sa * TableStrings.flushleft(11, mpString[1:10];  na=5)
            sa = sa * @sprintf("%.4e", Defaults.convertUnits("cross section: from atomic", line.crossSection.Coulomb))     * "   "
            sa = sa * @sprintf("%.4e", Defaults.convertUnits("cross section: from atomic", line.crossSection.Babushkin))   * "    "
            ##x sa = sa * @sprintf("%.6e", line.anisotropy.Coulomb)      * "    "
            ##x sa = sa * @sprintf("%.6e", line.anisotropy.Babushkin)    * "    "
            println(stream, sa)
        end
        println(stream, "  ", TableStrings.hLine(nx))
        #
        #
        if  settings.calcForStokes
            stokes = settings.stokes
            println(stream, " ")
            println(stream, "  Photoexcitation cross sections for incident plane-wave photons with given $stokes:  ... not yet correct !!")
            println(stream, " ")
            println(stream, "  ", TableStrings.hLine(nx))
            sa = "  ";   sb = "  "
            sa = sa * TableStrings.center(18, "i-level-f"; na=0);                         sb = sb * TableStrings.hBlank(18)
            sa = sa * TableStrings.center(18, "i--J^P--f"; na=4);                         sb = sb * TableStrings.hBlank(22)
            sa = sa * TableStrings.center(14, "Energy"   ; na=3);
            sb = sb * TableStrings.center(14,TableStrings.inUnits("energy"); na=3)
            sa = sa * TableStrings.center(10, "Multipoles"; na=3);                        sb = sb * TableStrings.hBlank(14)
            sa = sa * TableStrings.center(30, "Cou -- Cross section -- Bab"; na=2);
            sb = sb * TableStrings.center(30, TableStrings.inUnits("cross section")*"          "*
                                                  TableStrings.inUnits("cross section"); na=2)
            println(stream, sa);    println(stream, sb);    println(stream, "  ", TableStrings.hLine(nx))
            for  line in lines
                sa  = "";    isym = LevelSymmetry( line.initialLevel.J, line.initialLevel.parity)
                             fsym = LevelSymmetry( line.finalLevel.J,   line.finalLevel.parity)
                sa = sa * TableStrings.center(18, TableStrings.levels_if(line.initialLevel.index, line.finalLevel.index); na=2)
                sa = sa * TableStrings.center(18, TableStrings.symmetries_if(isym, fsym); na=4)
                sa = sa * @sprintf("%.8e", Defaults.convertUnits("energy: from atomic", line.omega)) * "    "
                multipoles = EmMultipole[]
                for  ch in line.channels
                    multipoles = push!( multipoles, ch.multipole)
                end
                multipoles = unique(multipoles);   mpString = TableStrings.multipoleList(multipoles) * "          "
                sa = sa * TableStrings.flushleft(11, mpString[1:10];  na=3)
                crossSection = PhotoExcitation.computeCrossSection(line, stokes)
                sa = sa * @sprintf("%.6e", Defaults.convertUnits("cross section: from atomic", crossSection.Coulomb))     * "    "
                sa = sa * @sprintf("%.6e", Defaults.convertUnits("cross section: from atomic", crossSection.Babushkin))   * "    "
                println(stream, sa)
            end
            println(stream, "  ", TableStrings.hLine(nx))
        end
        #
        #
        if  settings.calcTensors
            nx = 145
            stokes = settings.stokes
            println(stream, " ")
            println(stream, "  Statistical tensors rho_kq  and alignment parameters A_kq for the excitation by incident plane-wave photons")
            println(stream, "  with given $stokes:  ... not yet correct !!")
            println(stream, " ")
            println(stream, "  ", TableStrings.hLine(nx))
            sa = "  ";   sb = "  "
            sa = sa * TableStrings.center(18, "i-level-f"; na=0);                         sb = sb * TableStrings.hBlank(18)
            sa = sa * TableStrings.center(18, "i--J^P--f"; na=4);                         sb = sb * TableStrings.hBlank(22)
            sa = sa * TableStrings.center(14, "Energy"   ; na=3);
            sb = sb * TableStrings.center(14,TableStrings.inUnits("energy"); na=3)
            sa = sa * TableStrings.center(10, "Multipoles"; na=1);                        sb = sb * TableStrings.hBlank(12)
            sa = sa * TableStrings.center(12, "     k    q    "; na=0);                   sb = sb * TableStrings.hBlank(14)
            sa = sa * TableStrings.center(26, "Cou -- rho_kq -- Bab"; na=5);              sb = sb * TableStrings.hBlank(14)
            sa = sa * TableStrings.center(26, "Cou --  A_kq  -- Bab"; na=2);              sb = sb * TableStrings.hBlank(14)
            println(stream, sa);    println(stream, sb);    println(stream, "  ", TableStrings.hLine(nx))
            for  line in lines
                sa  = "";    isym = LevelSymmetry( line.initialLevel.J, line.initialLevel.parity)
                             fsym = LevelSymmetry( line.finalLevel.J,   line.finalLevel.parity)
                sa = sa * TableStrings.center(18, TableStrings.levels_if(line.initialLevel.index, line.finalLevel.index); na=2)
                sa = sa * TableStrings.center(18, TableStrings.symmetries_if(isym, fsym); na=4)
                sa = sa * @sprintf("%.8e", Defaults.convertUnits("energy: from atomic", line.omega)) * "    "
                multipoles = EmMultipole[]
                for  ch in line.channels
                    multipoles = push!( multipoles, ch.multipole)
                end
                multipoles = unique(multipoles);   mpString = TableStrings.multipoleList(multipoles) * "          "
                sa = sa * TableStrings.flushleft(11, mpString[1:10];  na=3)
                println(stream, sa);   sc = TableStrings.hBlank( length(sa)-1 )
                tc00   = PhotoExcitation.computeStatisticalTensor(0, 0, line, stokes)
                for  k = 2:2
                    for  q = -k:k
                        tc = PhotoExcitation.computeStatisticalTensor(k, q, line, stokes)
                        if  abs(tc.Coulomb) == abs(tc.Babushkin) == 0.    continue    end
                        #
                        aCoulomb = tc.Coulomb / tc00.Coulomb;    aBabushkin = tc.Babushkin / tc00.Babushkin
                        sd = sc * TableStrings.level(k) * TableStrings.level(q) * "   "
                        sd = sd * @sprintf("%.6e", tc.Coulomb)     * "  "
                        sd = sd * @sprintf("%.6e", tc.Babushkin)   * "    "
                        sd = sd * @sprintf("%.6e", aCoulomb)       * "  "
                        sd = sd * @sprintf("%.6e", aBabushkin)     * "    "
                        println(stream, sd)
                    end
                end
            end
            println(stream, "  ", TableStrings.hLine(nx))
        end
        #
        #
        if  settings.calcTensors   &&   settings.calcPhotonDm
            stokes = settings.stokes
            println(stream, " ")
            println(stream, "  Reduced density matrix of the fluorescence photons at selected solid angles for the prior excitation by incident plane-wave photons")
            println(stream, "  with given $stokes:")
            println(stream, " ")
            println(stream, "  Not yet implemented !!")
        end
        #
        return( nothing )
    end

end # module