# integration with AbstractFFTs.jl

@reexport using AbstractFFTs

import AbstractFFTs: plan_fft, plan_fft!, plan_bfft, plan_bfft!, plan_ifft,
    plan_rfft, plan_brfft, plan_inv, normalization, fft, bfft, ifft, rfft,
    Plan, ScaledPlan

using LinearAlgebra

Base.:(*)(p::Plan{T}, x::DenseCuArray) where {T} = p * copy1(T, x)
Base.:(*)(p::ScaledPlan, x::DenseCuArray) = rmul!(p.p * x, p.scale)

## plan structure

# K is an integer flag for forward/backward
# also used as an alias for r2c/c2r

# inplace is a boolean flag

abstract type CuFFTPlan{T<:cufftNumber, K, inplace} <: Plan{T} end

# for some reason, cufftHandle is an integer and not a pointer...
Base.convert(::Type{cufftHandle}, p::CuFFTPlan) = p.handle

function CUDA.unsafe_free!(plan::CuFFTPlan, stream::CuStream=stream())
    cufftDestroy(plan)
    unsafe_free!(plan.workarea, stream)
end

function unsafe_finalize!(plan::CuFFTPlan)
    context!(plan.ctx; skip_destroyed=true) do
        unsafe_free!(plan, default_stream())
    end
end

mutable struct cCuFFTPlan{T<:cufftNumber,K,inplace,N} <: CuFFTPlan{T,K,inplace}
    handle::cufftHandle
    ctx::CuContext
    stream::CuStream
    workarea::CuVector{Int8}
    sz::NTuple{N,Int} # Julia size of input array
    osz::NTuple{N,Int} # Julia size of output array
    xtype::cufftType
    region::Any
    pinv::ScaledPlan # required by AbstractFFT API

    function cCuFFTPlan{T,K,inplace,N}(handle::cufftHandle, workarea::CuVector{Int8},
                                       X::DenseCuArray{T,N}, sizey::Tuple, region, xtype;
                                       stream::CuStream=stream()) where {T<:cufftNumber,K,inplace,N}
        # maybe enforce consistency of sizey
        p = new(handle, context(), stream, workarea, size(X), sizey, xtype, region)
        finalizer(unsafe_finalize!, p)
        p
    end
end

mutable struct rCuFFTPlan{T<:cufftNumber,K,inplace,N} <: CuFFTPlan{T,K,inplace}
    handle::cufftHandle
    ctx::CuContext
    stream::CuStream
    workarea::CuVector{Int8}
    sz::NTuple{N,Int} # Julia size of input array
    osz::NTuple{N,Int} # Julia size of output array
    xtype::cufftType
    region::Any
    pinv::ScaledPlan # required by AbstractFFT API

    function rCuFFTPlan{T,K,inplace,N}(handle::cufftHandle, workarea::CuVector{Int8},
                                       X::DenseCuArray{T,N}, sizey::Tuple, region, xtype;
                                       stream::CuStream=stream()) where {T<:cufftNumber,K,inplace,N}
        # maybe enforce consistency of sizey
        p = new(handle, context(), stream, workarea, size(X), sizey, xtype, region)
        finalizer(unsafe_finalize!, p)
        p
    end
end

const xtypenames = Dict{cufftType,String}(CUFFT_R2C => "real-to-complex",
                                          CUFFT_C2R => "complex-to-real",
                                          CUFFT_C2C => "complex",
                                          CUFFT_D2Z => "d.p. real-to-complex",
                                          CUFFT_Z2D => "d.p. complex-to-real",
                                          CUFFT_Z2Z => "d.p. complex")

function showfftdims(io, sz, T)
    if isempty(sz)
        print(io,"0-dimensional")
    elseif length(sz) == 1
        print(io, sz[1], "-element")
    else
        print(io, join(sz, "×"))
    end
    print(io, " CuArray of ", T)
end

function Base.show(io::IO, p::CuFFTPlan{T,K,inplace}) where {T,K,inplace}
    print(io, inplace ? "CUFFT in-place " : "CUFFT ",
          xtypenames[p.xtype],
          K == CUFFT_FORWARD ? " forward" : " backward",
          " plan for ")
    showfftdims(io, p.sz, T)
end

Base.size(p::CuFFTPlan) = p.sz

# FFT plans can be user-created on a different task, whose stream might be different from
# the one used in the current task. call this function before every API call that performs
# operations on a stream to ensure the plan is using the correct task-local stream.
@inline function update_stream(plan::CuFFTPlan)
    new_stream = stream()
    if plan.stream != new_stream
        plan.stream = new_stream
        cufftSetStream(plan, new_stream)

        # replace the workarea by one (asynchronously) allocated on the current stream
        new_workarea = similar(plan.workarea)
        cufftSetWorkArea(plan, new_workarea)
        CUDA.unsafe_free!(plan.workarea)
        plan.workarea = new_workarea
    end
    return
end

## plan methods

# Note: we don't implement padded storage dimensions
function create_plan(xtype, xdims, region)
    nrank = length(region)
    sz = [xdims[i] for i in region]
    csz = copy(sz)
    csz[1] = div(sz[1],2) + 1
    batch = prod(xdims) ÷ prod(sz)

    # initialize the plan handle
    handle_ref = Ref{cufftHandle}()
    cufftCreate(handle_ref)
    handle = handle_ref[]

    # take control over the workarea
    cufftSetAutoAllocation(handle, 0)
    cufftSetStream(handle, stream())

    # make the plan
    worksize_ref = Ref{Csize_t}()
    if (nrank == 1) && (batch == 1)
        cufftMakePlan1d(handle, sz[1], xtype, 1, worksize_ref)
    elseif (nrank == 2) && (batch == 1)
        cufftMakePlan2d(handle, sz[2], sz[1], xtype, worksize_ref)
    elseif (nrank == 3) && (batch == 1)
        cufftMakePlan3d(handle, sz[3], sz[2], sz[1], xtype, worksize_ref)
    else
        rsz = (length(sz) > 1) ? rsz = reverse(sz) : sz
        if ((region...,) == ((1:nrank)...,))
            # handle simple case ... simply! (for robustness)
           cufftMakePlanMany(handle, nrank, Cint[rsz...], C_NULL, 1, 1, C_NULL, 1, 1,
                             xtype, batch, worksize_ref)
        else
            if nrank==1 || all(diff(collect(region)) .== 1)
                # _stride: successive elements in innermost dimension
                # _dist: distance between first elements of batches
                if region[1] == 1
                    istride = 1
                    idist = prod(sz)
                    cdist = prod(csz)
                else
                    if region[end] != length(xdims)
                        throw(ArgumentError("batching dims must be sequential"))
                    end
                    istride = prod(xdims[1:region[1]-1])
                    idist = 1
                    cdist = 1
                end
                inembed = Cint[rsz...]
                cnembed = (length(csz) > 1) ? Cint[reverse(csz)...] : Cint[csz[1]]
                ostride = istride
                if xtype == CUFFT_R2C || xtype == CUFFT_D2Z
                    odist = cdist
                    onembed = cnembed
                else
                    odist = idist
                    onembed = inembed
                end
                if xtype == CUFFT_C2R || xtype == CUFFT_Z2D
                    idist = cdist
                    inembed = cnembed
                end
            else
                if any(diff(collect(region)) .< 1)
                    throw(ArgumentError("region must be an increasing sequence"))
                end
                cdims = collect(xdims)
                cdims[region[1]] = div(cdims[region[1]],2)+1

                if region[1] == 1
                    istride = 1
                    ii=1
                    while (ii < nrank) && (region[ii] == region[ii+1]-1)
                        ii += 1
                    end
                    idist = prod(xdims[1:ii])
                    cdist = prod(cdims[1:ii])
                    ngaps = 0
                else
                    istride = prod(xdims[1:region[1]-1])
                    idist = 1
                    cdist = 1
                    ngaps = 1
                end
                nem = ones(Int,nrank)
                cem = ones(Int,nrank)
                id = 1
                for ii=1:nrank-1
                    if region[ii+1] > region[ii]+1
                        ngaps += 1
                    end
                    while id < region[ii+1]
                        nem[ii] *= xdims[id]
                        cem[ii] *= cdims[id]
                        id += 1
                    end
                    @assert nem[ii] >= sz[ii]
                end
                if region[end] < length(xdims)
                    ngaps += 1
                end
                # CUFFT represents batches by a single stride (_dist)
                # so we must verify that region is consistent with this:
                if ngaps > 1
                    throw(ArgumentError("batch regions must be sequential"))
                end

                inembed = Cint[reverse(nem)...]
                cnembed = Cint[reverse(cem)...]
                ostride = istride
                if xtype == CUFFT_R2C || xtype == CUFFT_D2Z
                    odist = cdist
                    onembed = cnembed
                else
                    odist = idist
                    onembed = inembed
                end
                if xtype == CUFFT_C2R || xtype == CUFFT_Z2D
                    idist = cdist
                    inembed = cnembed
                end
            end
            cufftMakePlanMany(handle, nrank, Cint[rsz...],
                              inembed, istride, idist, onembed, ostride, odist,
                              xtype, batch, worksize_ref)
        end
    end

    # assign the workarea
    workarea = CuArray{Int8}(undef, worksize_ref[])
    cufftSetWorkArea(handle, workarea)

    handle, workarea
end

# promote to a complex floating-point type (out-of-place only),
# so implementations only need Complex{Float} methods
for f in (:fft, :bfft, :ifft)
    pf = Symbol("plan_", f)
    @eval begin
        $f(x::DenseCuArray{<:Real}, region=1:ndims(x)) = $f(complexfloat(x), region)
        $pf(x::DenseCuArray{<:Real}, region) = $pf(complexfloat(x), region)
        $f(x::DenseCuArray{<:Complex{<:Union{Integer,Rational}}}, region=1:ndims(x)) = $f(complexfloat(x), region)
        $pf(x::DenseCuArray{<:Complex{<:Union{Integer,Rational}}}, region) = $pf(complexfloat(x), region)
    end
end
rfft(x::DenseCuArray{<:Union{Integer,Rational}}, region=1:ndims(x)) = rfft(realfloat(x), region)
plan_rfft(x::DenseCuArray{<:Real}, region) = plan_rfft(realfloat(x), region)

# region is an iterable subset of dimensions
# spec. an integer, range, tuple, or array

# inplace complex
function plan_fft!(X::DenseCuArray{T,N}, region) where {T<:cufftComplexes,N}
    K = CUFFT_FORWARD
    inplace = true
    xtype = (T == cufftComplex) ? CUFFT_C2C : CUFFT_Z2Z

    handle, workarea = create_plan(xtype, size(X), region)

    cCuFFTPlan{T,K,inplace,N}(handle, workarea, X, size(X), region, xtype)
end

function plan_bfft!(X::DenseCuArray{T,N}, region) where {T<:cufftComplexes,N}
    K = CUFFT_INVERSE
    inplace = true
    xtype =  (T == cufftComplex) ? CUFFT_C2C : CUFFT_Z2Z

    handle, workarea = create_plan(xtype, size(X), region)

    cCuFFTPlan{T,K,inplace,N}(handle, workarea, X, size(X), region, xtype)
end

# out-of-place complex
function plan_fft(X::DenseCuArray{T,N}, region) where {T<:cufftComplexes,N}
    K = CUFFT_FORWARD
    xtype =  (T == cufftComplex) ? CUFFT_C2C : CUFFT_Z2Z
    inplace = false

    handle, workarea = create_plan(xtype, size(X), region)

    cCuFFTPlan{T,K,inplace,N}(handle, workarea, X, size(X), region, xtype)
end

function plan_bfft(X::DenseCuArray{T,N}, region) where {T<:cufftComplexes,N}
    K = CUFFT_INVERSE
    inplace = false
    xtype =  (T == cufftComplex) ? CUFFT_C2C : CUFFT_Z2Z

    handle, workarea = create_plan(xtype, size(X), region)

    cCuFFTPlan{T,K,inplace,N}(handle, workarea, X, size(X), region, xtype)
end

# out-of-place real-to-complex
function plan_rfft(X::DenseCuArray{T,N}, region) where {T<:cufftReals,N}
    K = CUFFT_FORWARD
    inplace = false
    xtype =  (T == cufftReal) ? CUFFT_R2C : CUFFT_D2Z

    handle, workarea = create_plan(xtype, size(X), region)

    ydims = collect(size(X))
    ydims[region[1]] = div(ydims[region[1]],2)+1

    rCuFFTPlan{T,K,inplace,N}(handle, workarea, X, (ydims...,), region, xtype)
end

function plan_brfft(X::DenseCuArray{T,N}, d::Integer, region::Any) where {T<:cufftComplexes,N}
    K = CUFFT_INVERSE
    inplace = false
    xtype =  (T == cufftComplex) ? CUFFT_C2R : CUFFT_Z2D
    ydims = collect(size(X))
    ydims[region[1]] = d

    handle, workarea = create_plan(xtype, (ydims...,), region)

    rCuFFTPlan{T,K,inplace,N}(handle, workarea, X, (ydims...,), region, xtype)
end

# FIXME: plan_inv methods allocate needlessly (to provide type parameters)
# Perhaps use FakeArray types to avoid this.

function plan_inv(p::cCuFFTPlan{T,CUFFT_FORWARD,inplace,N}) where {T,N,inplace}
    X = CuArray{T}(undef, p.sz)
    handle, workarea = create_plan(p.xtype, p.sz, p.region)
    ScaledPlan(cCuFFTPlan{T,CUFFT_INVERSE,inplace,N}(handle, workarea, X, p.sz, p.region,
                                                     p.xtype),
               normalization(X, p.region))
end

function plan_inv(p::cCuFFTPlan{T,CUFFT_INVERSE,inplace,N}) where {T,N,inplace}
    X = CuArray{T}(undef, p.sz)
    handle, workarea = create_plan(p.xtype, p.sz, p.region)
    ScaledPlan(cCuFFTPlan{T,CUFFT_FORWARD,inplace,N}(handle, workarea, X, p.sz, p.region,
                                                     p.xtype),
               normalization(X, p.region))
end

function plan_inv(p::rCuFFTPlan{T,CUFFT_INVERSE,inplace,N}
                  ) where {T<:cufftComplexes,N,inplace}
    X = CuArray{real(T)}(undef, p.osz)
    Y = CuArray{T}(undef, p.sz)
    xtype = p.xtype == CUFFT_C2R ? CUFFT_R2C : CUFFT_D2Z
    handle, workarea = create_plan(xtype, p.osz, p.region)
    ScaledPlan(rCuFFTPlan{real(T),CUFFT_FORWARD,inplace,N}(handle, workarea, X, p.sz, p.region, xtype),
               normalization(X, p.region))
end

function plan_inv(p::rCuFFTPlan{T,CUFFT_FORWARD,inplace,N}
                  ) where {T<:cufftReals,N,inplace}
    X = CuArray{complex(T)}(undef, p.osz)
    Y = CuArray{T}(undef, p.sz)
    xtype = p.xtype == CUFFT_R2C ? CUFFT_C2R : CUFFT_Z2D
    handle, workarea = create_plan(xtype, p.sz, p.region)
    ScaledPlan(rCuFFTPlan{complex(T),CUFFT_INVERSE,inplace,N}(handle, workarea, X, p.sz,
                                                              p.region, xtype),
               normalization(Y, p.region))
end

## plan execution

# NOTE: "in-place complex-to-real FFTs may overwrite arbitrary imaginary input point values
#       [...]. Out-of-place complex-to-real FFT will always overwrite input buffer."
#       see # JuliaGPU/CuArrays.jl#345, NVIDIA/cuFFT#2714055.

function assert_applicable(p::CuFFTPlan{T}, X::DenseCuArray{T}) where {T}
    (size(X) == p.sz) ||
        throw(ArgumentError("CuFFT plan applied to wrong-size input"))
end

function assert_applicable(p::CuFFTPlan{T,K,inplace}, X::DenseCuArray{T},
                           Y::DenseCuArray) where {T,K,inplace}
    assert_applicable(p, X)
    if size(Y) != p.osz
        throw(ArgumentError("CuFFT plan applied to wrong-size output"))
    elseif inplace != (pointer(X) == pointer(Y))
        throw(ArgumentError(string("CuFFT ",
                                   inplace ? "in-place" : "out-of-place",
                                   " plan applied to ",
                                   inplace ? "out-of-place" : "in-place",
                                   " data")))
    end
end

function unsafe_execute!(plan::cCuFFTPlan{cufftComplex,K,<:Any,N},
                         x::DenseCuArray{cufftComplex,N},
                         y::DenseCuArray{cufftComplex,N}) where {K,N}
    @assert plan.xtype == CUFFT_C2C
    update_stream(plan)
    cufftExecC2C(plan, x, y, K)
end

function unsafe_execute!(plan::rCuFFTPlan{cufftComplex,K,true,N},
                         x::DenseCuArray{cufftComplex,N},
                         y::DenseCuArray{cufftReal,N}) where {K,N}
    @assert plan.xtype == CUFFT_C2R
    update_stream(plan)
    cufftExecC2R(plan, x, y)
end
function unsafe_execute!(plan::rCuFFTPlan{cufftComplex,K,false,N},
                         x::DenseCuArray{cufftComplex,N},
                         y::DenseCuArray{cufftReal}) where {K,N}
    @assert plan.xtype == CUFFT_C2R
    x = copy(x)
    update_stream(plan)
    cufftExecC2R(plan, x, y)
    unsafe_free!(x)
end

function unsafe_execute!(plan::rCuFFTPlan{cufftReal,K,<:Any,N},
                         x::DenseCuArray{cufftReal,N},
                         y::DenseCuArray{cufftComplex,N}) where {K,N}
    @assert plan.xtype == CUFFT_R2C
    update_stream(plan)
    cufftExecR2C(plan, x, y)
end

function unsafe_execute!(plan::cCuFFTPlan{cufftDoubleComplex,K,<:Any,N},
                         x::DenseCuArray{cufftDoubleComplex,N},
                         y::DenseCuArray{cufftDoubleComplex}) where {K,N}
    @assert plan.xtype == CUFFT_Z2Z
    update_stream(plan)
    cufftExecZ2Z(plan, x, y, K)
end

function unsafe_execute!(plan::rCuFFTPlan{cufftDoubleComplex,K,true,N},
                         x::DenseCuArray{cufftDoubleComplex,N},
                         y::DenseCuArray{cufftDoubleReal}) where {K,N}
    update_stream(plan)
    @assert plan.xtype == CUFFT_Z2D
    cufftExecZ2D(plan, x, y)
end
function unsafe_execute!(plan::rCuFFTPlan{cufftDoubleComplex,K,false,N},
                         x::DenseCuArray{cufftDoubleComplex,N},
                         y::DenseCuArray{cufftDoubleReal}) where {K,N}
    @assert plan.xtype == CUFFT_Z2D
    x = copy(x)
    update_stream(plan)
    cufftExecZ2D(plan, x, y)
    unsafe_free!(x)
end

function unsafe_execute!(plan::rCuFFTPlan{cufftDoubleReal,K,<:Any,N},
                         x::DenseCuArray{cufftDoubleReal,N},
                         y::DenseCuArray{cufftDoubleComplex,N}) where {K,N}
    @assert plan.xtype == CUFFT_D2Z
    update_stream(plan)
    cufftExecD2Z(plan, x, y)
end

## high-level integrations

function LinearAlgebra.mul!(y::DenseCuArray{Ty}, p::CuFFTPlan{T}, x::DenseCuArray{T}
                           ) where {Ty, T}
    assert_applicable(p,x,y)
    unsafe_execute!(p,x,y)
    return y
end

function Base.:(*)(p::cCuFFTPlan{T,K,true,N}, x::DenseCuArray{T,N}) where {T,K,N}
    assert_applicable(p,x)
    unsafe_execute!(p,x,x)
    x
end

function Base.:(*)(p::rCuFFTPlan{T,CUFFT_FORWARD,false,N}, x::DenseCuArray{T,N}
           ) where {T<:cufftReals,N}
    assert_applicable(p,x)
    @assert p.xtype ∈ [CUFFT_R2C,CUFFT_D2Z]
    y = CuArray{complex(T),N}(undef, p.osz)
    unsafe_execute!(p,x,y)
    y
end

function Base.:(*)(p::rCuFFTPlan{T,CUFFT_INVERSE,false,N}, x::DenseCuArray{T,N}
           ) where {T<:cufftComplexes,N}
    assert_applicable(p,x)
    @assert p.xtype ∈ [CUFFT_C2R,CUFFT_Z2D]
    y = CuArray{real(T),N}(undef, p.osz)
    unsafe_execute!(p,x,y)
    y
end

function Base.:(*)(p::cCuFFTPlan{T,K,false,N}, x::DenseCuArray{T,N}) where {T,K,N}
    assert_applicable(p,x)
    y = CuArray{T,N}(undef, p.osz)
    unsafe_execute!(p,x,y)
    y
end


fetchvcat(J::judiVector) = J
fetchvcat(x) = vcat(x...)

function JUDI.time_modeling(model::Model, srcGeometry, srcData, recGeometry, recData, dm, srcnum::UnitRange{Int64}, op::Char, mode::Int64, options)

    # Broadcast common parameters
    _model = @bcast model
    _dm = isnothing(dm) ? dm : @bcast dm
    opts = make_opt([model, dm, op, srcGeometry])
    iter = make_parts(srcnum)
    # Run on azure
    results = @batchexec pmap(sx -> time_modeling_azure(_model, subsample(srcGeometry, sx), subsample(srcData, sx),
                                                        subsample(recGeometry, sx), subsample(recData, sx), _dm,
                                                        op, mode, subsample(options, sx)), iter) opts
    # Gather results
    if op=='F' || (op=='J' && mode==1)
        temp = fetch(results)
        argout1 = fetchvcat(temp)
    elseif op=='J' && mode==-1
        argout1 = fetchreduce(results; op=+, remote=true)
    else
        error("operation no defined")
    end
    return argout1
end

# FWI
function JUDI.fwi_objective(model::Model, source::judiVector, dObs::judiVector; options=JUDI.Options())
    # fwi_objective function for multiple sources. The function distributes the sources and the input data amongst the available workers.
    # Broadcast common parameters
    _model = @bcast model
    opts = make_opt([model, dObs, source])
    iter = make_parts(1:dObs.nsrc)
    results = @batchexec pmap(j -> fwi_objective_azure(_model, source[j], dObs[j], subsample(options, j)), iter) opts

    # Collect and reduce gradients
    obj, gradient = fetchreduce(results; op=+, remote=true)
    # first value corresponds to function value, the rest to the gradient
    return obj, gradient
end

# lsrtm

function JUDI.lsrtm_objective(model::Model, source::judiVector, dObs::judiVector, dm; options=JUDI.Options(), nlind=false)
    # lsrtm_objective function for multiple sources. The function distributes the sources and the input data amongst the available workers.

    # Broadcast common parameters
    _model = @bcast model
    _dm = isnothing(dm) ? dm : @bcast dm

    opts = make_opt([model, dObs, source, dm])
    iter = make_parts(1:dObs.nsrc)
    results = @batchexec pmap(j -> lsrtm_objective_azure(_model, source[j], dObs[j], _dm, subsample(options, j); nlind=nlind), iter) opts
    # Collect and reduce gradients
    obj, gradient = fetchreduce(results; op=+, remote=false)
    delete_job(results)

    # first value corresponds to function value, the rest to the gradient
    return obj, gradient
end

function JUDI.lsrtm_objective(model::Array{Model,1}, source::Array{judiVector{T,Array{T,2}},1}, dObs::Array{judiVector{T,Array{T,2}},1}, dm::Union{Array{Array{T,1},1}, Array{PhysicalParameter{T},1}}; options=Options(), nlind=false) where T
# lsrtm_objective function for multiple sources and multiple vintages. The function distributes the sources and the input data amongst the available workers.

    println("serial lsrtm_objective on each vintage now")
    ctrl = Array{Any, 1}(undef, length(dObs))
    obj = Array{T, 1}(undef, length(dObs))
    gradient = Array{Array{T,2}, 1}(undef, length(dObs))

    for i = 1:length(dObs)
        _model = model[i]
        _source = source[i]
        _dObs = dObs[i]
        _dm = dm[i]
        _model1 = @bcast _model
        _dm1 = @bcast _dm
        opts = make_opt([_model, dObs, source, _dm])
        ctrl[i] = @batchexec pmap(j -> lsrtm_objective_azure(_model1, _source[j], _dObs[j], _dm1, subsample(options, j); nlind=nlind), 1:dObs[i].nsrc) opts
    end
    for i = 1:length(dObs)
        results = fetch(ctrl[i])
        delete_job(ctrl[i])
        obj[i] = sum([results[j][1] for j =  1:dObs[i].nsrc])
        gradient[i] = sum([results[j][2] for j = 1:dObs[i].nsrc])
    end

    obj1 = sum(obj)

    # first value corresponds to function value, the rest to the gradient
    return obj1, gradient
end

# Make options for unique batch ids
make_opt(a::Array) = AzureClusterlessHPC.Options(;task_name="task_$(objectid(a))")

# Batchexec function
slicec(i::Integer, j::Integer) = i==j ? i : (i:j)
chunk(arr, n) = [arr[slicec(i, min(i + n - 1, length(arr)))] for i in 1:n:length(arr)]
make_parts(iter) = _njpi > 1 ? chunk(iter, _njpi) : iter


function set_wally(nodes, bdcon, wally, cellcenter, cellxmax, cellymax, icell)
	# serch wall
	swith_wall = zeros(4)   # x-, x+, y-, y+を検索する
							# 1:壁
	for i in 1:4
		if Int(bdcon[i][1]) == 2 || Int(bdcon[i][1]) == 7
			swith_wall[i] = 1
		end
	end

	# cal number of ponits
	nop = 0
	if swith_wall[1] == 1
		nop = nop + (cellymax-1)
	end
	if swith_wall[2] == 1
		nop = nop + (cellymax-1)
	end
	if swith_wall[3] == 1
		nop = nop + (cellxmax-1)
	end
	if swith_wall[4] == 1
		nop = nop + (cellxmax-1)
	end

	wallpoint = zeros(nop, 2) # (格子数, 座標)
	ite = 1

	# x-
	if swith_wall[1] == 1
		for j in 1+icell:cellymax-icell
			wallpoint[ite,1] = nodes[icell+1,j,1]
			wallpoint[ite,2] = nodes[icell+1,j,2]
			ite = ite + 1
		end
	end
	# x+
	if swith_wall[2] == 1
		for j in 1+icell:cellymax-icell
			wallpoint[ite,1] = nodes[cellxmax-icell+1,j,1]
			wallpoint[ite,2] = nodes[cellxmax-icell+1,j,2]
			ite = ite + 1
		end
	end
	# y-
	if swith_wall[3] == 1
		for i in 2:cellxmax-1
			wallpoint[ite,1] = nodes[i,icell+1,1]
			wallpoint[ite,2] = nodes[i,icell+1,2]
			ite = ite + 1
		end
	end
	# y+
	if swith_wall[4] == 1
		for i in 2:cellxmax-1
			wallpoint[ite,1] = nodes[i,cellymax-icell+1,1]
			wallpoint[ite,2] = nodes[i,cellymax-icell+1,2]
			ite = ite + 1
		end
	end

	# 距離格納用
	distance_list = zeros(nop,2)
	distance = 0

	# 最大値を2つ抽出
	for j in 1+icell:cellymax-icell
		for i in 1+icell:cellxmax-icell
			for n in 1:nop
				distance = ((cellcenter[i,j,1]-wallpoint[n,1])^2 +
							(cellcenter[i,j,2]-wallpoint[n,2])^2
							)^0.5
				distance_list[n,1] = n
				distance_list[n,2] = distance
			end

			qdistance_list = copy(distance_list)
			qdistance_list = quicksort(qdistance_list,1,nop,2,2)

			# wall_point
			point1 = Int(qdistance_list[1,1])
			point2 = Int(qdistance_list[2,1])

			dis_bottom = ((wallpoint[point1, 1] - wallpoint[point2, 1])^2 + (wallpoint[point1, 2] - wallpoint[point2, 2])^2)^0.5

			# O型格子の関係で同じ点が二つあるため
			if dis_bottom == 0.0
				point1 = Int(qdistance_list[1,1])
				point2 = Int(qdistance_list[3,1])

				dis_bottom = ((wallpoint[point1, 1] - wallpoint[point2, 1])^2 + (wallpoint[point1, 2] - wallpoint[point2, 2])^2)^0.5
			end

			dis1 = distance_list[point1, 2]
			dis2 = distance_list[point2, 2]

			# ヘロンの公式から距離を算出
			s = (dis_bottom + dis1 + dis2)*0.5
			surface = (s*(s-dis_bottom)*(s-dis1)*(s-dis2))^0.5
			wally[i,j] = surface/dis_bottom * 2
		end
	end

	# 仮想セルへの代入
	for i in 1+icell:cellxmax-icell
		for j in 1:icell
			wally[i,j] = wally[i,icell+1]
			wally[i,cellymax-j+1] = wally[i,cellymax-icell]
		end
	end
	for j in 1+icell:cellymax-icell
		for i in 1:icell
			wally[i,j] = wally[icell+1,j]
			wally[cellxmax-i+1,j] = wally[cellxmax-icell,j]
		end
	end

	return wally, swith_wall
end

function cal_yplus(yplus, Qbase, wally, swith_wall, mu, cellxmax, cellymax, vecAx, vecAy, volume, icell)
	tvAxx = 0.0
    tvAxy = 0.0
    Axx = 0.0
    Axy = 0.0
    tvAyx = 0.0
    tvAyy = 0.0
    Ayx = 0.0
    Ayy = 0.0

	if swith_wall[1] == 1
		for j in 1+icell:cellymax-icell
			# 壁に並行な方向の速度を求めるためx-の壁面に対しvecAyを使用
			for i in 1:cellxmax
				tvAyx = 0.5*(vecAy[i,j-1,1]+vecAy[i,j,1])
				tvAyy = 0.5*(vecAy[i,j-1,2]+vecAy[i,j,2])
				Ayx = tvAyx / (tvAyx^2+tvAyy^2)^0.5
				Ayy = tvAyy / (tvAyx^2+tvAyy^2)^0.5
				nu = mu[i,j] / Qbase[i,j,1]
				ubar = abs(Ayx*Qbase[i,j,2] + Ayy*Qbase[i,j,3])
				yplus[i,j], up = cal_wall_val_spalding(ubar, wally[i,j], nu)
			end
		end
	end
	if swith_wall[2] == 1
		for j in 1+:cellymax-icell
			# 壁に並行な方向の速度を求めるためx+の壁面に対しvecAyを使用
			for i in 1:cellxmax
				tvAyx = 0.5*(vecAy[i,j-1,1]+vecAy[i,j,1])
				tvAyy = 0.5*(vecAy[i,j-1,2]+vecAy[i,j,2])
				Ayx = tvAyx / (tvAyx^2+tvAyy^2)^0.5
				Ayy = tvAyy / (tvAyx^2+tvAyy^2)^0.5
				nu = mu[i,j] / Qbase[i,j,1]
				ubar = abs(Ayx*Qbase[i,j,2] + Ayy*Qbase[i,j,3])
				yplus[i,j], up = cal_wall_val_spalding(ubar, wally[i,j], nu)
			end
		end
	end
	if swith_wall[3] == 1
		for j in 1:cellymax
			for i in 1+icell:cellxmax-icell
			# 壁に並行な方向の速度を求めるためy-の壁面に対しvecAxを使用

				tvAxx = 0.5*(vecAx[i-1,j,1]+vecAx[i,j,1])
				tvAxy = 0.5*(vecAx[i-1,j,2]+vecAx[i,j,2])
				Axx = tvAxx / (tvAxx^2+tvAxy^2)^0.5
				Axy = tvAxy / (tvAxx^2+tvAxy^2)^0.5
				nu = mu[i,j] / Qbase[i,j,1]
				ubar = abs(Axx*Qbase[i,j,2] + Axy*Qbase[i,j,3])
				yplus[i,j], up = cal_wall_val_spalding(ubar, wally[i,j], nu)

				if yplus[i,j] > 500
					break
				end
			end
		end
	end
	if swith_wall[4] == 1
		for j in 1:cellymax
			# 壁に並行な方向の速度を求めるためy-の壁面に対しvecAxを使用
			for i in 1+icell:cellxmax-icell
				tvAxx = 0.5*(vecAx[i-1,j,1]+vecAx[i,j,1])
				tvAxy = 0.5*(vecAx[i-1,j,2]+vecAx[i,j,2])
				Axx = tvAxx / (tvAxx^2+tvAxy^2)^0.5
				Axy = tvAxy / (tvAxx^2+tvAxy^2)^0.5
				nu = mu[i,j] / Qbase[i,j,1]
				ubar = abs(Axx*Qbase[i,j,2] + Axy*Qbase[i,j,3])
				yplus[i,j], up = cal_wall_val_spalding(ubar, wally[i,j], nu)
			end
		end
	end

	return yplus
end

function cal_wall_val_spalding(u,y,nu)
	# spalding 定数
	k = 0.4
	B = 5.5

	# ニュートン法
	up = 200.0
	for i in 1:100
		old = up

		up = up - spalding(up,u,y,nu,k,B)/spalding_dash(up,u,y,nu,k,B)

		if abs(up-old)/up < 10^(-4)
			break
		end
	end

	yplus = u*y/nu/up
	return yplus, up
end

function spalding(up,u,y,nu,k,B)
	t = k*up
	F = - 1/up* u*y/nu + up + exp(-k*B)*(exp(t)-1-t-t^2/2-t^3/6)
	return F
end

function spalding_dash(up,u,y,nu,k,B)
	t = k*up
	Fdash = 1/up^2* u*y/nu + 1 + exp(-k*B)*(k*exp(t)-k-k^2*up-k^3*up^2/2)
	return Fdash
end

# test
#=
uu = 6.64037328
yy = 0.06299976860
nunu = 1.406107133e-5
println(cal_wall_val_spalding(uu,yy,nunu))
throw(UndefVarError(:x))
=#

function wallf_Van_Driest(yplus)
	f = 1 - exp(-yplus/26)
	return f
end


using Juriba.Radiation
using Juriba.Constants
using StaticArrays, QuadGK, Interpolations

const RD = Juriba.Radiation

@Base.kwdef mutable struct Compton <: AbstractRadiation

    size::Int					        # size of arrays
    en_phot::Vector{Float64}			# array of photon energies
    num_phot::Vector{Float64} 			# array of number of photons in units of erg/s/Hz
    en_phot_obs::Vector{Float64}		# same as above but in observer frame
    num_phot_obs::Vector{Float64} 		# same as above but in observer frame

    seed_ph::Interpolations.Extrapolation
    iter_ph::InterruptException.Extrapolation
    esc_p_sph::Interpolations.Extrapolation
    esc_p_cyl::Interpolation.Extrapolation

    z::Float64 = 0.				            # Dimensions of emitting region
    r::Float64 = 0.
    vol::Float64 = 0.					    # Volume of emitting region
    beta::Float64 = 0.				        # speed of the emitting region
    angle::Float64 = 0.		        # Viewing angle/Doppler factor of emitting region
    dopfac::Float64 = 0.
    dopnum::Float64 = 0.				        # Doppler boosting exponent, depends on geometry
    counterjet::Bool = false			# boolean switch if user wants to include counterjet emission
    geometry::String = ""   	        # string to track geometry of emitting region

    seed_size::Int					# size of seed photon field
    Niter::Int						# number of IC iterations
    tau::Float64                    # optical depth/comtpon Y of emitting region
    ypar::Float64
    rphot::Float64 					# photospheric radius when tau > 1, used to renormalize volume
    escape_corr::Float64				# escape term, used to renormalize our spectra to CompPS

    seed_energ::Float64				# array of seed frequencies in Hz
    seed_urad::Float64 				# array of seed photon number density in log10(#/erg/cm^3)
    iter_urad::Float64 				# array of iterated photon number density in log10(#/erg/cm^3)

end

const Te_table = @SVector [19.999,30.,50.,70.,90.,120.,150.,200.,
                  280.,400.,600.,900.,1300.,1800.,2500.001]
const tau_table	= @SVector [0.04999,0.07,0.1,0.19,0.27,0.38,0.54,0.76,
                   0.94,1.08,1.34,1.75,2.1,2.6,3.0001]

esc_table_sph = SMatrix{15, 15}(
                    0.02,0.07,0.17,0.27,0.33,0.4,0.45,0.45,0.55,0.55,0.54,0.51,0.48,0.55,0.55,
                    0.02,0.07,0.18,0.275,0.33,0.4,0.45,0.47,0.55,0.53,0.56,0.51,0.5,0.55,0.55,
                    0.02,0.07,0.18,0.275,0.34,0.4,0.45,0.47,0.55,0.54,0.52,0.51,0.53,0.55,0.58,
                    0.025,0.085,0.18,0.275,0.35,0.4,0.45,0.47,0.55,0.54,0.59,0.57,0.5,0.55,0.6,
                    0.028,0.085,0.2,0.275,0.35,0.4,0.45,0.49,0.54,0.54,0.6,0.59,0.53,0.55,0.65,
                    0.035,0.09,0.21,0.29,0.35,0.4,0.45,0.49,0.55,0.57,0.58,0.6,0.53,0.55,0.7,
                    0.045,0.12,0.23,0.3,0.35,0.41,0.42,0.49,0.54,0.57,0.58,0.58,0.5,0.55,0.7,
                    0.055,0.135,0.23,0.3,0.35,0.4,0.42,0.49,0.52,0.57,0.59,0.61,0.58,0.57,0.85,
                    0.07,0.14,0.235,0.3,0.35,0.4,0.41,0.49,0.5,0.57,0.58,0.61,0.61,0.57,0.85,
                    0.08,0.145,0.235,0.3,0.35,0.39,0.4,0.49,0.5,0.57,0.58,0.63,0.7,0.6,0.85,
                    0.09,0.15,0.245,0.3,0.35,0.4,0.4,0.45,0.49,0.55,0.62,0.64,0.73,0.6,0.85,
                    0.105,0.17,0.26,0.31,0.35,0.4,0.4,0.45,0.47,0.56,0.62,0.65,0.78,0.65,0.85,
                    0.115,0.19,0.28,0.32,0.35,0.4,0.4,0.45,0.44,0.5,0.67,0.68,0.8,0.7,1.,
                    0.14,0.215,0.3,0.34,0.38,0.4,0.4,0.45,0.42,0.48,0.68,0.73,0.85,0.8,1.,
                    0.16,0.23,0.31,0.35,0.38,0.4,0.4,0.45,0.42,0.48,0.67,0.78,0.9,0.9,1.
                    )

esc_table_cyl = SMatrix{15, 15}(
                    0.02,0.06,0.22,0.3,0.5,0.5,0.57,0.7,0.75,0.75,0.75,0.77,0.75,0.76,0.7,
                    0.02,0.06,0.19,0.28,0.4,0.5,0.57,0.65,0.65,0.75,0.7,0.75,0.7,0.7,0.75,
                    0.017,0.04,0.17,0.27,0.32,0.5,0.57,0.6,0.6,0.7,0.65,0.65,0.73,0.67,0.7,
                    0.01,0.04,0.12,0.17,0.23,0.35,0.4,0.55,0.5,0.55,0.55,0.65,0.7,0.67,0.68,
                    0.01,0.04,0.1,0.17,0.19,0.28,0.32,0.4,0.45,0.5,0.5,0.5,0.68,0.66,0.68,
                    0.015,0.04,0.09,0.17,0.19,0.25,0.25,0.35,0.45,0.4,0.45,0.5,0.64,0.64,0.68,
                    0.02,0.05,0.09,0.15,0.17,0.22,0.24,0.3,0.35,0.35,0.45,0.4,0.62,0.61,0.64,
                    0.02,0.05,0.095,0.15,0.17,0.22,0.22,0.25,0.3,0.35,0.4,0.38,0.4,0.4,0.6,
                    0.024,0.05,0.095,0.14,0.16,0.21,0.21,0.25,0.27,0.32,0.4,0.42,0.4,0.4,0.5,
                    0.025,0.055,0.1,0.15,0.17,0.21,0.21,0.25,0.27,0.3,0.35,0.4,0.35,0.4,0.43,
                    0.03,0.065,0.11,0.15,0.18,0.21,0.21,0.25,0.27,0.3,0.35,0.4,0.4,0.45,0.47,
                    0.045,0.085,0.135,0.17,0.19,0.22,0.24,0.27,0.3,0.35,0.4,0.45,0.5,0.5,0.55,
                    0.06,0.1,0.16,0.2,0.22,0.25,0.26,0.29,0.33,0.35,0.45,0.5,0.5,0.55,0.65,
                    0.09,0.14,0.2,0.23,0.26,0.29,0.29,0.33,0.4,0.4,0.48,0.55,0.6,0.65,0.75,
                    0.11,0.165,0.23,0.26,0.29,0.31,0.33,0.36,0.4,0.45,0.52,0.6,0.65,0.75,0.8
                    )

function compton(s1::Int, s2::Int)

    size = s1
    seed_size = s2

    en_phot = zeros(Float64, size)
    num_phot = similar(en_phot)
    iter_urad = similar(en_phot)
    en_phot_obs = zeros(Float64, 2*size)
    num_phot_obs = similar(en_phot_obs)

    seed_energ = zeros(Float64, seed_size)
    seed_urad = zeros(Float64, seed_size)

    Niter = 20
    ypad = 0
    escape_corr = 1.0

    counterjet = false

    # xsize = 15
    # ysize = 15

    esc_p_sph = CubicSplineInterpolation((Te_table, tau_table),esc_table_sph)
    esc_p_cyl = CubicSplineInterpolation((Te_table, tau_table),esc_table_cyl)

    return Compton(size=size, seed_size=seed_size, en_phot=en_phot,
                   num_phot=num_phot, iter_urad=iter_urad, en_phot_obs=en_phot_obs,
                   num_phot_obs=num_phot_obs, seed_energ=seed_energ,
                   seed_urad=seed_urad, Niter=Niter, ypar=ypar,
                   escape_corr=escape_corr, esc_p_cyl=esc_p_cyl, esc_p_sph=esc_p_sph)
end

# This function is the kernel of eq 2.48 in Blumenthal & Gould(1970), represents the scattered photon spectrum
# for a given electron and includes the Klein-Nishina cross section.
function comfnc(ein, p::RD.comfnc_params)
    game = p.game
    e1 = p.e1

    phodis = p.phodis

    einit = exp(ein)
    btst = einit/(game*emerg)
    eg4 = 3.0*einit*game
    utst = eg4*(emerg + eg4)

    if ((e1 < btst && (btst-e1)/btst >= 4.0e-8) ||
        (e1 > utst && (e1-utst)/utst >= 4.0e-8))
        return 0.0
    else
        biggam = eg4/emerg
        q = e1/(biggam*(1.0 - e1))
        phonum = phodis(einit)
        tm1 = 2.0*q*log(q)
        tm2 = (1.0 + 2.0*q)*(1.0 - q)
        tm3 = 0.5*((biggam*q)^2)*(1.0 - q)/(1.0 + biggam*q)

        return (tm1 + tm2 + tm3)*10^phonum
    end

end

# This function is the integral of comfnc above over the total seed photon distribution
function comint(gam, p::comint_params)

    eph = p.eph
    ephmin = p.ephmin
    ephmax = p.ephmax

    eldis = p.eldis
    phodis = p.phodis

    game = exp(gam)
    e1 = eph/(game*emerg)
    econst = 2π*re0^2*cee
    blim = log(max(eph/(4.0*game*(game-eph/emerg)), ephmin))
    ulim = log(min(eph, ephmax))

    (ulim <= blim) && return 0.0

    F2params = RD.comfnc_params(game, e1, phodis)
    F2func(x) = comfnc(x, F2params)

    result, error = quadgk(F2func, blim, ulim, rtol=1.0, atol=1.0, maxevals=100)

    den = eldis(game)

    return econst*den*result/game
end

"""This integrates the individual electron spectrum from
comint over the total electron distribution"""
function comintegral(com::Compton, it, blim, ulim, enphot, enphmin, enphmax,
                     eldis)

    F1params = RD.comint_params(enphot, enphmin, enphmax,
                                eldis, com.seed_ph)
    F1params_it = RD.comint_params(enphot, enphmin, enphmax,
                                   eldis, com.iter_ph)

    if iszero(it)
        params = F1params
    else
        params = F1params_it
    end

    F1func(gam) = comint(gam, params)

    result, errror = quadgk(F1func, blim, ulim, rtol=1.0, atol=1.0, maxevals=100)
    return result
end

"""
This function calculates the final spectrum in all frequency bins,
including multiple scatters.
Note: the reason the Doppler boosting is a factor of 2 instead of 3 is because the calculations are done for
a conical jet in the Lind&BLanford 1985 prescription
"""
function compton_spectrum(com::Compton, gmin, gmax, eldis)

    ephmin = com.seed_energ[1]
    ephmax = com.seed_energ[end]

    dopfac_cj = com.dopfac*(1.0 - com.beta*cos(com.angle))/
                (1.0 +  com.beta*cos(com.angle))

    for it = 1:com.Niter
        for i = 1:com.size
            blim = log(max(gmin, com.en_phot[i]/emerg))
            ulim = log(gmax)

            if blim >= ulim
                comres = 1e-100
            else
                comres = comintegral(com, it, blim, ulim, com.en_phot[i],
                                  ephmin, ephmax, eldis)
            end

            com.num_phot[i] = com.num_phot[i]+comres*com.vol*com.en_phot[i]*herg
            com.en_phot_obs[i] = com.en_phot[i]*com.dopfac
            com.num_phot_obs[i] = com.num_phot[i]*dopfac_cj^com.dopnum

            if com.counterjet
                com.en_phot_obs[i+com.size] = com.en_phot[i]*dopfac_cj
                com.num_phot_obs[i+com.size] = com.num_phot[i]*dopfac_cj^com.dopnum
            else
                com.en_phot_obs[i+com.size] = 0.0
                com.num_phot_obs[i+com.size] = 0.0
            end

            if iszero(comres)
                com.iter_urad[i] = -50.0
            else
                com.iter_urad[i] = log10(com.escape_corr*comres*com.vol)/(π*com.r^2*cee)
            end

        end

        ephmin = com.en_phot[1]
        ephmax = com.en_phot[end]
        com.iter_ph = LinearInterpolation(com.en_phot, com.iter_urad)
    end
end

"""Method to include cyclosynchrotron array from `Cyclosyn.jl` to the seed field. The two input arrays should be
`get_energ()` and `get_nphot()` methods from `Cyclosyn.jl`. If the seed_urad array wasn't empty, the contribution
is automatically added to the existing one.
note: the second condition has a <= sign to dodge numerical errors when the seed field energy density is
extremely low, which can result in negative values/nan for the seed_urad """
function cyclosyn_seed(com::Compton, seed_arr, seed_lum)
    chr2 = cee*herg*π*com.r^2
    for i in eachindex(seed_arr)
        com.seed_energ[i] = seed_arr[i]

        if !iszero(com.seed_urad[i])
            com.seed_urad[i] = log10(10^com.seed_urad[i] + seed_lum[i]/(chr2*com.seed_energ[i]))
        elseif seed_lum[i]/(chr2*com.seed_energ[i]) <= 0
            com.seed_urad[i] = -100
        else
            com.seed_urad[i] = log10(seed_lum[i]/(chr2*com.seed_energ[i]))
        end

    end

    com.seed_ph = LinearInterpolation(seed_energ, seed_urad)

end

"""
Method to include black body to seed field for IC;
Note: Urad and Tbb need to be passed in the co-moving frame, the function does NOT account for beaming
"""
function bb_seed_k(com::Compton, seed_arr, Urad, Tbb)

    ulim = 100.0*Tbb*kboltz
    bfield_const = herg*cee^2*sbconst*Tbb^4

    seed_freq_array(com, seed_arr)
    for i = 1:com.seed_size
        if (com.seed_energ[i] < ulim)
            bbfield = (2.0*Urad*(seed_energ[i]/herg)^2.)/(bfield_const*(exp(seed_energ[i]/(kboltz*Tbb)) - 1.0))
        else
            bbfield = 1.0e-100
        end

        if !iszero(com.seed_urad[i])
            com.seed_urad[i] = log10(10^com.seed_urad[i] + bbfield)
        elseif bbfield > 0
            com.seed_urad[i] = log10(bbfield)
        else
            com.seed_urad[i] = -100.0
        end

    end

    com.seed_ph = LinearInterpolation(com.seed_energ, com.seed_urad)

end

function bb_seed_kev(com::Compton, seed_arr, Urad, Tbb)

    ulim = 100.0*Tbb*kboltz_kev2erg
    bfield_const = herg*cee^2*sbconst*(Tbb*kboltz_kev2erg/kboltz)^4

    seed_freq_array(com, seed_arr)
    for i = 1:com.seed_size
        if (com.seed_energ[i] < ulim)
            bbfield = (2.0*Urad*(seed_energ[i]/herg)^2.)/(bfield_const*(exp(seed_energ[i]/(kboltz*Tbb)) - 1.0))
        else
            bbfield = 1.0e-100
        end

        if !iszero(com.seed_urad[i])
            com.seed_urad[i] = log10(10^com.seed_urad[i] + bbfield)
        elseif bbfield > 0
            com.seed_urad[i] = log10(bbfield)
        else
            com.seed_urad[i] = -100.0
        end

    end

    com.seed_ph = LinearInterpolation(com.seed_energ, com.seed_urad)

end

"""
Calculates incident photon field for a Shakura-Sunyaev disk with aspect ratio H, inner temperature Tin,
inner radius Rin, at a distance z from the disk, taking beaming into account
"""
function disk_integral(alfa, p)

end

function seed_freq_array(com::Compton, seed_arr)
    com.seed_energ .= seed_arr
end


"""
`module  JAC.ImpactIonization`
    ... a submodel of JAC that contains all methods for computing electron impact ionization cross sections.
"""
module ImpactIonization

    using JAC, JAC.ManyElectron, JAC.ImpactExcitation, JAC.AutoIonization

    """
    `struct  ImpactIonization.Settings`  ... defines a type for the settings in estimating electron-impact ionization cross sections.

        + electronEnergies   ::Array{Float64,1}             ... List of impact-energies of the incoming elecgtrons.
        + calcShellDependent ::Bool                         ... True of shell-dependent cross sections are to be calculated, and false otherwise.
        + shells             ::Array{Shell,1}               ... List of shells for which shell-dependent cross sections are calculated.
    """
    struct Settings
        electronEnergies     ::Array{Float64,1}
        calcShellDependent   ::Bool
        shells               ::Array{Shell,1}
    end

    """
    `JAC.ImpactIonization.Settings()`  ... constructor for the default values of electron-impact ionization cross section estimates.
    """
    function Settings()
       Settings( Float64[], false, Shell[] )
    end

    ## `Base.show(io::IO, settings::ImpactIonization.Settings)`  ... prepares a proper printout of the variable settings::ImpactIonization.Settings.
    function Base.show(io::IO, settings::ImpactIonization.Settings)
        println(io, "electronEnergies:       $(settings.electronEnergies)  ")
        println(io, "calcShellDependent:     $(settings.calcShellDependent)  ")
        println(io, "shells:                 $(settings.shells)  ")
    end

    """
    `struct  ImpactIonization.Channel`
        ... defines a type for a electron-impact ionization channel to help characterize the incoming and two outgoing (continuum)
            states of many electron-states with (two) free electron

        + initialEpsilon   ::Float64            ... energy of the incoming free-electron
        + finalEpsilon     ::Float64            ... energy of the outgoing free-electron
        + releasedEpsilon  ::Float64            ... energy of the released free-electron
        + initialKappa     ::Int64              ... partial-wave of the incoming free electron
        + finalKappa       ::Int64              ... partial-wave of the outgoing free electron
        + releasedKappa    ::Int64              ... partial-wave of the released free electron
        + symmetry         ::LevelSymmetry      ... total angular momentum and parity of the scattering state
        + initialPhase     ::Float64            ... phase of the incoming partial wave
        + finalPhase       ::Float64            ... phase of the outgoing partial wave
        + releasedPhase    ::Float64            ... phase of the released partial wave
        + amplitude        ::Complex{Float64}   ... Collision amplitude associated with the given channel.
    """
    struct  Channel
        initialEpsilon     ::Float64
        finalEpsilon       ::Float64
        releasedEpsilon    ::Float64
        initialKappa       ::Int64
        finalKappa         ::Int64
        releasedKappa      ::Int64
        symmetry           ::LevelSymmetry
        initialPhase       ::Float64
        finalPhase         ::Float64
        releasedPhase      ::Float64
        amplitude          ::Complex{Float64}
    end

    # `Base.show(io::IO, channel::ImpactIonization.Channel)`  ... prepares a proper printout of the variable channel::ImpactIonization.Channel.
    function Base.show(io::IO, channel::ImpactIonization.Channel)
        println(io, "initialEpsilon:     $(channel.initialEpsilon)  ")
        println(io, "finalEpsilon:       $(channel.finalEpsilon)  ")
        println(io, "releasedEpsilon:    $(channel.releasedEpsilon)  ")
        println(io, "initialKappa:       $(channel.initialKappa)  ")
        println(io, "finalKappa:         $(channel.finalKappa)  ")
        println(io, "releasedKappa:      $(channel.releasedKappa)  ")
        println(io, "symmetry:           $(channel.symmetry)  ")
        println(io, "initialPhase:       $(channel.initialPhase)  ")
        println(io, "finalPhase:         $(channel.finalPhase)  ")
        println(io, "releasedPhase:      $(channel.releasedPhase)  ")
        println(io, "amplitude:          $(channel.amplitude)  ")
    end

    """
    `struct  ImpactIonization.Line`  ... defines a type for a electron-impact ionization line that may include the definition of channels and
                                         their corresponding amplitudes.

        + initialLevel   ::Level          ... initial-(state) level
        + finalLevel     ::Level          ... final-(state) level
        + crossSection   ::Float64        ... total cross section of this line
        + hasChannels    ::Bool           ... Determines whether the individual scattering (sub-) channels are defined in terms of their free-
                                              electron energies, kappa's, phases and the total angular momentum/parity as well as the amplitude,
                                              or not.
        + channels       ::Array{ImpactIonization.Line,1}  ... List of Eimex channels of this line.
    """
    struct  Line
        initialLevel     ::Level
        finalLevel       ::Level
        crossSection     ::Float64
        hasChannels      ::Bool
        channels         ::Array{ImpactIonization.Channel,1}
    end

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

end # module


export md_hmac, fp_prime_get

# Module initialization function
function __init__()
    # make sure the native lib was configure correctly
    @assert sizeof(Limb) == LIMB_SIZE
    @assert MD_LEN == 32  # SHA256
    @assert BN_SIZE >= 32  # 256-bits
    @assert BN_ST_SIZE == sizeof(BN)
    @assert FP_ST_SIZE == sizeof(FP)
    @assert FP2_ST_SIZE == sizeof(FP2)
    #@assert FP3_ST_SIZE == sizeof(FP3)
    @assert G1_ST_SIZE == sizeof(EP)
    @assert G2_ST_SIZE == sizeof(EP2)
    Threads.@threads for i = 1:Threads.nthreads()
        iszero(ccall((:core_init, LIB), Cint, ())) || error("core_init failed")
        iszero(ccall((:ep_param_set_any_pairf, LIB), Cint, ())) || error("ep_param_set_any_pairf failed")
        @info "$(basename(LIB)) initialised successfully (on Thread $(Threads.threadid()))!"
    end
end

# the data is always last so this should be safe for both BN and BigInt
unsafe_bnptr(bn::BN) = Ptr{Limb}(pointer_from_objref(bn) + sizeof(bn) - sizeof(bn.dp))

function bn_add_dig!(c::BN, a::BN, b::Limb)
    ccall((:bn_add_dig, LIB), Cvoid, (Ref{BN}, Ref{BN}, Limb), c, a, b)
    return c
end
function bn_add!(c::BN, a::BN, b::BN)
    ccall((:bn_add, LIB), Cvoid, (Ref{BN}, Ref{BN}, Ref{BN}), c, a, b)
    return c
end

function bn_mod_basic!(c::BN, a::BN, m::BN)
    ccall((:bn_mod_basic, LIB), Cvoid, (Ref{BN}, Ref{BN}, Ref{BN}), c, a, m)
    return c
end

function bn_mul_dig!(c::BN, a::BN, b::Limb)
    ccall((:bn_mul_dig, LIB), Cvoid, (Ref{BN}, Ref{BN}, Limb), c, a, b)
    return c
end
function bn_mul_comba!(c::BN, a::BN, b::BN)
    ccall((:bn_mul_comba, LIB), Cvoid, (Ref{BN}, Ref{BN}, Ref{BN}), c, a, b)
    return c
end
function bn_neg!(c::BN, a::BN)
    ccall((:bn_neg, LIB), Cvoid, (Ref{BN}, Ref{BN}), c, a)
    return c
end

function bn_rand!(a::BN, bits::Int)
    if bits < 1 || bits > typemax(Cint)
        error("invalid bits specified: $bits")
    end
    ccall((:bn_rand, LIB), Cvoid, (Ref{BN}, Cint, Cint), a, 0, bits)
    return a
end
function bn_read_bin!(a::BN, bin::Vector{UInt8})
    ccall((:bn_read_bin, LIB), Cvoid, (Ref{BN}, Ptr{UInt8}, Cint), a, bin, length(bin))
    return a
end

bn_size_bin(a::BN) = Int(ccall((:bn_size_bin, LIB), Cint, (Ref{BN},), a))

function bn_sub_dig!(c::BN, a::BN, b::Limb)
    ccall((:bn_sub_dig, LIB), Cvoid, (Ref{BN}, Ref{BN}, Limb), c, a, b)
    return c
end
function bn_sub!(c::BN, a::BN, b::BN)
    ccall((:bn_sub, LIB), Cvoid, (Ref{BN}, Ref{BN}, Ref{BN}), c, a, b)
    return c
end

function bn_write_bin!(bin::Vector{UInt8}, a::BN)
    ccall((:bn_write_bin, LIB), Cvoid, (Ptr{UInt8}, Cint, Ref{BN}), bin, length(bin), a)
    return bin
end

function ep_add_basic!(c::EP, a::EP, b::EP)
    ccall((:ep_add_basic, LIB), Cvoid, (Ref{EP}, Ref{EP}, Ref{EP}), c, a, b)
    return c
end

function ep_dbl_basic!(c::EP, a::EP)
    ccall((:ep_dbl_basic, LIB), Cvoid, (Ref{EP}, Ref{EP}), c, a)
    return c
end

function ep_curve_get_gen!(g::EP)
    ccall((:ep_curve_get_gen, LIB), Cvoid, (Ref{EP},), g)
    return g
end

function ep_curve_get_ord!(n::BN)
    ccall((:ep_curve_get_ord, LIB), Cvoid, (Ref{BN},), n)
    return n
end

function ep_map!(p::EP, msg::Vector{UInt8})
    ccall((:ep_map, LIB), Cvoid, (Ref{EP}, Ptr{UInt8}, Cint), p, msg, length(msg))
    return p
end

function ep_mul_lwnaf!(c::EP, a::EP, b::BN)
    ccall((:ep_mul_lwnaf, LIB), Cvoid, (Ref{EP}, Ref{EP}, Ref{BN}), c, a, b)
    return c
end

function ep_neg_basic!(c::EP, a::EP)
    ccall((:ep_neg_basic, LIB), Cvoid, (Ref{EP}, Ref{EP}), c, a)
    return c
end

function ep_read_bin!(a::EP, bin::Vector{UInt8})
    ccall((:ep_read_bin, LIB), Cvoid, (Ref{EP}, Ptr{UInt8}, Cint), a, bin, length(bin))
    return a
end

function ep_rand!(ep::EP)
    ccall((:ep_rand, LIB), Cvoid, (Ref{EP},), ep)
    return ep
end

ep_size_bin(ep::EP, pack::Bool) = Int(ccall((:ep_size_bin, LIB), Cint, (Ref{EP}, Cint), ep, pack ? one(Cint) : zero(Cint)))

function ep_sub_basic!(c::EP, a::EP, b::EP)
    ccall((:ep_sub_basic, LIB), Cvoid, (Ref{EP}, Ref{EP}, Ref{EP}), c, a, b)
    return c
end

function ep_write_bin!(bin::Vector{UInt8}, a::EP, pack::Bool)
    flags = pack ? one(Cint) : zero(Cint)
    ccall((:ep_write_bin, LIB), Cvoid, (Ptr{UInt8}, Cint, Ref{EP}, Cint), bin, length(bin), a, flags)
    return bin
end

function ep2_add_basic!(c::EP2, a::EP2, b::EP2)
    ccall((:ep2_add_basic, LIB), Cvoid, (Ref{EP2}, Ref{EP2}, Ref{EP2}), c, a, b)
    return c
end

function ep2_dbl_basic!(c::EP2, a::EP2)
    ccall((:ep2_dbl_basic, LIB), Cvoid, (Ref{EP2}, Ref{EP2}), c, a)
    return c
end

function ep2_curve_get_gen!(g::EP2)
    ccall((:ep2_curve_get_gen, LIB), Cvoid, (Ref{EP2},), g)
    return g
end

function ep2_curve_get_ord!(n::BN)
    ccall((:ep2_curve_get_ord, LIB), Cvoid, (Ref{BN},), n)
    return n
end

function ep2_map!(p::EP2, msg::Vector{UInt8})
    ccall((:ep2_map, LIB), Cvoid, (Ref{EP2}, Ptr{UInt8}, Cint), p, msg, length(msg))
    return p
end

function ep2_mul_lwnaf!(c::EP2, a::EP2, b::BN)
    ccall((:ep2_mul_lwnaf, LIB), Cvoid, (Ref{EP2}, Ref{EP2}, Ref{BN}), c, a, b)
    return c
end

function ep2_neg_basic!(c::EP2, a::EP2)
    ccall((:ep2_neg_basic, LIB), Cvoid, (Ref{EP2}, Ref{EP2}), c, a)
    return c
end

function ep2_rand!(ep::EP2)
    ccall((:ep2_rand, LIB), Cvoid, (Ref{EP2},), ep)
    return ep
end

function ep2_read_bin!(a::EP2, bin::Vector{UInt8})
    ccall((:ep2_read_bin, LIB), Cvoid, (Ref{EP2}, Ptr{UInt8}, Cint), a, bin, length(bin))
    return a
end

ep2_size_bin(ep::EP2, pack::Bool) = Int(ccall((:ep2_size_bin, LIB), Cint, (Ref{EP2}, Cint), ep, pack ? one(Cint) : zero(Cint)))

function ep2_sub_basic!(c::EP2, a::EP2, b::EP2)
    ccall((:ep2_sub_basic, LIB), Cvoid, (Ref{EP2}, Ref{EP2}, Ref{EP2}), c, a, b)
    return c
end

function ep2_write_bin!(bin::Vector{UInt8}, a::EP2, pack::Bool)
    flags = pack ? one(Cint) : zero(Cint)
    ccall((:ep2_write_bin, LIB), Cvoid, (Ptr{UInt8}, Cint, Ref{EP2}, Cint), bin, length(bin), a, flags)
    return bin
end

function fp_add_basic!(c::FP, a::FP, b::FP)
    ccall((:fp_add_basic, LIB), Cvoid, (Ref{FP}, Ref{FP}, Ref{FP}), c, a, b)
    return c
end

function fp_exp_slide!(c::FP, a::FP, b::BN)
    ccall((:fp_exp_slide, LIB), Cvoid, (Ref{FP}, Ref{FP}, Ref{BN}), c, a, b)
    return c
end

function fp_inv_lower!(c::FP, a::FP)
    ccall((:fp_inv_lower, LIB), Cvoid, (Ref{FP}, Ref{FP}), c, a)
    return c
end

function fp_mul_comba!(c::FP, a::FP, b::FP)
    ccall((:fp_mul_comba, LIB), Cvoid, (Ref{FP}, Ref{FP}, Ref{FP}), c, a, b)
    return c
end

function fp_neg_basic!(c::FP, a::FP)
    ccall((:fp_neg_basic, LIB), Cvoid, (Ref{FP}, Ref{FP}), c, a)
    return c
end

function fp_prime_get()
    fp = ccall((:fp_prime_get, LIB), Ptr{Limb}, ())
    bn = BN(used=FP_SIZE)
    unsafe_copyto!(unsafe_bnptr(bn), fp, bn.used)
    return bn
end

function fp_rand!(fp::FP)
    ccall((:fp_rand, LIB), Cvoid, (Ref{FP},), fp)
    return fp
end

function fp_sqr_comba!(c::FP, a::FP)
    ccall((:fp_sqr_comba, LIB), Cvoid, (Ref{FP}, Ref{FP}), c, a)
    return c
end

function fp_sub_basic!(c::FP, a::FP, b::FP)
    ccall((:fp_sub_basic, LIB), Cvoid, (Ref{FP}, Ref{FP}, Ref{FP}), c, a, b)
    return c
end

function fp_zero(c::FP)
    ccall((:fp_zero, LIB), Cvoid, (Ref{FP},), c)
    return c
end

function fp2_add_integ!(c::FP2, a::FP2, b::FP2)
    ccall((:fp2_add_integ, LIB), Cvoid, (Ref{FP2}, Ref{FP2}, Ref{FP2}), c, a, b)
    return c
end

function fp2_exp!(c::FP2, a::FP2, b::BN)
    ccall((:fp2_exp, LIB), Cvoid, (Ref{FP2}, Ref{FP2}, Ref{BN}), c, a, b)
    return c
end

function fp2_inv!(c::FP2, a::FP2)
    ccall((:fp2_inv, LIB), Cvoid, (Ref{FP2}, Ref{FP2}), c, a)
    return c
end

function fp2_mul_integ!(c::FP2, a::FP2, b::FP2)
    ccall((:fp2_mul_integ, LIB), Cvoid, (Ref{FP2}, Ref{FP2}, Ref{FP2}), c, a, b)
    return c
end

function fp2_neg!(c::FP2, a::FP2)
    ccall((:fp2_neg, LIB), Cvoid, (Ref{FP2}, Ref{FP2}), c, a)
    return c
end

function fp2_rand!(a::FP2)
    ccall((:fp2_rand, LIB), Cvoid, (Ref{FP2},), a)
    return a
end

function fp2_sub_integ!(c::FP2, a::FP2, b::FP2)
    ccall((:fp2_sub_integ, LIB), Cvoid, (Ref{FP2}, Ref{FP2}, Ref{FP2}), c, a, b)
    return c
end

function fp2_sqr_integ!(c::FP2, a::FP2)
    ccall((:fp2_sqr_integ, LIB), Cvoid, (Ref{FP2}, Ref{FP2}), c, a)
    return c
end

function fp2_zero(c::FP2)
    ccall((:fp2_zero, LIB), Cvoid, (Ref{FP2},), c)
    return c
end

function fp6_add!(c::FP6, a::FP6, b::FP6)
    ccall((:fp6_add, LIB), Cvoid, (Ref{FP6}, Ref{FP6}, Ref{FP6}), c, a, b)
    return c
end

function fp6_exp!(c::FP6, a::FP6, b::BN)
    ccall((:fp6_exp, LIB), Cvoid, (Ref{FP6}, Ref{FP6}, Ref{BN}), c, a, b)
    return c
end

function fp6_inv!(c::FP6, a::FP6)
    ccall((:fp6_inv, LIB), Cvoid, (Ref{FP6}, Ref{FP6}), c, a)
    return c
end

function fp6_mul_lazyr!(c::FP6, a::FP6, b::FP6)
    ccall((:fp6_mul_lazyr, LIB), Cvoid, (Ref{FP6}, Ref{FP6}, Ref{FP6}), c, a, b)
    return c
end

function fp6_neg!(c::FP6, a::FP6)
    ccall((:fp6_neg, LIB), Cvoid, (Ref{FP6}, Ref{FP6}), c, a)
    return c
end

function fp6_rand!(a::FP6)
    ccall((:fp6_rand, LIB), Cvoid, (Ref{FP6},), a)
    return a
end

function fp6_sub!(c::FP6, a::FP6, b::FP6)
    ccall((:fp6_sub, LIB), Cvoid, (Ref{FP6}, Ref{FP6}, Ref{FP6}), c, a, b)
    return c
end

function fp6_sqr_lazyr!(c::FP6, a::FP6)
    ccall((:fp6_sqr_lazyr, LIB), Cvoid, (Ref{FP6}, Ref{FP6}), c, a)
    return c
end

function fp6_zero(c::FP6)
    ccall((:fp6_zero, LIB), Cvoid, (Ref{FP6},), c)
    return c
end

function fp12_add!(c::FP12, a::FP12, b::FP12)
    ccall((:fp12_add, LIB), Cvoid, (Ref{FP12}, Ref{FP12}, Ref{FP12}), c, a, b)
    return c
end

function fp12_exp!(c::FP12, a::FP12, b::BN)
    ccall((:fp12_exp, LIB), Cvoid, (Ref{FP12}, Ref{FP12}, Ref{BN}), c, a, b)
    return c
end

function fp12_inv!(c::FP12, a::FP12)
    ccall((:fp12_inv, LIB), Cvoid, (Ref{FP12}, Ref{FP12}), c, a)
    return c
end

function fp12_mul_lazyr!(c::FP12, a::FP12, b::FP12)
    ccall((:fp12_mul_lazyr, LIB), Cvoid, (Ref{FP12}, Ref{FP12}, Ref{FP12}), c, a, b)
    return c
end

function fp12_neg!(c::FP12, a::FP12)
    ccall((:fp12_neg, LIB), Cvoid, (Ref{FP12}, Ref{FP12}), c, a)
    return c
end

function fp12_rand!(a::FP12)
    ccall((:fp12_rand, LIB), Cvoid, (Ref{FP12},), a)
    return a
end

function fp12_sub!(c::FP12, a::FP12, b::FP12)
    ccall((:fp12_sub, LIB), Cvoid, (Ref{FP12}, Ref{FP12}, Ref{FP12}), c, a, b)
    return c
end

function fp12_sqr_lazyr!(c::FP12, a::FP12)
    ccall((:fp12_sqr_lazyr, LIB), Cvoid, (Ref{FP12}, Ref{FP12}), c, a)
    return c
end

function fp12_zero(c::FP12)
    ccall((:fp12_zero, LIB), Cvoid, (Ref{FP12},), c)
    return c
end

function md_hmac(in::Vector{UInt8}, key::Vector{UInt8})
    mac = Vector{UInt8}(undef, MD_LEN)
    ccall((:md_hmac, LIB), Cvoid, (Ref{UInt8}, Ptr{UInt8}, Cint, Ptr{UInt8}, Cint),
        mac, in, length(in), key, length(key))
    return mac
end

function md_map_sh256(msg::Vector{UInt8})
    hash = Vector{UInt8}(undef, 32)
    ccall((:md_map_sh256, LIB), Cvoid, (Ref{UInt8}, Ptr{UInt8}, Cint), hash, msg, length(msg))
    return hash
end

function pp_exp_k12!(c::FP12, a::FP12)
    ccall((:pp_exp_k12, LIB), Cvoid, (Ref{FP12}, Ref{FP12}), c, a)
    return c
end

function pp_map_oatep_k12!(r::FP12, p::EP, q::EP2)
    ccall((:pp_map_oatep_k12, LIB), Cvoid, (Ref{FP12}, Ref{EP}, Ref{EP2}), r, p, q)
    return r
end


## Functions to compute the reduced shape

# for reductions that expand 0 dims to 1
reduced_dims(a::AbstractArray, region) = reduced_dims(size(a), region)

# for reductions that keep 0 dims as 0
reduced_dims0(a::AbstractArray, region) = reduced_dims0(size(a), region)

reduced_dims{N}(siz::NTuple{N,Int}, d::Int, rd::Int) = (d == 1 ? tuple(rd, siz[d+1:N]...) :
                                                        d == N ? tuple(siz[1:N-1]..., rd) :
                                                        1 < d < N ? tuple(siz[1:d-1]..., rd, siz[d+1:N]...) :
                                                        siz)::typeof(siz)

reduced_dims{N}(siz::NTuple{N,Int}, d::Int) = reduced_dims(siz, d, 1)

reduced_dims0{N}(siz::NTuple{N,Int}, d::Int) = 1 <= d <= N ? reduced_dims(siz, d, (siz[d] == 0 ? 0 : 1)) : siz

function reduced_dims{N}(siz::NTuple{N,Int}, region)
    rsiz = [siz...]
    for i in region
        if 1 <= i <= N
            rsiz[i] = 1
        end
    end
    tuple(rsiz...)::typeof(siz)
end

function reduced_dims0{N}(siz::NTuple{N,Int}, region)
    rsiz = [siz...]
    for i in region
        if i <= i <= N
            rsiz[i] = (rsiz[i] == 0 ? 0 : 1)
        end
    end
    tuple(rsiz...)::typeof(siz)
end

function regionsize(a, region)
    s = 1
    for d in region
        s *= size(a,d)
    end
    s
end

### Generic reduction functions

reducedim(f::Function, A, region, initial) = reducedim!(f, reduction_init(A, region, initial), A)

reducedim(f::Function, A, region, initial, R) = reducedim!(f, fill!(R, initial), A)

function reducedim!_function(N::Int, f::Function)
    body = gen_reduction_body(N, f)
    @eval begin
        local _F_
        function _F_(R, A)
            $body
        end
        _F_
    end
end

let reducedim_cache = Dict()
# reducedim! assumes that R has already been initialized with a seed value
global reducedim!
function reducedim!(f::Function, R, A)
    if isempty(R)
        return R
    end
    ndimsA = ndims(A)
    key = (ndimsA, f)
    if !haskey(reducedim_cache,key)
        func = reducedim!_function(ndimsA, f)
        reducedim_cache[key] = func
    else
        func = reducedim_cache[key]
    end
    func(R, A)::typeof(R)
end
end  # let reducedim_cache

# Generate the body for a reduction function reduce!(f, R, A), using binary operation f,
# where R is the output and A is the input.
# R must have already been set to the appropriate size and initialized with the seed value
function gen_reduction_body(N, f::Function)
    F = Expr(:quote, f)
    quote
        (isempty(R) || isempty(A)) && return R
        for i = 1:$N
            (size(R, i) == size(A, i) || size(R, i) == 1) || throw(DimensionMismatch("Reduction on array of size $(size(A)) with output of size $(size(R))"))
        end
        @nextract $N sizeR d->size(R,d)
        # If we're reducing along dimension 1, for efficiency we can make use of a temporary.
        # Otherwise, keep the result in R so that we traverse A in storage order.
        if size(R, 1) < size(A, 1)
            @nloops $N i d->(d>1? (1:size(A,d)) : (1:1)) d->(j_d = sizeR_d==1 ? 1 : i_d) begin
                @inbounds tmp = (@nref $N R j)
                for i_1 = 1:size(A,1)
                    @inbounds tmp = ($F)(tmp, (@nref $N A i))
                end
                @inbounds (@nref $N R j) = tmp
            end
        else
            @nloops $N i A d->(j_d = sizeR_d==1 ? 1 : i_d) begin
                @inbounds (@nref $N R j) = ($F)((@nref $N R j), (@nref $N A i))
            end
        end
        R
    end
end

reduction_init{T}(A::AbstractArray, region, initial::T) = fill!(similar(A,T,reduced_dims(A,region)), initial)

### Pre-generated cases
# For performance, these bypass reducedim_cache

function initarray!{T}(a::AbstractArray{T}, v::T, init::Bool)
    if init
        fill!(a, v)
    end
    return a
end

eval(ngenerate(:N, :(typeof(R)), :(_all!{N}(R::AbstractArray, A::AbstractArray{Bool,N})), N->gen_reduction_body(N, &)))
all!(r::AbstractArray, A::AbstractArray{Bool}; init::Bool=true) = _all!(initarray!(r, true, init), A)
all(A::AbstractArray{Bool}, region) = _all!(reduction_init(A, region, true), A)

eval(ngenerate(:N, :(typeof(R)), :(_any!{N}(R::AbstractArray, A::AbstractArray{Bool,N})), N->gen_reduction_body(N, |)))
any!(r::AbstractArray, A::AbstractArray{Bool}; init::Bool=true) = _any!(initarray!(r, false, init), A)
any(A::AbstractArray{Bool}, region) = any!(reduction_init(A, region, false), A)

eval(ngenerate(:N, :(typeof(R)), :(_maximum!{T,N}(R::AbstractArray, A::AbstractArray{T,N})), N->gen_reduction_body(N, scalarmax)))
maximum!{R}(r::AbstractArray{R}, A::AbstractArray; init::Bool=true) = _maximum!(initarray!(r, typemin(R), init), A)
maximum{T}(A::AbstractArray{T}, region) =
    isempty(A) ? similar(A,reduced_dims0(A,region)) : _maximum!(reduction_init(A, region, typemin(T)), A)

eval(ngenerate(:N, :(typeof(R)), :(_minimum!{T,N}(R::AbstractArray, A::AbstractArray{T,N})), N->gen_reduction_body(N, scalarmin)))
minimum!{R}(r::AbstractArray{R}, A::AbstractArray; init::Bool=true) = _minimum!(initarray!(r, typemax(R), init), A)
minimum{T}(A::AbstractArray{T}, region) =
    isempty(A) ? similar(A, reduced_dims0(A, region)) : _minimum!(reduction_init(A, region, typemax(T)), A)

eval(ngenerate(:N, :(typeof(R)), :(_sum!{T,N}(R::AbstractArray, A::AbstractArray{T,N})), N->gen_reduction_body(N, +)))
sum!{R}(r::AbstractArray{R}, A::AbstractArray; init::Bool=true) = _sum!(initarray!(r, zero(R), init), A)
sum{T}(A::AbstractArray{T}, region) = _sum!(reduction_init(A, region, zero(T)+zero(T)), A)

eval(ngenerate(:N, :(typeof(R)), :(_prod!{T,N}(R::AbstractArray, A::AbstractArray{T,N})), N->gen_reduction_body(N, *)))
prod!{R}(r::AbstractArray{R}, A::AbstractArray; init::Bool=true) = _prod!(initarray!(r, one(R), init), A)
prod{T}(A::AbstractArray{T}, region) = _prod!(reduction_init(A, region, one(T)*one(T)), A)

prod(A::AbstractArray{Bool}, region) = error("use all() instead of prod() for boolean arrays")


struct FVLAE{E,D,C,F,G,XD,ZD,V,XDIST<:Val} <: AbstractVLAE
    e::E
    d::D
    c::C # critic
    f::F # extracts latent variables
    g::G # concatenates latent vars with the rest
    xdim::XD # (h,w,c)
    zdim::ZD # scalar
    var::V # dense or conv last layer (ignore for dense fvlae)
    xdist::XDIST # gaussian or bernoulli
end
FVLAE(e,d,c,f,g,xdim,zdim,var,xdist::Symbol=:gaussian) =
	FVLAE(e,d,c,f,g,xdim,zdim,var,Val(xdist))
Flux.@functor FVLAE
(m::FVLAE)(x) = reconstruct(m, x)

function Base.show(io::IO, m::FVLAE)
    msg = """$(nameof(typeof(m))):
        encoder =
            e - $(repr(m.e))
            f - $(repr(m.f))
        decoder =
            d - $(repr(m.d))
            g - $(repr(m.g))
        critic =
            $(repr(m.c))
    """
    print(io, msg)
end

"""
	FVLAE(zdim::Int, hdim::Int, ks, ncs, stride, datasize; discriminator_nlayers::Int=3,
		layer_depth::Int=1, var=:dense, activation="relu", xdist=:gaussian)

    FVLAE(zdim::Int, hdims, discriminator_hdim::Int, datasize; discriminator_nlayers::Int=2,
    	activation=relu, layer_depth=1)

Factor VLAE constructor.
"""
function FVLAE(zdim::Int, hdim::Int, ks, ncs, strd, datasize; discriminator_nlayers::Int=3,
		var=:dense, xdist=:gaussian, kwargs...)
	(xdist in [:gaussian, :bernoulli]) ? nothing :
		error("xdist must be either :gaussian or :bernoulli")

	# get encoder, decoder and latent extractors and reshapers
	e,d,f,g = basic_model_constructor(zdim, ks, ncs, strd, datasize; var=var, xdist=xdist,
		kwargs...)

	# now construct the critic (discriminator)
	c = discriminator_constructor(zdim*length(ks), hdim, discriminator_nlayers; kwargs...)

    return FVLAE(e,d,c,f,g,datasize[1:3],zdim,var,xdist)
end
function FVLAE(zdim::Int, hdims, d_hdim::Int, datasize; discriminator_nlayers::Int=2, kwargs...)
    # get encoder, decoder and latent extractors and reshapers
    e,d,f,g = basic_model_constructor(zdim, hdims, datasize[1]; kwargs...)

	# now construct the critic (discriminator)
	c = HierarchicalAD.discriminator_constructor(zdim*length(hdims), d_hdim,
		discriminator_nlayers; kwargs...)

    return HierarchicalAD.FVLAE(e,d,c,f,g,(datasize[1],),zdim,nothing,:gaussian)
end

function logpdf(m::FVLAE{E,D,C,F,G,X,Z,V,XD}, x::AbstractArray{T,2}, zs...) where
		{E,D,C,F,G,X,Z,V,XD<:Val{:gaussian}} where T
    μx, σx = _decoded_mu_var(m, zs...)
    _x = (m.var == :dense) ? vectorize(x) : x
    return  gauss_logpdf(_x, μx, σx)
end
function logpdf(m::FVLAE{E,D,C,F,G,X,Z,V,XD}, x::AbstractArray{T,4}, zs...) where
	{E,D,C,F,G,X,Z,V,XD<:Val{:gaussian}} where T
    μx, σx = _decoded_mu_var1(m, zs...)
    _x = (m.var == :dense) ? vectorize(x) : x
    return  gauss_logpdf(_x, μx, σx)
end
function logpdf(m::FVLAE{E,D,C,F,G,X,Z,V,XD}, x::AbstractArray{T,4}, zs...) where
	{E,D,C,F,G,X,Z,V,XD<:Val{:bernoulli}} where T
    _x = _decoder_out(m, zs...)
    return  bernoulli_logpdf(_x, x)
end

dloss(d,z::AbstractArray{T,2},zp::AbstractArray{T,2}) where T  =
	-0.5f0*Flux.mean(log.(d(z) .+ eps(Float32)) .+ log.(1 .- d(zp) .+ eps(Float32)))
function total_correlation(d,z::AbstractArray{T,2}) where T
	dz = d(z)
	Flux.mean(log.(dz .+ eps(Float32)) .- log.(1 .- dz .+ eps(Float32)))
end

function factor_aeloss(m, x, γ::Float32)
	# encoder pass - KL divergence
	μzs_σzs = _encoded_mu_vars(m, x)
	zs = map(y->rptrick(y...), μzs_σzs)
	kldl = sum(map(y->Flux.mean(kld(y...)), μzs_σzs))

	# decoder pass - logpdf
	lpdf = Flux.mean(logpdf(m, x, zs...))
	elbo = -kldl + lpdf

	# now the discriminator loss
	_zs = cat(zs...,dims=1)
	tc = total_correlation(m.c, _zs)

	-elbo + γ*tc
end

reconstruction_loss(m,x)  =
	-Flux.mean(logpdf(m, x, encode_all(m, x)...))
kld_loss(m, x) =
    kldl = sum(map(y->Flux.mean(kld(y...)), _encoded_mu_vars(m, x)))
tc_loss(m, x) =
	total_correlation(m.c, cat(encode_all(m, x)...,dims=1))

function init_perm_mat(x::AbstractArray{T,2}) where T
    n = size(x,2)
    fill!(similar(x, n, n), 0), n
end
function init_perm_mat(x::SubArray)
    n = size(x,1)
    fill!(similar(x.parent, n, n), 0), n
end
function permutation_mat(x)
	perm_mat, n = init_perm_mat(x)
	for (j,i) in zip(1:n, sample(1:n, n, replace=false))
		perm_mat[j,i] = 1
	end
	perm_mat
end
# permutation for individual latents
permute_cols(x) where T = x*permutation_mat(x)
# permutation for inidivdual dims
permute_rows(x::AbstractArray{T,2}) where T =
    vcat(map(x->reshape(permutation_mat(x)*x,1,size(x,1)), eachrow(x))...)
Flux.Zygote.@nograd permutation_mat # or less?

function factor_closs_per_dim(m, x)
	zs = _encode_all(m, x; mean=false)
	zps = map(zs) do z
		zp = permute_rows(z)
		zp
	end
	_zs = cat(zs..., dims=1)
	_zps = cat(zps..., dims=1)
	dloss(m.c, _zs, _zps)
end
function factor_closs_per_latent(m, x)
	zs = _encode_all(m, x; mean=false)
	zps = map(zs) do z
		zp = permute_cols(z)
		zp
	end
	_zs = cat(zs..., dims=1)
	_zps = cat(zps..., dims=1)
	dloss(m.c, _zs, _zps)
end

#### THIS IS FOR IMAGE DATA ####
"""
	train_fvlae(zdim, hdim, batchsize, ks, ncs, stride, nepochs, tr_x, val_x;
	λ=0.0f0, γ=1.0f0, epochsize = size(tr_x,4), layer_depth=1, lr=0.001f0, var=:dense,
	activation="relu", discriminator_nlayers=3, xdist=:gaussian, pad=0,
	initial_convergence_threshold=0)

Train a factored VLAE.
"""
function train_fvlae(zdim, hdim, batchsize, ks, ncs, strd, nepochs, tr_x::AbstractArray{T,4},
	val_x::AbstractArray{T,4}; tr_y=nothing, val_y=nothing, factors=nothing,
	disentangle_per_latent=true,
	λ=0.0f0, γ=1.0f0, epochsize = size(tr_x,4),
	layer_depth=1, lr=0.001f0, var=:dense,
	activation="relu", discriminator_nlayers=3, xdist=:gaussian, pad=0,
	convergence_threshold=0f0, initial_convergence_epochs=10,
    max_retrain_tries=10, early_stopping=true, kwargs...) where T
    # this is to ensure that the model converges to something meaningful
	hist, rdata, model, zs, aeopt, copt = nothing, nothing, nothing, nothing, nothing, nothing

    ntries = 0

    while isnothing(hist) && ntries<=max_retrain_tries
        if ntries > 0
            @info "Restarting training automatically..."
        end

        # construct the model
        model = gpu(FVLAE(zdim, hdim, ks, ncs, strd, size(tr_x), layer_depth=layer_depth, var=var,
            activation=activation, nlayers=discriminator_nlayers, xdist=xdist, pad=pad))
        aeopt = ADAM(lr)
        copt = ADAM(lr)

        # train it
        hist, rdata, zs = train!(model, nepochs, batchsize, tr_x, val_x, aeopt, copt;
        	tr_y=tr_y, val_y=val_y, factors=factors,
        	disentangle_per_latent=disentangle_per_latent,
            λ=λ, γ=γ, epochsize = epochsize,
            convergence_threshold=convergence_threshold,
            initial_convergence_epochs=initial_convergence_epochs,
            early_stopping=early_stopping)
        ntries += 1
    end

    # if the training was still unsuccesfull
    if isnothing(hist)
        @info "Gave up the training after $ntries tries."
        return model, hist, rdata, zs, aeopt, copt
    end

    return model, hist, rdata, zs, aeopt, copt
end

"""
	train!(model::FVLAE, nepochs, batchsize, tr_x::AbstractArray{T,4},
		val_x::AbstractArray{T,4}, aeopt, copt;
		tr_y=nothing, val_y=nothing, factors=nothing, disentangle_per_latent=true,
		λ=0.0f0, γ=1.0f0, epochsize = size(tr_x,4), convergence_threshold=0f0,
    	initial_convergence_epochs=2)
"""
function train!(model::FVLAE, nepochs, batchsize, tr_x::AbstractArray{T,4},
	val_x::AbstractArray{T,4}, aeopt, copt;
	tr_y=nothing, val_y=nothing, factors=nothing, disentangle_per_latent=true,
	disentanglement_repeats::Int=1,
	λ=0.0f0, γ=1.0f0, epochsize = size(tr_x,4), convergence_threshold=0f0,
    initial_convergence_epochs=10, early_stopping=true,
    kwargs...) where T
	if !isnothing(tr_y) && !(isnothing(val_y)) && isnothing(factors)
		error("specify factors you want to disentangle")
	end

	# data an initial reconstruction loss
	gval_x = gpu(val_x[:,:,:,1:min(1000, size(val_x,4))]);
	local data_itr
	@suppress begin # suppress the batchsize warning
		data_itr = Flux.Data.DataLoader(sample_tensor(tr_x, epochsize), batchsize=batchsize)
	end
    init_rloss = batched_loss(x->reconstruction_loss(model,x), gval_x, batchsize)
    control_rloss = init_rloss

    # params
    # this is according to their code right
    aeps = Flux.params(model.e, model.d, model.f, model.g)
    cps = Flux.params(model.c)

    # losses
    aeloss(x) = factor_aeloss(model, gpu(x), γ) + λ*sum(l2, aeps)
    closs(x) = disentangle_per_latent ?
    	factor_closs_per_latent(model, gpu(x)) : factor_closs_per_dim(model, gpu(x))

    # rest of the setup
    rdata = []
    hist = MVHistory()
    nl = length(model.e)
    zs = [[] for _ in 1:nl]

	println("Training in progress...")
	for epoch in 1:nepochs
		for x in data_itr
			# first update the autoencoder
			param_update!(aeloss, aeps, x, aeopt)

			# then update the critic/discriminator
			param_update!(closs, cps, x, copt)
		end

		 # sometimes the model is stuck in some local minima and cant get out - better restart the training
        rloss = batched_loss(x->reconstruction_loss(model, x), gval_x, batchsize)
        if epoch == initial_convergence_epochs
            if abs((init_rloss - rloss)/init_rloss) < convergence_threshold
                @info "\nInitial improvement after $(initial_convergence_epochs) epochs is $(init_rloss) => $rloss, model probably stuck in local minima, terminating training. Try restarting the model."
                return nothing, nothing, nothing
            end
        end

		# logging
		Flux.Zygote.ignore() do
			# compute loss values
			ael=round(batched_loss(aeloss, gval_x, batchsize), digits=2)
			cl=round(batched_loss(closs, gval_x, batchsize), digits=4)
			rl=round(batched_loss(x->reconstruction_loss(model, x), gval_x, batchsize), digits=2)
			kldl=round(batched_loss(x->kld_loss(model, x), gval_x, batchsize), digits=2)
			tcl=round(batched_loss(x->tc_loss(model, x), gval_x, batchsize), digits=4)
			println("Epoch $(epoch)/$(nepochs), validation loss: AE=$ael | C=$cl | R=$rl | KLD=$kldl | TC=$tcl")

			push!(hist, :autoencoder_loss, epoch, ael)
			push!(hist, :critic_loss, epoch, cl)
			push!(hist, :reconstruction_loss, epoch, rl)
			push!(hist, :kld, epoch, kldl)
			push!(hist, :total_correlation, epoch, tcl)
			push!(rdata, cpu(reconstruct(model, gval_x)))

			# also, compute the disentanglement metric
			dmls = []
			dmds = []
			if !isnothing(factors)
				zs_all = cpu(encode_all(model, cat(tr_x, val_x, dims=4)))
				y = (isnothing(tr_y) || isnothing(val_y)) ? nothing : vcat(tr_y, val_y)
				# in case factors is just one vector
				if typeof(factors[1]) == Symbol || length(factors[1]) == 1
					factors = [factors]
				end
				for factor_vec in factors
					dml = disentanglement_per_latent(zs_all, y, factor_vec, disentanglement_repeats)
					dmd = disentanglement_per_dim(vcat(zs_all...), y, factor_vec, disentanglement_repeats)
					println("Disentanglement of $(factor_vec): $dml (per latent) | $dmd (per dim)")
					push!(dmls, dml)
					push!(dmds, dmd)
				end
				println("")
			end
			push!(hist, :disentanglement_per_latent, epoch, dmls)
			push!(hist, :disentanglement_per_dim, epoch, dmds)

			#
			for i in 1:nl
				z = encode(model, gval_x, i)
				push!(zs[i], cpu(z))
			end
		end

		# early stopping
		if early_stopping && control_rloss < rloss
			@info "Stoppping early after $epoch epochs due to no improvement."
			return hist, rdata, zs
		else
			control_rloss = rloss
		end
	end

	return hist, rdata, zs
end

#### THIS IS FOR TAB DATA ####
"""
	train_fvlae(zdim, hdim, batchsize, discriminator_hdim, nepochs, tr_x, val_x;
	λ=0.0f0, γ=1.0f0, epochsize = size(tr_x,2), layer_depth=1, lr=0.001f0,
	activation="relu", discriminator_nlayers=3,
	initial_convergence_threshold=0, initial_convergence_epochs=10,
    max_retrain_tries=10)

Train a factored dense VLAE.
"""
function train_fvlae(zdim, hdims, batchsize, d_hdim, nepochs, tr_x::AbstractArray{T,2},
	val_x::AbstractArray{T,2}; λ=0.0f0, γ=1.0f0, epochsize = size(tr_x,2),
	layer_depth=1, lr=0.001f0, activation="relu", discriminator_nlayers=3,
	initial_convergence_threshold=0f0, initial_convergence_epochs=nepochs+1,
    max_retrain_tries=10, verb=true, kwargs...) where T

    # this is to ensure that the model converges to something meaningful
	hist, rdata, model, zs, aeopt, copt = nothing, nothing, nothing, nothing, nothing, nothing

    ntries = 0

    while isnothing(hist) && ntries<=max_retrain_tries
        if ntries > 0
            @info "Restarting training automatically..."
        end

        # construct the model
        model = FVLAE(zdim, hdims, d_hdim, size(tr_x); activation=activation,
        	discriminator_nlayers=discriminator_nlayers,
        	layer_depth=layer_depth)
        aeopt = ADAM(lr)
        copt = ADAM(lr)

        # train it
        hist, rdata, zs = train!(model, nepochs, batchsize, tr_x, val_x, aeopt, copt;
            λ=λ, γ=γ, epochsize = epochsize,
            initial_convergence_threshold=initial_convergence_threshold,
            initial_convergence_epochs=initial_convergence_epochs,
            verb=verb)
        ntries += 1
    end

    # if the training was still unsuccesfull
    if isnothing(hist)
        @info "Gave up the training after $ntries tries."
        return model, hist, rdata, zs, aeopt, copt
    end

    return model, hist, rdata, zs, aeopt, copt
end

"""
	train!(model::FVLAE, nepochs, batchsize, tr_x::AbstractArray{T,2},
		val_x::AbstractArray{T,2}, aeopt, copt;
		λ=0.0f0, γ=1.0f0, epochsize = size(tr_x,2), initial_convergence_threshold=0f0,
    	initial_convergence_epochs=50)
"""
function train!(model::FVLAE, nepochs, batchsize, tr_x::AbstractArray{T,2},
	val_x::AbstractArray{T,2}, aeopt, copt;
	λ=0.0f0, γ=1.0f0, epochsize = size(tr_x,2), initial_convergence_threshold=0f0,
    initial_convergence_epochs=50, verb=true) where T

	# data an initial reconstruction loss
	subtr_x = tr_x[:,1:min(1000, size(tr_x,2))]
	subval_x = val_x[:,1:min(1000, size(val_x,2))]
	local data_itr
	@suppress begin # suppress the batchsize warning
		data_itr = Flux.Data.DataLoader(sample_tensor(tr_x, epochsize), batchsize=batchsize)
    end
    init_rloss = batched_loss(x->reconstruction_loss(model,x), subval_x, batchsize)

    # params
    aeps = Flux.params(model.e, model.d, model.f, model.g)
    cps = Flux.params(model.c)

    # losses
    aeloss(x) = factor_aeloss(model, x, γ) + λ*sum(l2, aeps)
    closs(x) = factor_closs(model, x)

    # rest of the setup
    rdata = []
    hist = MVHistory()
    nl = length(model.e)
    zs = [[] for _ in 1:nl]

    if verb
		println("Training in progress...")
	end
	for epoch in 1:nepochs
		for x in data_itr
			# first update the autoencoder
			param_update!(aeloss, aeps, x, aeopt)

			# then update the critic/discriminator
			param_update!(closs, cps, x, copt)
		end

		 # sometimes the model is stuck in some local minima and cant get out - better restart the training
        if epoch == initial_convergence_epochs
            rloss = batched_loss(x->reconstruction_loss(model, x), subval_x, batchsize)
            if abs((init_rloss - rloss)/init_rloss) < initial_convergence_threshold
                @info "Initial improvement after $(initial_convergence_epochs) epochs is $(init_rloss) => $rloss, model probably stuck in local minima, terminating training. Try restarting the model."
                return nothing, nothing, nothing
            end
        end

		# logging
		Flux.Zygote.ignore() do
			ael=round(batched_loss(aeloss, subval_x, batchsize), digits=4)
			cl=round(batched_loss(closs, subval_x, batchsize), digits=4)
			rl=round(batched_loss(x->reconstruction_loss(model, x), subval_x, batchsize), digits=4)
			kldl=round(batched_loss(x->kld_loss(model, x), subval_x, batchsize), digits=2)
			tcl=round(batched_loss(x->tc_loss(model, x), subval_x, batchsize), digits=4)

			if verb
				println("Epoch $(epoch)/$(nepochs), validation loss: AE=$ael | C=$cl | R=$rl | KLD=$kldl | TC=$tcl")
			end
			for i in 1:nl
				z = encode(model, subval_x, i)
				push!(zs[i], z)
			end
			push!(hist, :autoencoder_loss, epoch, ael)
			push!(hist, :critic_loss, epoch, cl)
			push!(hist, :reconstruction_loss, epoch, rl)
			push!(hist, :kld, epoch, kldl)
			push!(hist, :total_correlation, epoch, tcl)
			push!(rdata, reconstruct(model, subval_x))
		end
	end

	return hist, rdata, zs
end


#=
This file is auto-generated. Do not edit.
=#
"""
    mutable struct OneDOneQMachine <: Machine
        R::Float64
        Xd::Float64
        Xq::Float64
        Xd_p::Float64
        Xq_p::Float64
        Td0_p::Float64
        Tq0_p::Float64
        MVABase::Float64
        ext::Dict{String, Any}
        states::Vector{Symbol}
        n_states::Int64
        internal::InfrastructureSystemsInternal
    end

Parameters of 4-states synchronous machine: Simplified Marconato model
 The derivative of stator fluxes (ψd and ψq) is neglected and ωψd = ψd and
 ωψq = ψq is assumed (i.e. ω=1.0). This is standard when
 transmission network dynamics is neglected.

# Arguments
- `R::Float64`: Resistance after EMF in machine per unit, validation range: (0, nothing)
- `Xd::Float64`: Reactance after EMF in d-axis per unit, validation range: (0, nothing)
- `Xq::Float64`: Reactance after EMF in q-axis per unit, validation range: (0, nothing)
- `Xd_p::Float64`: Transient reactance after EMF in d-axis per unit, validation range: (0, nothing)
- `Xq_p::Float64`: Transient reactance after EMF in q-axis per unit, validation range: (0, nothing)
- `Td0_p::Float64`: Time constant of transient d-axis voltage, validation range: (0, nothing)
- `Tq0_p::Float64`: Time constant of transient q-axis voltage, validation range: (0, nothing)
- `MVABase::Float64`: Nominal Capacity in MVA, validation range: (0, nothing)
- `ext::Dict{String, Any}`
- `states::Vector{Symbol}`
- `n_states::Int64`
- `internal::InfrastructureSystemsInternal`: power system internal reference, do not modify
"""
mutable struct OneDOneQMachine <: Machine
    "Resistance after EMF in machine per unit"
    R::Float64
    "Reactance after EMF in d-axis per unit"
    Xd::Float64
    "Reactance after EMF in q-axis per unit"
    Xq::Float64
    "Transient reactance after EMF in d-axis per unit"
    Xd_p::Float64
    "Transient reactance after EMF in q-axis per unit"
    Xq_p::Float64
    "Time constant of transient d-axis voltage"
    Td0_p::Float64
    "Time constant of transient q-axis voltage"
    Tq0_p::Float64
    "Nominal Capacity in MVA"
    MVABase::Float64
    ext::Dict{String, Any}
    states::Vector{Symbol}
    n_states::Int64
    "power system internal reference, do not modify"
    internal::InfrastructureSystemsInternal
end

function OneDOneQMachine(R, Xd, Xq, Xd_p, Xq_p, Td0_p, Tq0_p, MVABase, ext=Dict{String, Any}(), )
    OneDOneQMachine(R, Xd, Xq, Xd_p, Xq_p, Td0_p, Tq0_p, MVABase, ext, [:eq_p, :ed_p], 2, InfrastructureSystemsInternal(), )
end

function OneDOneQMachine(; R, Xd, Xq, Xd_p, Xq_p, Td0_p, Tq0_p, MVABase, ext=Dict{String, Any}(), )
    OneDOneQMachine(R, Xd, Xq, Xd_p, Xq_p, Td0_p, Tq0_p, MVABase, ext, )
end

# Constructor for demo purposes; non-functional.
function OneDOneQMachine(::Nothing)
    OneDOneQMachine(;
        R=0,
        Xd=0,
        Xq=0,
        Xd_p=0,
        Xq_p=0,
        Td0_p=0,
        Tq0_p=0,
        MVABase=0,
        ext=Dict{String, Any}(),
    )
end

"""Get OneDOneQMachine R."""
get_R(value::OneDOneQMachine) = value.R
"""Get OneDOneQMachine Xd."""
get_Xd(value::OneDOneQMachine) = value.Xd
"""Get OneDOneQMachine Xq."""
get_Xq(value::OneDOneQMachine) = value.Xq
"""Get OneDOneQMachine Xd_p."""
get_Xd_p(value::OneDOneQMachine) = value.Xd_p
"""Get OneDOneQMachine Xq_p."""
get_Xq_p(value::OneDOneQMachine) = value.Xq_p
"""Get OneDOneQMachine Td0_p."""
get_Td0_p(value::OneDOneQMachine) = value.Td0_p
"""Get OneDOneQMachine Tq0_p."""
get_Tq0_p(value::OneDOneQMachine) = value.Tq0_p
"""Get OneDOneQMachine MVABase."""
get_MVABase(value::OneDOneQMachine) = value.MVABase
"""Get OneDOneQMachine ext."""
get_ext(value::OneDOneQMachine) = value.ext
"""Get OneDOneQMachine states."""
get_states(value::OneDOneQMachine) = value.states
"""Get OneDOneQMachine n_states."""
get_n_states(value::OneDOneQMachine) = value.n_states
"""Get OneDOneQMachine internal."""
get_internal(value::OneDOneQMachine) = value.internal


"""
`module  JAC.GreenFunction`
    ... a submodel of JAC that contains all methods for computing approximate many-electron Green functions.
"""
module GreenFunction

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

    """
    `abstract type GreenFunction.AbstractGreenApproach`
        ... defines an abstract and a number of singleton types for approximating a many-electron Green
            function for calculating second-order processes.

      + struct SingleCSFwithoutCI
        ... to approximate the many-electron multiplets (gMultiplet) for every chosen level symmetry by dealing with each CSF
            independently, and without any configuration interaction. This is a fast but also very rough approximation.

      + struct CoreSpaceCI
        ... to approximate the many-electron multiplets (gMultiplet) by taking the electron-electron interaction between the
            bound-state orbitals into account.

      + struct DampedSpaceCI
        ... to approximate the many-electron multiplets (gMultiplet) by taking the electron-electron interaction for all, the
            bound and free-electron, orbitals into account but by including a damping factor e^{tau*r}, tau > 0 into the
            electron densities rho_ab (r) --> rho_ab (r) * e^{tau*r}
    """
    abstract type  AbstractGreenApproach                           end
    struct         SingleCSFwithoutCI  <:  AbstractGreenApproach   end
    struct         CoreSpaceCI         <:  AbstractGreenApproach   end
    struct         DampedSpaceCI       <:  AbstractGreenApproach   end

    """
    `struct  GreenFunction.Settings`
        ... defines a type for defining the details and parameters of the approximated Green functions.

        + approach                 ::AbstractGreenApproach            ... Approach used to approximate the representation.
        + excitationScheme         ::Basics.AbstractExcitationScheme  ... Applied excitation scheme with regard to the
                                                                          given bound configuration.
        + nMax                     ::Int64            ... maximum principal quantum numbers of (single-electron)
                                                          excitations that are to be included into the representation.
        + lValues                  ::Array{Int64,1}   ... List of (non-relativistic) orbital angular momenta for which
                                                          (single-electron) excitations are to be included.
        + levelSymmetries          ::Array{LevelSymmetry,1}           ... Total level symmetries J^P to be included into
                                                                          the  Green function representation.
        + printBeforeComputation   ::Bool             ... True if a short overview is to be printed before.
        + selectLevels             ::Bool             ... True if individual levels are selected for the computation.
        + selectedLevels           ::Array{Int64,1}   ... List of selected levels.
    """
    struct Settings
        approach                   ::AbstractGreenApproach
        excitationScheme           ::Basics.AbstractExcitationScheme
        nMax                       ::Int64
        lValues                    ::Array{Int64,1}
        levelSymmetries            ::Array{LevelSymmetry,1}
        printBeforeComputation     ::Bool
        selectLevels               ::Bool
        selectedLevels             ::Array{Int64,1}
    end

    """
    `GreenFunction.Settings()`  ... constructor for an `empty` instance of GreenFunction.
    """
    function Settings()
        Settings( GreenFunction.SingleCSFwithoutCI(), Basics.NoExcitationScheme(), 0, Int64[], LevelSymmetry[], false, false, Int64[])
    end

    # `Base.show(io::IO, settings::GreenFunction.Settings)`  ... prepares a proper printout of the variable settings::GreenFunction.Settings.
    function Base.show(io::IO, settings::GreenFunction.Settings)
        println(io, "approach:                 $(settings.approach)  ")
        println(io, "excitationScheme:         $(settings.excitationScheme)  ")
        println(io, "nMax:                     $(settings.nMax)  ")
        println(io, "lValues:                  $(settings.lValues)  ")
        println(io, "levelSymmetries:          $(settings.levelSymmetries)  ")
        println(io, "printBeforeComputation:   $(settings.printBeforeComputation)  ")
        println(io, "selectLevels:             $(settings.selectLevels)  ")
        println(io, "selectedLevels:           $(settings.selectedLevels)  ")
    end

    """
    `struct  GreenFunction.Channel`
        ... defines a type for a single symmetry channel of an (approximate) Green function.

        + symmetry          ::LevelSymmetry    ... Level symmetry of this part of the representation.
        + gMultiplet        ::Multiplet        ... Multiplet of (scattering) levels of this symmetry.
    """
    struct Channel
        symmetry            ::LevelSymmetry
        gMultiplet          ::Multiplet
    end

    """
    `GreenFunction.Channel()`  ... constructor for an `empty` instance of GreenFunction.Channel.
    """
    function Channel()
        Channel( LevelSymmetry(0, Basics.plus), ManyElectron.Multiplet)
    end

    # `Base.show(io::IO, channel::GreenFunction.Channel)`  ... prepares a proper printout of the variable channel::GreenFunction.Channel.
    function Base.show(io::IO, channel::GreenFunction.Channel)
        println(io, "symmetry:                $(channel.symmetry)  ")
        println(io, "gMultiplet:              $(channel.gMultiplet)  ")
    end

    """
    `struct  GreenFunction.Representation`
        ... defines a type to keep a representation of an (approximate) Green function that is associated
            with a given set of bound configurations.

        + approach          ::AbstractGreenApproach           ... Approach used to approximate the representation.
        + excitationScheme  ::Basics.AbstractExcitationScheme ... Applied excitation scheme with regard to the
                                                                  given bound configuration.
        + boundConfigs      ::Array{Configuration,1}          ... List of bound configurations to which the Green
                                                                  function refers to.
        + NoElectrons       ::Int64                           ... Number of electrons.
        + channels          ::Array{GreenFunction.Channel,1}  ... List of channels with different level symmetry.
    """
    struct Representation
        approach            ::AbstractGreenApproach
        excitationScheme    ::Basics.AbstractExcitationScheme
        boundConfigs        ::Array{Configuration,1}
        NoElectrons         ::Int64
        channels            ::Array{GreenFunction.Channel,1}
    end

    """
    `GreenFunction.Representation()`  ... constructor for an `empty` instance of GreenFunction.Representation.
    """
    function Representation()
        Representation( GreenFunction.SingleCSFwithoutCI(), Basics.NoExcitationScheme(), Configuration[], 0, GreenFunction.Channel[])
    end

    # `Base.show(io::IO, representation::GreenFunction.Representation)`  ... prepares a proper printout of the
    #       variable representation::GreenFunction.Representation.
    function Base.show(io::IO, representation::GreenFunction.Representation)
        println(io, "approach:                $(representation.approach)  ")
        println(io, "excitationScheme:        $(representation.excitationScheme)  ")
        println(io, "boundConfigs:            $(representation.boundConfigs)  ")
        println(io, "NoElectrons:             $(representation.NoElectrons)  ")
        println(io, "channels:                $(representation.channels)  ")
    end

    """
    `GreenFunction.computeRepresentation(configs::Array{Configuration,1}, multiplet::Multiplet, nm::Nuclear.Model,
                                         grid::Radial.Grid, settings::GreenFunction.Settings; output=true)`
        ... to compute an approximate Green function representation for the levels of the given bound configurations and for the
            given excitation scheme.
    """
    function computeRepresentation(configs::Array{Configuration,1}, multiplet::Multiplet, nm::Nuclear.Model,
                                   grid::Radial.Grid, settings::GreenFunction.Settings; output=true)
        println("")
        printstyled("GreenFunction.computeRepresentation(): The computation of approximate Green functions starts now ... \n", color=:light_green)
        printstyled("---------------------------------------------------------------------------------------------------- \n", color=:light_green)

        gChannels = GreenFunction.Channel[]

        # Generate all (non-relativistic) configurations from the bound configurations due to the given excitation scheme
        confList = Basics.generateConfigurationsForExcitationScheme(configs, settings.excitationScheme, settings.nMax, settings.lValues)
        # Print (if required) information about the generated configuration list
        if  settings.printBeforeComputation    GreenFunction.displayConfigurationList(confList)    end
        # Determine all (selected) levels and generate an averaged potential for these levels; use a ... potential for the selected levels
        pot = GreenFunction.generateMeanPotential(multiplet, nm, grid, settings)
        # Generate a full single-electron spectrum for this potential
        spectrum = Bspline.generateSpectrum()
        # Cycle over all selected level symmetries to generate in turn all channels of the requires Green function representation
        for  levelSymmetry  in  settings.levelSymmetries
            # Generate all CSFs for this symmetry
            csfs  = Basics.generateCsfList(levelSymmetry, confList)
            # Compute a multiplet representation due to the given approach
            gMultiplet = GreenFunction.computeMultipletForApproach(settings.approach, csfs, spectrum)
            channel    = GreenFunction.Channel( levelSymmetry, gMultiplet )
        end

        gfRep = GreenFunction.Representation( settings.approach, settings.excitationScheme, configs, confList[1].NoElectrons , gChannels)
        #
        # Print all results to screen
        GreenFunction.displayRepresentation(stdout, gfRep)
        printSummary, iostream = Defaults.getDefaults("summary flag/stream")
        if  printSummary    GreenFunction.displayRepresentation(iostream, gfRep)   end
        #
        if    output    return( gfRep )
        else            return( nothing )
        end
    end

    """
    `GreenFunction.displayConfigurationList(configs::Array{Configuration,1})`
        ... displays the gnerated list of configurations in a compact form; nothing is returned in this case.
    """
    function  displayConfigurationList(configs::Array{Configuration,1})
        println(" ")
        println("  Generated list of (non-relativistic) configurations that contribute to the Green functions representation:")
        println(" ")
        println("  ", TableStrings.hLine(60))
        for  i = 1:length(configs)
            sa = string("(", i, ")" );    sa = TableStrings.flushright(9, sa, na=3);   sa = sa * string(configs[i])
            println(sa)
        end
        println("  ", TableStrings.hLine(60))

        return( nothing )
    end

    """
    `GreenFunction.displayRepresentation(representation::GreenFunction.Representation)`
        ... to display
    """
    function  displayRepresentation(representation::GreenFunction.Representation)
        println(" ")
        println("  Selected GreenFunction levels:  ... not yet implemented")
        println(" ")
        println("  ", TableStrings.hLine(43))
        sa = "  ";   sb = "  "
        sa = sa * TableStrings.center(10, "Level"; na=2);                             sb = sb * TableStrings.hBlank(12)
        sa = sa * TableStrings.center(10, "J^P";   na=4);                             sb = sb * TableStrings.hBlank(14)
        sa = sa * TableStrings.center(14, "Energy"; na=4);
        sb = sb * TableStrings.center(14, TableStrings.inUnits("energy"); na=4)
        println(sa);    println(sb);    println("  ", TableStrings.hLine(43))
        #
        for  outcome in outcomes
            sa  = "  ";    sym = LevelSymmetry( outcome.level.J, outcome.level.parity)
            sa = sa * TableStrings.center(10, TableStrings.level(outcome.level.index); na=2)
            sa = sa * TableStrings.center(10, string(sym); na=4)
            sa = sa * @sprintf("%.8e", Defaults.convertUnits("energy: from atomic", outcome.level.energy)) * "    "
            println( sa )
        end
        println("  ", TableStrings.hLine(43))
        #
        return( nothing )
    end

    """
    `GreenFunction.generateMeanPotential(multiplet::Multiplet, nm::Nuclear.Model, grid::Radial.Grid, settings::GreenFunction.Settings)`
        ... generates a mean radial potential for the selected levels of the multiplet that is used to compute the
            single-electron spectrum for the representation of the Green function; a potential::Radial.Potential
            is returned.
    """
    function  generateMeanPotential(multiplet::Multiplet, nm::Nuclear.Model, grid::Radial.Grid, settings::GreenFunction.Settings)
        nuclearPotential  = Nuclear.nuclearPotential(nm, grid)
        ## wp1 = compute("radial potential: core-Hartree", grid, wLevel)
        ## wp2 = compute("radial potential: Hartree-Slater", grid, wLevel)
        ## wp3 = compute("radial potential: Kohn-Sham", grid, wLevel)
        ## wp4 = compute("radial potential: Dirac-Fock-Slater", grid, wLevel)

        if  settings.selectLevels
            wpx = zeros[grid.NoPoints];   nx = 0
            for  i = 1:length(multiplet.levels)
                if haskey(settings.selectedLevels, i)
                    wp = compute("radial potential: Dirac-Fock-Slater", grid, multiplet.level[i]);   np = length(wp.Zr)
                    wpx[1:np] = wpx[1:np] + wp.Zr[1:np];    nx = nx + 1
                end
            end
            wpx = wpx / nx
            wp = Radial.Potential(wp.name, wpx, grid)
        else
            wp = compute("radial potential: Dirac-Fock-Slater", grid, multiplet.levels[1].basis)
        end

        pot = Basics.add(nuclearPotential, wp)
        println("Mean potential for the generation of the (quasi-complete) single-electron spectrum:  $(pot)")

        return( pot )
    end

end # module