export Experiment

using .ArcadeLearningEnvironment
using Dates
using ReinforcementLearningCore
using .ReinforcementLearningEnvironments
using Flux
using BSON
using TensorBoardLogger
using Logging
using Statistics

clip_reward(obs) = RewardOverriddenObs(obs, clamp(get_reward(obs), -1, 1))

function atari_env_factory(name, state_size, n_frames, max_episode_steps=100_000; seed = nothing)
    WrappedEnv(
        env = AtariEnv(;
            name = string(name),
            grayscale_obs = true,
            noop_max = 30,
            frame_skip = 4,
            terminal_on_life_loss = false,
            repeat_action_probability = 0.25,
            max_num_frames_per_episode = n_frames * max_episode_steps,
            color_averaging = false,
            full_action_space = false,
            seed = seed,
        ),
        preprocessor = ComposedPreprocessor(
            ResizeImage(state_size...),  # this implementation is different from cv2.resize https://github.com/google/dopamine/blob/e7d780d7c80954b7c396d984325002d60557f7d1/dopamine/discrete_domains/atari_lib.py#L629
            StackFrames(state_size..., n_frames),
            clip_reward,
        ),
    )
end

function RLCore.Experiment(
    ::Val{:Dopamine},
    ::Val{:DQN},
    ::Val{:Atari},
    name::AbstractString;
    save_dir = nothing,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyymmddHHMMSS")
        save_dir = joinpath(pwd(), "checkpoints", "dopamine_DQN_atari_$(name)_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    N_FRAMES = 4
    STATE_SIZE = (84, 84)
    env = atari_env_factory(name, STATE_SIZE, N_FRAMES;)#= seed=nothing =#
    N_ACTIONS = length(get_action_space(env))
    init = seed_glorot_uniform()#= seed=341 =#

    create_model() =
        Chain(
            x -> x ./ 255,
            CrossCor((8, 8), N_FRAMES => 32, relu; stride = 4, pad = 2, init = init),
            CrossCor((4, 4), 32 => 64, relu; stride = 2, pad = 2, init = init),
            CrossCor((3, 3), 64 => 64, relu; stride = 1, pad = 1, init = init),
            x -> reshape(x, :, size(x)[end]),
            Dense(11 * 11 * 64, 512, relu),
            Dense(512, N_ACTIONS),
        ) |> gpu

    agent = Agent(
        policy = QBasedPolicy(
            learner = DQNLearner(
                approximator = NeuralNetworkApproximator(
                    model = create_model(),
                    optimizer = RMSProp(0.00025, 0.95),
                ),  # unlike TF/PyTorch RMSProp doesn't support center
                target_approximator = NeuralNetworkApproximator(model = create_model()),
                update_freq = 4,
                γ = 0.99f0,
                update_horizon = 1,
                batch_size = 32,
                stack_size = N_FRAMES,
                min_replay_history = 20_000,
                loss_func = huber_loss,
                target_update_freq = 8_000,
            ),
            explorer = EpsilonGreedyExplorer(
                ϵ_init = 1.0,
                ϵ_stable = 0.01,
                decay_steps = 250_000,
                kind = :linear,
            ),
        ),
        trajectory = CircularCompactSARTSATrajectory(
            capacity = 1_000_000,
            state_type = Float32,
            state_size = STATE_SIZE,
        ),
    )

    evaluation_result = []
    EVALUATION_FREQ = 250_000
    MAX_EPISODE_STEPS_EVAL = 27_000
    N_CHECKPOINTS = 3

    total_reward_per_episode = TotalRewardPerEpisode()
    time_per_step = TimePerStep()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        DoEveryNStep() do t, agent, env, obs
            with_logger(lg) do
                @info "training" loss = agent.policy.learner.loss
            end
        end,
        DoEveryNEpisode() do t, agent, env, obs
            with_logger(lg) do
                @info "training" reward = total_reward_per_episode.rewards[end] log_step_increment =
                    0
            end
        end,
        DoEveryNStep(EVALUATION_FREQ) do t, agent, env, obs
            @info "evaluating agent at $t step..."
            flush(stdout)
            Flux.testmode!(agent)
            old_explorer = agent.policy.explorer
            agent.policy.explorer = EpsilonGreedyExplorer(0.001)  # set evaluation epsilon
            h = ComposedHook(TotalRewardPerEpisode(), StepsPerEpisode())
            s = @elapsed run(
                agent,
                atari_env_factory(name, STATE_SIZE, N_FRAMES, MAX_EPISODE_STEPS_EVAL;),#= seed=nothing =#
                StopAfterStep(125_000; is_show_progress = false),
                h,
            )
            res = (
                avg_length = mean(h[2].steps[1:end-1]),
                avg_score = mean(h[1].rewards[1:end-1]),
            )
            push!(evaluation_result, res)
            agent.policy.explorer = old_explorer
            Flux.trainmode!(agent)
            @info "finished evaluating agent in $s seconds" avg_length = res.avg_length avg_score =
                res.avg_score
            with_logger(lg) do
                @info "evaluating" avg_length = res.avg_length avg_score = res.avg_score log_step_increment = 0
            end
            flush(stdout)

            RLCore.save(joinpath(save_dir, string(t)), agent; is_save_trajectory = false)  # saving trajectory will take about 27G disk space each time
            BSON.@save joinpath(save_dir, string(t), "stats.bson") total_reward_per_episode time_per_step evaluation_result

            # only keep recent 3 checkpoints
            old_checkpoint_folder =
                joinpath(save_dir, string(t - EVALUATION_FREQ * N_CHECKPOINTS))
            if isdir(old_checkpoint_folder)
                rm(old_checkpoint_folder; force = true, recursive = true)
            end
        end,
    )

    N_TRAINING_STEPS = 50_000_000
    stop_condition = StopAfterStep(N_TRAINING_STEPS)

    description = """
    This experiment uses almost the same config in [dopamine](https://github.com/google/dopamine/blob/master/dopamine/agents/dqn/configs/dqn.gin). But do notice that there are some minor differences:

    - The RMSProp in Flux do not support center option (also the epsilon is not the same).
    - The image resize method used here is provided by ImageTransformers, which is not the same with the one in cv2.
    - `max_steps_per_episode` is not set, this might affect the evaluation result slightly.

    The testing environment is $name.
    Agent and statistic info will be saved to: $(joinpath(save_dir, string(N_TRAINING_STEPS)))
    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`

    To load the agent and statistic info:
    ```
    agent = RLCore.load("$(joinpath(save_dir, string(N_TRAINING_STEPS)))", Agent)
    BSON.@load joinpath("$(joinpath(save_dir, string(N_TRAINING_STEPS)))", "stats.bson") total_reward_per_episode time_per_step evaluation_result
    ```
    """

    Experiment(agent, env, stop_condition, hook, description)
end

function RLCore.Experiment(
    ::Val{:Dopamine},
    ::Val{:Rainbow},
    ::Val{:Atari},
    name::AbstractString;
    save_dir = nothing,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyymmddHHMMSS")
        save_dir = joinpath(pwd(), "checkpoints", "Dopamine_Rainbow_Atari_$(name)_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    N_FRAMES = 4
    STATE_SIZE = (84, 84)
    env = atari_env_factory(name, STATE_SIZE, N_FRAMES;)#= seed=(135, 246) =#
    N_ACTIONS = length(get_action_space(env))
    N_ATOMS = 51
    init = seed_glorot_uniform()#= seed=341 =#

    create_model() =
        Chain(
            x -> x ./ 255,
            CrossCor((8, 8), N_FRAMES => 32, relu; stride = 4, pad = 2, init = init),
            CrossCor((4, 4), 32 => 64, relu; stride = 2, pad = 2, init = init),
            CrossCor((3, 3), 64 => 64, relu; stride = 1, pad = 1, init = init),
            x -> reshape(x, :, size(x)[end]),
            Dense(11 * 11 * 64, 512, relu),
            Dense(512, N_ATOMS * N_ACTIONS),
        ) |> gpu

    agent = Agent(
        policy = QBasedPolicy(
            learner = RainbowLearner(
                approximator = NeuralNetworkApproximator(
                    model = create_model(),
                    optimizer = ADAM(0.0000625),
                ),  # epsilon is not set here
                target_approximator = NeuralNetworkApproximator(model = create_model()),
                n_actions = N_ACTIONS,
                n_atoms = N_ATOMS,
                Vₘₐₓ = 10.0f0,
                Vₘᵢₙ = -10.0f0,
                update_freq = 4,
                γ = 0.99f0,
                update_horizon = 3,
                batch_size = 32,
                stack_size = N_FRAMES,
                min_replay_history = 20_000,
                loss_func = logitcrossentropy_unreduced,
                target_update_freq = 8_000,
                # seed=89,
            ),
            explorer = EpsilonGreedyExplorer(
                ϵ_init = 1.0,
                ϵ_stable = 0.01,
                decay_steps = 250_000,
                kind = :linear,
                # seed=97,
            ),
        ),
        trajectory = CircularCompactPSARTSATrajectory(
            capacity = 1_000_000,
            state_type = Float32,
            state_size = STATE_SIZE,
        ),
    )

    evaluation_result = []
    EVALUATION_FREQ = 250_000
    MAX_EPISODE_STEPS_EVAL = 27_000
    N_CHECKPOINTS = 3

    total_reward_per_episode = TotalRewardPerEpisode()
    time_per_step = TimePerStep()
    steps_per_episode = StepsPerEpisode()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        steps_per_episode,
        DoEveryNStep() do t, agent, env, obs
            with_logger(lg) do
                @info "training" loss = agent.policy.learner.loss
            end
        end,
        DoEveryNEpisode() do t, agent, env, obs
            with_logger(lg) do
                @info "training" reward = total_reward_per_episode.rewards[end] episode_length =
                    steps_per_episode.steps[end] log_step_increment = 0
            end
        end,
        DoEveryNStep(EVALUATION_FREQ) do t, agent, env, obs
            @info "evaluating agent at $t step..."
            flush(stdout)
            Flux.testmode!(agent)
            old_explorer = agent.policy.explorer
            agent.policy.explorer = EpsilonGreedyExplorer(0.001)  # set evaluation epsilon
            h = ComposedHook(TotalRewardPerEpisode(), StepsPerEpisode())
            s = @elapsed run(
                agent,
                atari_env_factory(name, STATE_SIZE, N_FRAMES, MAX_EPISODE_STEPS_EVAL;),#= seed=nothing =#
                StopAfterStep(125_000; is_show_progress = false),
                h,
            )
            res = (
                avg_length = mean(h[2].steps[1:end-1]),
                avg_score = mean(h[1].rewards[1:end-1]),
            )
            push!(evaluation_result, res)
            agent.policy.explorer = old_explorer
            Flux.trainmode!(agent)
            @info "finished evaluating agent in $s seconds" avg_length = res.avg_length avg_score =
                res.avg_score
            with_logger(lg) do
                @info "evaluating" avg_length = res.avg_length avg_score = res.avg_score log_step_increment = 0
            end
            flush(stdout)

            RLCore.save(joinpath(save_dir, string(t)), agent; is_save_trajectory = false)  # saving trajectory will take about 27G disk space each time
            BSON.@save joinpath(save_dir, string(t), "stats.bson") total_reward_per_episode time_per_step evaluation_result

            # only keep recent 3 checkpoints
            old_checkpoint_folder =
                joinpath(save_dir, string(t - EVALUATION_FREQ * N_CHECKPOINTS))
            if isdir(old_checkpoint_folder)
                rm(old_checkpoint_folder; force = true, recursive = true)
            end
        end,
    )

    N_TRAINING_STEPS = 50_000_000
    stop_condition = StopAfterStep(N_TRAINING_STEPS)

    description = """
    This experiment uses alomost the same config in [dopamine](https://github.com/google/dopamine/blob/master/dopamine/agents/rainbow/configs/rainbow.gin). But do notice that there are some minor differences:

    - The epsilon in ADAM optimizer is not changed
    - The image resize method used here is provided by ImageTransformers, which is not the same with the one in cv2.

    The testing environment is $name.
    Agent and statistic info will be saved to: `$(joinpath(save_dir, string(N_TRAINING_STEPS)))`
    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`

    To load the agent and statistic info:
    ```
    agent = RLCore.load("$(joinpath(save_dir, string(N_TRAINING_STEPS)))", Agent)
    BSON.@load joinpath("$(joinpath(save_dir, string(N_TRAINING_STEPS)))", "stats.bson") total_reward_per_episode time_per_step evaluation_result
    ```
    """

    Experiment(agent, env, stop_condition, hook, description)
end

function RLCore.Experiment(
    ::Val{:Dopamine},
    ::Val{:IQN},
    ::Val{:Atari},
    name::AbstractString;
    save_dir = nothing,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyymmddHHMMSS")
        save_dir = joinpath(pwd(), "checkpoints", "Dopamine_IQN_Atari_$(name)_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    N_FRAMES = 4
    STATE_SIZE = (84, 84)
    MAX_STEPS_PER_EPISODE = 27_000

    env = atari_env_factory(name, STATE_SIZE, N_FRAMES;)#= seed=(135, 246) =#
    N_ACTIONS = length(get_action_space(env))
    Nₑₘ = 64

    init = seed_glorot_uniform()#= seed=341 =#

    create_model() =
        ImplicitQuantileNet(
            ψ = Chain(
                x -> x ./ 255,
                CrossCor((8, 8), N_FRAMES => 32, relu; stride = 4, pad = 2, init = init),
                CrossCor((4, 4), 32 => 64, relu; stride = 2, pad = 2, init = init),
                CrossCor((3, 3), 64 => 64, relu; stride = 1, pad = 1, init = init),
                x -> reshape(x, :, size(x)[end]),
            ),
            ϕ = Dense(Nₑₘ, 11 * 11 * 64, relu),
            header = Chain(Dense(11 * 11 * 64, 512, relu), Dense(512, N_ACTIONS)),
        ) |> gpu

    agent = Agent(
        policy = QBasedPolicy(
            learner = IQNLearner(
                approximator = NeuralNetworkApproximator(
                    model = create_model(),
                    optimizer = ADAM(0.00005),  # epsilon is not set here
                ),
                target_approximator = NeuralNetworkApproximator(model = create_model()),
                κ = 1.0f0,
                N = 64,
                N′ = 64,
                Nₑₘ = Nₑₘ,
                K = 32,
                γ = 0.99f0,
                stack_size = 4,
                batch_size = 32,
                update_horizon = 3,
                min_replay_history = 20_000,
                update_freq = 4,
                target_update_freq = 8_000,
                default_priority = 1.0f2,
                # seed=105,
                # device_seed=237,
            ),
            explorer = EpsilonGreedyExplorer(
                ϵ_init = 1.0,
                ϵ_stable = 0.01,
                decay_steps = 250_000,
                kind = :linear,
                # seed=99,
            ),
        ),
        trajectory = CircularCompactSARTSATrajectory(
            capacity = 1_000_000,
            state_type = Float32,
            state_size = STATE_SIZE,
        ),
    )

    evaluation_result = []
    EVALUATION_FREQ = 250_000
    MAX_EPISODE_STEPS_EVAL = 27_000
    N_CHECKPOINTS = 3

    total_reward_per_episode = TotalRewardPerEpisode()
    time_per_step = TimePerStep()
    steps_per_episode = StepsPerEpisode()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        steps_per_episode,
        DoEveryNStep() do t, agent, env, obs
            with_logger(lg) do
                @info "training" loss = agent.policy.learner.loss
            end
        end,
        DoEveryNEpisode() do t, agent, env, obs
            with_logger(lg) do
                @info "training" reward = total_reward_per_episode.rewards[end] episode_length =
                    steps_per_episode.steps[end] log_step_increment = 0
            end
        end,
        DoEveryNStep(EVALUATION_FREQ) do t, agent, env, obs
            @info "evaluating agent at $t step..."
            flush(stdout)
            Flux.testmode!(agent)
            old_explorer = agent.policy.explorer
            agent.policy.explorer = EpsilonGreedyExplorer(0.001)  # set evaluation epsilon
            h = ComposedHook(TotalRewardPerEpisode(), StepsPerEpisode())
            s = @elapsed run(
                agent,
                atari_env_factory(name, STATE_SIZE, N_FRAMES, MAX_EPISODE_STEPS_EVAL;),#= seed=nothing =#
                StopAfterStep(125_000; is_show_progress = false),
                h,
            )
            res = (
                avg_length = mean(h[2].steps[1:end-1]),
                avg_score = mean(h[1].rewards[1:end-1]),
            )
            push!(evaluation_result, res)
            agent.policy.explorer = old_explorer
            Flux.trainmode!(agent)
            @info "finished evaluating agent in $s seconds" avg_length = res.avg_length avg_score =
                res.avg_score
            with_logger(lg) do
                @info "evaluating" avg_length = res.avg_length avg_score = res.avg_score log_step_increment = 0
            end
            flush(stdout)

            RLCore.save(joinpath(save_dir, string(t)), agent; is_save_trajectory = false)  # saving trajectory will take about 27G disk space each time
            BSON.@save joinpath(save_dir, string(t), "stats.bson") total_reward_per_episode time_per_step evaluation_result

            # only keep recent 3 checkpoints
            old_checkpoint_folder =
                joinpath(save_dir, string(t - EVALUATION_FREQ * N_CHECKPOINTS))
            if isdir(old_checkpoint_folder)
                rm(old_checkpoint_folder; force = true, recursive = true)
            end
        end,
    )

    N_TRAINING_STEPS = 50_000_000
    stop_condition = StopAfterStep(N_TRAINING_STEPS)

    description = """
    This experiment uses alomost the same config in [dopamine](https://github.com/google/dopamine/blob/master/dopamine/agents/implicit_quantile/configs/implicit_quantile.gin). But do notice that there are some minor differences:

    - The epsilon in ADAM optimizer is not changed
    - The image resize method used here is provided by ImageTransformers, which is not the same with the one in cv2.

    The testing environment is $name.
    Agent and statistic info will be saved to: `$(joinpath(save_dir, string(N_TRAINING_STEPS)))`
    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`

    To load the agent and statistic info:
    ```
    agent = RLCore.load("$(joinpath(save_dir, string(N_TRAINING_STEPS)))", Agent)
    BSON.@load joinpath("$(joinpath(save_dir, string(N_TRAINING_STEPS)))", "stats.bson") total_reward_per_episode time_per_step evaluation_result
    ```
    """

    Experiment(agent, env, stop_condition, hook, description)
end

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:A2C},
    ::Val{:Atari},
    name::AbstractString;
    save_dir = nothing,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyymmddHHMMSS")
        save_dir = joinpath(pwd(), "checkpoints", "JuliaRL_A2C_Atari_$(name)_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    N_ENV = 16
    UPDATE_FREQ = 5
    N_FRAMES = 4
    STATE_SIZE = (84, 84)
    env = MultiThreadEnv([atari_env_factory(name, STATE_SIZE, N_FRAMES;) for i in 1:N_ENV]) #= seed = i =#
    N_ACTIONS = length(get_action_space(env[1]))

    init = seed_glorot_uniform()#= seed=341 =#

    # share model
    model = Chain(
        x -> x ./ 255,
        CrossCor((8, 8), N_FRAMES => 32, relu; stride = 4, pad = 2, init = init),
        CrossCor((4, 4), 32 => 64, relu; stride = 2, pad = 2, init = init),
        CrossCor((3, 3), 64 => 64, relu; stride = 1, pad = 1, init = init),
        x -> reshape(x, :, size(x)[end]),
        Dense(11 * 11 * 64, 512, relu),
    )

    agent = Agent(
        policy = QBasedPolicy(
            learner = A2CLearner(
                approximator = ActorCritic(
                    actor = Chain(model, Dense(512, N_ACTIONS; initW = init)),
                    critic = Chain(model, Dense(512, 1; initW = init)),
                    optimizer = RMSProp(7e-4, 0.99),
                ) |> gpu,
                γ = 0.99f0,
                max_grad_norm = 0.5f0,
                actor_loss_weight = 1.0f0,
                critic_loss_weight = 0.25f0,
                entropy_loss_weight = 0.01f0,
            ),
            explorer = BatchExplorer(GumbelSoftmaxExplorer(;)),#= seed = s=#
        ),
        trajectory = CircularCompactSARTSATrajectory(;
            capacity = UPDATE_FREQ,
            state_type = Float32,
            state_size = (STATE_SIZE..., N_FRAMES, N_ENV),
            action_type = Int,
            action_size = (N_ENV,),
            reward_type = Float32,
            reward_size = (N_ENV,),
            terminal_type = Bool,
            terminal_size = (N_ENV,),
        ),
    )

    N_TRAINING_STEPS = 10_000_000
    EVALUATION_FREQ = 100_000
    MAX_EPISODE_STEPS_EVAL = 27_000
    N_CHECKPOINTS = 3
    stop_condition = StopAfterStep(N_TRAINING_STEPS)

    total_batch_reward_per_episode = TotalBatchRewardPerEpisode(N_ENV)
    batch_steps_per_episode = BatchStepsPerEpisode(N_ENV)
    evaluation_result = []

    hook = ComposedHook(
        total_batch_reward_per_episode,
        batch_steps_per_episode,
        DoEveryNStep(UPDATE_FREQ) do t, agent, env, obs
            learner = agent.policy.learner
            with_logger(lg) do
                @info "training" loss = learner.loss actor_loss = learner.actor_loss critic_loss =
                    learner.critic_loss entropy_loss = learner.entropy_loss norm =
                    learner.norm log_step_increment = UPDATE_FREQ
            end
        end,
        DoEveryNStep() do t, agent, env, obs
            with_logger(lg) do
                rewards = [
                    total_batch_reward_per_episode.rewards[i][end]
                    for i in 1:length(obs) if get_terminal(obs[i])
                ]
                if length(rewards) > 0
                    @info "training" rewards = mean(rewards) log_step_increment = 0
                end
                steps = [
                    batch_steps_per_episode.steps[i][end]
                    for i in 1:length(obs) if get_terminal(obs[i])
                ]
                if length(steps) > 0
                    @info "training" steps = mean(steps) log_step_increment = 0
                end
            end
        end,
        DoEveryNStep(EVALUATION_FREQ) do t, agent, env, obs
            @info "evaluating agent at $t step..."
            flush(stdout)
            Flux.testmode!(agent)
            h = TotalBatchRewardPerEpisode(N_ENV)
            s = @elapsed run(
                agent,
                MultiThreadEnv([
                    atari_env_factory(name, STATE_SIZE, N_FRAMES, MAX_EPISODE_STEPS_EVAL;) for i in 1:N_ENV
                ]),
                StopAfterStep(27_000; is_show_progress = false),
                h,
            )
            res = (avg_score = mean(Iterators.flatten(h.rewards)),)
            push!(evaluation_result, res)
            Flux.trainmode!(agent)
            @info "finished evaluating agent in $s seconds" avg_score = res.avg_score
            with_logger(lg) do
                @info "evaluating" avg_score = res.avg_score log_step_increment = 0
            end
            flush(stdout)

            RLCore.save(joinpath(save_dir, string(t)), agent;)
            BSON.@save joinpath(save_dir, string(t), "stats.bson") total_batch_reward_per_episode evaluation_result

            # only keep recent 3 checkpoints
            old_checkpoint_folder =
                joinpath(save_dir, string(t - EVALUATION_FREQ * N_CHECKPOINTS))
            if isdir(old_checkpoint_folder)
                rm(old_checkpoint_folder; force = true, recursive = true)
            end
        end,
    )

    description = """
    # Play Atari($name) with A2C

    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`
    """

    Experiment(agent, env, stop_condition, hook, description)
end

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:A2CGAE},
    ::Val{:Atari},
    name::AbstractString;
    save_dir = nothing,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyymmddHHMMSS")
        save_dir = joinpath(pwd(), "checkpoints", "JuliaRL_A2CGAE_Atari_$(name)_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    N_ENV = 16
    UPDATE_FREQ = 32
    N_FRAMES = 4
    STATE_SIZE = (84, 84)
    env = MultiThreadEnv([atari_env_factory(name, STATE_SIZE, N_FRAMES;) for i in 1:N_ENV]) #= seed = i =#
    N_ACTIONS = length(get_action_space(env[1]))

    init = seed_glorot_uniform()#= seed=341 =#

    # share model
    model = Chain(
        x -> x ./ 255,
        CrossCor((8, 8), N_FRAMES => 32, relu; stride = 4, pad = 2, init = init),
        CrossCor((4, 4), 32 => 64, relu; stride = 2, pad = 2, init = init),
        CrossCor((3, 3), 64 => 64, relu; stride = 1, pad = 1, init = init),
        x -> reshape(x, :, size(x)[end]),
        Dense(11 * 11 * 64, 512, relu),
    )

    agent = Agent(
        policy = QBasedPolicy(
            learner = A2CGAELearner(
                approximator = ActorCritic(
                    actor = Chain(model, Dense(512, N_ACTIONS; initW = init)),
                    critic = Chain(model, Dense(512, 1; initW = init)),
                    optimizer = RMSProp(7e-4, 0.99),
                ) |> gpu,
                γ = 0.99f0,
                λ = 0.95f0,
                max_grad_norm = 0.5f0,
                actor_loss_weight = 1.0f0,
                critic_loss_weight = 0.25f0,
                entropy_loss_weight = 0.01f0,
            ),
            explorer = BatchExplorer(GumbelSoftmaxExplorer(;)),#= seed = s=#
        ),
        trajectory = CircularCompactSARTSATrajectory(;
            capacity = UPDATE_FREQ,
            state_type = Float32,
            state_size = (STATE_SIZE..., N_FRAMES, N_ENV),
            action_type = Int,
            action_size = (N_ENV,),
            reward_type = Float32,
            reward_size = (N_ENV,),
            terminal_type = Bool,
            terminal_size = (N_ENV,),
        ),
    )

    N_TRAINING_STEPS = 10_000_000
    EVALUATION_FREQ = 100_000
    MAX_EPISODE_STEPS_EVAL = 27_000
    N_CHECKPOINTS = 3
    stop_condition = StopAfterStep(N_TRAINING_STEPS)

    total_batch_reward_per_episode = TotalBatchRewardPerEpisode(N_ENV)
    batch_steps_per_episode = BatchStepsPerEpisode(N_ENV)
    evaluation_result = []

    hook = ComposedHook(
        total_batch_reward_per_episode,
        batch_steps_per_episode,
        DoEveryNStep(UPDATE_FREQ) do t, agent, env, obs
            learner = agent.policy.learner
            with_logger(lg) do
                @info "training" loss = learner.loss actor_loss = learner.actor_loss critic_loss =
                    learner.critic_loss entropy_loss = learner.entropy_loss norm =
                    learner.norm log_step_increment = UPDATE_FREQ
            end
        end,
        DoEveryNStep() do t, agent, env, obs
            with_logger(lg) do
                rewards = [
                    total_batch_reward_per_episode.rewards[i][end]
                    for i in 1:length(obs) if get_terminal(obs[i])
                ]
                if length(rewards) > 0
                    @info "training" rewards = mean(rewards) log_step_increment = 0
                end
                steps = [
                    batch_steps_per_episode.steps[i][end]
                    for i in 1:length(obs) if get_terminal(obs[i])
                ]
                if length(steps) > 0
                    @info "training" steps = mean(steps) log_step_increment = 0
                end
            end
        end,
        DoEveryNStep(EVALUATION_FREQ) do t, agent, env, obs
            @info "evaluating agent at $t step..."
            flush(stdout)
            Flux.testmode!(agent)
            h = TotalBatchRewardPerEpisode(N_ENV)
            s = @elapsed run(
                agent,
                MultiThreadEnv([
                    atari_env_factory(name, STATE_SIZE, N_FRAMES, MAX_EPISODE_STEPS_EVAL;) for i in 1:N_ENV
                ]),
                StopAfterStep(27_000; is_show_progress = false),
                h,
            )
            res = (avg_score = mean(Iterators.flatten(h.rewards)),)
            push!(evaluation_result, res)
            Flux.trainmode!(agent)
            @info "finished evaluating agent in $s seconds" avg_score = res.avg_score
            with_logger(lg) do
                @info "evaluating" avg_score = res.avg_score log_step_increment = 0
            end
            flush(stdout)

            RLCore.save(joinpath(save_dir, string(t)), agent;)
            BSON.@save joinpath(save_dir, string(t), "stats.bson") total_batch_reward_per_episode evaluation_result

            # only keep recent 3 checkpoints
            old_checkpoint_folder =
                joinpath(save_dir, string(t - EVALUATION_FREQ * N_CHECKPOINTS))
            if isdir(old_checkpoint_folder)
                rm(old_checkpoint_folder; force = true, recursive = true)
            end
        end,
    )

    description = """
    # Play Atari($name) with A2CGAE

    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`
    """

    Experiment(agent, env, stop_condition, hook, description)
end

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:PPO},
    ::Val{:Atari},
    name::AbstractString;
    save_dir = nothing,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyymmddHHMMSS")
        save_dir = joinpath(pwd(), "checkpoints", "JuliaRL_PPO_Atari_$(name)_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    N_ENV = 32
    UPDATE_FREQ = 64
    N_FRAMES = 4
    STATE_SIZE = (84, 84)
    env = MultiThreadEnv([atari_env_factory(name, STATE_SIZE, N_FRAMES;) for i in 1:N_ENV]) #= seed = i =#
    N_ACTIONS = length(get_action_space(env[1]))
    INIT_CLIP_RANGE = 0.1f0
    INIT_LEARNING_RATE = 1e-3

    init = seed_orthogonal()#= seed=341 =#

    # share model
    model = Chain(
        x -> x ./ 255,
        CrossCor((8, 8), N_FRAMES => 32, relu; stride = 4, pad = 2, init = init),
        CrossCor((4, 4), 32 => 64, relu; stride = 2, pad = 2, init = init),
        CrossCor((3, 3), 64 => 64, relu; stride = 1, pad = 1, init = init),
        x -> reshape(x, :, size(x)[end]),
        Dense(11 * 11 * 64, 512, relu),
    )

    agent = Agent(
        policy = QBasedPolicy(
            learner = PPOLearner(
                approximator = ActorCritic(
                    actor = Chain(model, Dense(512, N_ACTIONS; initW = init)),
                    critic = Chain(model, Dense(512, 1; initW = init)),
                    optimizer = ADAM(INIT_LEARNING_RATE),  # decrease learning rate with a hook
                ) |> gpu,
                γ = 0.99f0,
                λ = 0.98f0,
                clip_range = INIT_CLIP_RANGE,  # decrease with a hook
                max_grad_norm = 1.0f0,
                n_microbatches = 4,
                n_epochs = 4,
                actor_loss_weight = 1.0f0,
                critic_loss_weight = 1.0f0,
                entropy_loss_weight = 0.01f0,
            ),
            explorer = BatchExplorer(GumbelSoftmaxExplorer(;)),#= seed = s=#
        ),
        trajectory = PPOTrajectory(;
            capacity = UPDATE_FREQ,
            state_type = Float32,
            state_size = (STATE_SIZE..., N_FRAMES, N_ENV),
            action_type = Int,
            action_size = (N_ENV,),
            action_log_prob_type = Float32,
            action_log_prob_size = (N_ENV,),
            reward_type = Float32,
            reward_size = (N_ENV,),
            terminal_type = Bool,
            terminal_size = (N_ENV,),
        ),
    )

    N_TRAINING_STEPS = 10_000_000
    EVALUATION_FREQ = 100_000
    MAX_EPISODE_STEPS_EVAL = 27_000
    N_CHECKPOINTS = 3
    stop_condition = StopAfterStep(N_TRAINING_STEPS)

    total_batch_reward_per_episode = TotalBatchRewardPerEpisode(N_ENV)
    batch_steps_per_episode = BatchStepsPerEpisode(N_ENV)
    evaluation_result = []

    hook = ComposedHook(
        total_batch_reward_per_episode,
        batch_steps_per_episode,
        DoEveryNStep(UPDATE_FREQ) do t, agent, env, obs
            learner = agent.policy.learner
            with_logger(lg) do
                @info "training" loss = mean(learner.loss) actor_loss =
                    mean(learner.actor_loss) critic_loss = mean(learner.critic_loss) entropy_loss =
                    mean(learner.entropy_loss) norm = mean(learner.norm) log_step_increment =
                    UPDATE_FREQ
            end
        end,
        DoEveryNStep(UPDATE_FREQ) do t, agent, env, obs
            decay = (N_TRAINING_STEPS - t) / N_TRAINING_STEPS
            agent.policy.learner.approximator.optimizer.eta =
                INIT_LEARNING_RATE * decay
            agent.policy.learner.clip_range = INIT_CLIP_RANGE * Float32(decay)
        end,
        DoEveryNStep() do t, agent, env, obs
            with_logger(lg) do
                rewards = [
                    total_batch_reward_per_episode.rewards[i][end]
                    for i in 1:length(obs) if get_terminal(obs[i])
                ]
                if length(rewards) > 0
                    @info "training" rewards = mean(rewards) log_step_increment = 0
                end
                steps = [
                    batch_steps_per_episode.steps[i][end]
                    for i in 1:length(obs) if get_terminal(obs[i])
                ]
                if length(steps) > 0
                    @info "training" steps = mean(steps) log_step_increment = 0
                end
            end
        end,
        DoEveryNStep(EVALUATION_FREQ) do t, agent, env, obs
            @info "evaluating agent at $t step..."
            flush(stdout)
            Flux.testmode!(agent)
            h = TotalBatchRewardPerEpisode(N_ENV)
            s = @elapsed run(
                agent,
                MultiThreadEnv([
                    atari_env_factory(name, STATE_SIZE, N_FRAMES, MAX_EPISODE_STEPS_EVAL;) for i in 1:4
                ]),
                StopAfterStep(27_000; is_show_progress = false),
                h,
            )
            res = (avg_score = mean(Iterators.flatten(h.rewards)),)
            push!(evaluation_result, res)
            Flux.trainmode!(agent)
            @info "finished evaluating agent in $s seconds" avg_score = res.avg_score
            with_logger(lg) do
                @info "evaluating" avg_score = res.avg_score log_step_increment = 0
            end
            flush(stdout)

            RLCore.save(joinpath(save_dir, string(t)), agent;)
            BSON.@save joinpath(save_dir, string(t), "stats.bson") total_batch_reward_per_episode evaluation_result

            # only keep recent 3 checkpoints
            old_checkpoint_folder =
                joinpath(save_dir, string(t - EVALUATION_FREQ * N_CHECKPOINTS))
            if isdir(old_checkpoint_folder)
                rm(old_checkpoint_folder; force = true, recursive = true)
            end
        end,
    )

    description = """
    # Play Atari($name) with PPO

    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`
    """

    Experiment(agent, env, stop_condition, hook, description)
end


using Images, MXNet

### LOADING THE MODEL
const MODEL_NAME = "weights/mobilenet-v2/mobilenet_v2"
const MODEL_CLASS_NAMES = "weights/mobilenet-v2/synset.txt"

nnet = mx.load_checkpoint(MODEL_NAME, 0, mx.FeedForward; context = mx.gpu());
synset = readlines(MODEL_CLASS_NAMES);

### SEARCH FOR A LAYER OF INTERESTS
layers = mx.get_internals(nnet.arch);
layers_flatten = nothing
layers_to_remove = Symbol[]

# We iterate over all layers until we find the one matching our requirements
# and remove the ones to follow after
for i = 1:2000

    layer = layers[i];
    layer_name = mx.get_name(layer)

    if layers_flatten == nothing && layer_name == :pool6
        layers_flatten = layer
    elseif layers_flatten != nothing
        push!(layers_to_remove, layer_name)
        if layer_name in [:softmax, :label, :prob] break end
    end
end

nnet.arch = @mx.chain layers_flatten => Flatten()
map(x -> delete!(nnet.arg_params, x), layers_to_remove);
map(x -> delete!(nnet.aux_params, x), layers_to_remove);

mx.save_checkpoint(nnet, "weights/mobilenet-v2/MobiletNet-FE", mx.OptimizationState(1, 0, 0, 0))


"""
converts a primitive into a drawable for `window`
"""
function Drawable(window::WT, primitive::PT) where {WT, PT}
    error("Can't draw $PT with window type: $WT")
end

"""
Creates a rasterizer clojure for a certain window type
"""
function rasterizer{WT}(
        window::WT,
        vertexarray::AbstractArray,
        uniforms::Tuple,
        vertexshader::Function,
        fragmentshader::Function;
        geometryshader = nothing,
        max_primitives = 4,
        primitive_in = :points,
        primitive_out = :triangle_strip,
    )
    error("No rasterizer available for window type: $WT")
end

"Enables debugging information for a composable"
@field Debugging = false

"Name of something, usually a String or Symbol"
@field Name = gensym(:Visualize)

"Defines if an object is visible or not"
@field Visible = true

"Defines if an object is focused or not"
@field Focused = false

@field Rotation = Quaternion(1, 0, 0, 0)
@field Scale = Vec3f0(1)
@field Position = Vec3f0(0)
@field Translation = Vec3f0(0)

abstract type Transformation <: ReactiveComposable end

@composed type Transform <: Transformation
    Scale
    Rotation
    Translation
end

@field XAxis
@field YAxis
@field ZAxis
@field Primitive
"""
The Area of a viewable. Can be a Rectangle or a tuple with two elements.
"""
@field Area = IRect(0, 0, 500, 500)

"""
Resolution is the smaller brother of Area. It can be NTuple{2, <: Number} or a Vec{2, <: Number}.
"""
@field Resolution

function default(::Type{Resolution}, x::Partial)
    area = get(x, Area) # APIDOC this gets the area if in `x`, or creates a default area
    widths(area)
end

function default(::Type{Area}, x::Partial)
    haskey(x, Resolution) && return convertfor(Area, x, get(x, Resolution))
    default(Area)
end
function convertfor(::Type{Area}, x::Partial, value)
    isa(value, Tuple) && length(value) == 2 && return IRect(0, 0, value...)
    isa(value, SimpleRectangle) && return IRect(origin(value), widths(value))
    throw(UsageError(Area, value))
end

@field Spacetransform = identity

@field Projection = eye(Mat4f0)
@field View = eye(Mat4f0)
@field ProjectionView = eye(Mat4f0)

@field BoundingBox = AABB{Float32}()

@composed type Shared
    Transform
    BoundingBox
end

@field Color = Vec4f0(1f0, 1f0, 1f0, 1f0)

"""
Thickness e.g. of a line or stroke.
"""
@field Thickness = 2f0
@field Vertices = error("Needs default")


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

    using  Printf, ..AngularMomentum, ..Basics, ..Continuum, ..Defaults, ..InteractionStrength, ..ManyElectron, ..Nuclear,
                   ..PlasmaShift, ..Radial, ..SpinAngular, ..TableStrings

    """
    `struct  Settings`  ... defines a type for the details and parameters of computing Auger lines.

        + calcAnisotropy          ::Bool                         ... True, if the intrinsic alpha_2,4 angular parameters are to be calculated,
                                                                     and false otherwise.
        + printBefore             ::Bool                         ... True, if all energies and lines are printed before their evaluation.
        + lineSelection           ::LineSelection                ... Specifies the selected levels, if any.
        + minAugerEnergy          ::Float64                      ... Minimum energy of free (Auger) electrons to be included.
        + maxAugerEnergy          ::Float64                      ... Maximum energy of free (Auger) electrons to be included.
        + maxKappa                ::Int64                        ... Maximum kappa value of partial waves to be included.
        + operator                ::AbstractEeInteraction        ... Auger operator that is to be used for evaluating the Auger amplitudes:
                                                                     allowed values are: CoulombInteraction(), BreitInteraction(), ...
    """
    struct Settings
        calcAnisotropy            ::Bool
        printBefore               ::Bool
        lineSelection             ::LineSelection
        minAugerEnergy            ::Float64
        maxAugerEnergy            ::Float64
        maxKappa                  ::Int64
        operator                  ::AbstractEeInteraction
    end

    """
    `AutoIonization.Settings()`  ... constructor for the default values of AutoIonization line computations
    """
    function Settings()
        Settings(false, false, LineSelection(), 0., 10e5, 100, CoulombInteraction())
    end

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

            calcAnisotropy=..,      printBefore=..,
            minAugerEnergy=..,      maxAugerEnergy=..,          maxKappa=..,            operator=..)

        ... constructor for modifying the given AutoIonization.Settings by 'overwriting' the previously selected parameters.
    """
    function Settings(set::AutoIonization.Settings;
        calcAnisotropy::Union{Nothing,Bool}=nothing,            printBefore::Union{Nothing,Bool}=nothing,
        lineSelection::Union{Nothing,LineSelection}=nothing,
        minAugerEnergy::Union{Nothing,Float64}=nothing,         maxAugerEnergy::Union{Nothing,Float64}=nothing,
        maxKappa::Union{Nothing,Int64}=nothing,                 operator::Union{Nothing,String}=nothing)

        if  calcAnisotropy  == nothing   calcAnisotropyx  = set.calcAnisotropy    else  calcAnisotropyx  = calcAnisotropy   end
        if  printBefore     == nothing   printBeforex     = set.printBefore       else  printBeforex     = printBefore      end
        if  lineSelection   == nothing   lineSelectionx   = set.lineSelection     else  lineSelectionx   = lineSelection    end
        if  minAugerEnergy  == nothing   minAugerEnergyx  = set.minAugerEnergy    else  minAugerEnergyx  = minAugerEnergy   end
        if  maxAugerEnergy  == nothing   maxAugerEnergyx  = set.maxAugerEnergy    else  maxAugerEnergyx  = maxAugerEnergy   end
        if  maxKappa        == nothing   maxKappax        = set.maxKappa          else  maxKappax        = maxKappa         end
        if  operator        == nothing   operatorx        = set.operator          else  operatorx        = operator         end

        Settings( calcAnisotropyx, printBeforex, lineSelectionx, minAugerEnergyx, maxAugerEnergyx, maxKappax, operatorx)
    end

    # `Base.show(io::IO, settings::AutoIonization.Settings)`  ... prepares a proper printout of the variable settings::AutoIonization.Settings.
    function Base.show(io::IO, settings::AutoIonization.Settings)
        println(io, "calcAnisotropy:                $(settings.calcAnisotropy)  ")
        println(io, "printBefore:                   $(settings.printBefore)  ")
        println(io, "lineSelection:                 $(settings.lineSelection)  ")
        println(io, "minAugerEnergy:                $(settings.minAugerEnergy)  ")
        println(io, "maxAugerEnergy:                $(settings.maxAugerEnergy)  ")
        println(io, "maxKappa:                      $(settings.maxKappa)  ")
        println(io, "operator:                      $(settings.operator)  ")
    end

    """
    `struct  Channel`
        ... defines a type for a AutoIonization channel to help characterize a scattering (continuum) state of many
            electron-states with a single free electron.

        + kappa          ::Int64                ... partial-wave of the free electron
        + symmetry       ::LevelSymmetry        ... total angular momentum and parity of the scattering state
        + phase          ::Float64              ... phase of the partial wave
        + amplitude      ::Complex{Float64}     ... Auger amplitude associated with the given channel.
    """
    struct  Channel
        kappa            ::Int64
        symmetry         ::LevelSymmetry
        phase            ::Float64
        amplitude        ::Complex{Float64}
    end

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

        + initialLevel   ::Level           ... initial-(state) level
        + finalLevel     ::Level           ... final-(state) level
        + electronEnergy ::Float64         ... Energy of the (incoming free) electron.
        + totalRate      ::Float64         ... Total rate of this line.
        + angularAlpha   ::Float64         ... Angular alpha_2 coefficient.
        + hasChannels    ::Bool
            ... Determines whether the individual scattering (sub-) channels are defined in terms of their free-electron energy, kappa
                and the total angular momentum/parity as well as the amplitude, or not.
        + channels       ::Array{AutoIonization.Channel,1}  ... List of AutoIonization channels of this line.
    """
    struct  Line
        initialLevel     ::Level
        finalLevel       ::Level
        electronEnergy   ::Float64
        totalRate        ::Float64
        angularAlpha     ::Float64
        hasChannels      ::Bool
        channels         ::Array{AutoIonization.Channel,1}
    end

    """
    `AutoIonization.Line(initialLevel::Level, finalLevel::Level, totalRate::Float64)`
        ... constructor for an AutoIonization line between a specified initial and final level.
    """
    function Line(initialLevel::Level, finalLevel::Level, totalRate::Float64)
        Line(initialLevel, finalLevel, 0., totalRate, 0., false, AutoIonization.Channel[])
    end

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

    """
    `AutoIonization.amplitude(kind::String, channel::AutoIonization.Channel, continuumLevel::Level, initialLevel::Level,
                              grid::Radial.Grid; printout::Bool=true)`
        ... to compute the kind in  CoulombInteraction(), BreitInteraction(), CoulombBreit()   Auger amplitude
            <(alpha_f J_f, kappa) J_i || O^(Auger, kind) || alpha_i J_i>  due to the interelectronic interaction for the given
            final and initial level. A value::ComplexF64 is returned.
    """
    function amplitude(kind::AbstractEeInteraction, channel::AutoIonization.Channel, continuumLevel::Level, initialLevel::Level, grid::Radial.Grid;
                       printout::Bool=true)
        nt = length(continuumLevel.basis.csfs);    ni = length(initialLevel.basis.csfs);    partial = Subshell(9,channel.kappa)
        if  printout  printstyled("Compute ($kind) Auger matrix of dimension $nt x $ni in the continuum- and initial-state bases " *
                                  "for the transition [$(initialLevel.index)- ...] and for partial wave $(string(partial)[2:end]) ... ",
                                  color=:light_green)    end
        matrix = zeros(ComplexF64, nt, ni)
        #
        if      kind in [ CoulombInteraction(), BreitInteraction(), CoulombBreit()]        ## pure V^Coulomb interaction
        #--------------------------------------------------------------------------
            for  r = 1:nt
                for  s = 1:ni
                    if  initialLevel.basis.csfs[s].J != initialLevel.J  ||  initialLevel.basis.csfs[s].parity != initialLevel.parity      continue    end
                    ##x wa = compute("angular coefficients: e-e, Ratip2013", continuumLevel.basis.csfs[r], initialLevel.basis.csfs[s])
                    # Calculate the spin-angular coefficients
                    if  Defaults.saRatip()
                        waR = compute("angular coefficients: e-e, Ratip2013", continuumLevel.basis.csfs[r], initialLevel.basis.csfs[s])
                        wa  = waR
                    end
                    if  Defaults.saGG()
                        subshellList = continuumLevel.basis.subshells
                        opa  = SpinAngular.TwoParticleOperator(0, plus, true)
                        waG2 = SpinAngular.computeCoefficients(opa, continuumLevel.basis.csfs[r], initialLevel.basis.csfs[s], subshellList)
                        wa   = [1.0, waG2]
                    end
                    if  Defaults.saRatip() && Defaults.saGG() && true
                        if  length(waR[2]) != 0     println("\n>> Angular coeffients from Ratip2013   = $(waR[2]) ")    end
                        if  length(waG2)   != 0     println(  ">> Angular coeffients from SpinAngular = $waG2 ")        end
                    end
                    #
                    me = 0.
                    for  coeff in wa[2]
                        if   kind in [ CoulombInteraction(), CoulombBreit()]
                            me = me + coeff.V * InteractionStrength.XL_Coulomb(coeff.nu,
                                                    continuumLevel.basis.orbitals[coeff.a], continuumLevel.basis.orbitals[coeff.b],
                                                    initialLevel.basis.orbitals[coeff.c],   initialLevel.basis.orbitals[coeff.d], grid)   end
                        if   kind in [ BreitInteraction(), CoulombBreit()]
                            me = me + coeff.V * InteractionStrength.XL_Breit(coeff.nu,
                                                    continuumLevel.basis.orbitals[coeff.a], continuumLevel.basis.orbitals[coeff.b],
                                                    initialLevel.basis.orbitals[coeff.c],   initialLevel.basis.orbitals[coeff.d], grid)   end
                    end
                    matrix[r,s] = me
                end
            end
            if  printout  printstyled("done. \n", color=:light_green)    end
            amplitude = transpose(continuumLevel.mc) * matrix * initialLevel.mc
            amplitude = im^Basics.subshell_l(Subshell(101, channel.kappa)) * exp( -im*channel.phase ) * amplitude
            #
            #
         elseif  kind == "H-E"
        #--------------------
            iLevel = finalLevel;   fLevel = initialLevel
            amplitude = 0.;    error("stop a")
        else    error("stop b")
        end

        return( amplitude )
    end

    """
    `AutoIonization.channelAmplitude(kind::String, channel::AutoIonization.Channel, energy::Float64, finalLevel::Level,
                                     initialLevel::Level, grid::Radial.Grid)`
        ... to compute the kind = (CoulombInteraction(), BreitInteraction(), CoulombBreit())   Auger amplitude
            <(alpha_f J_f, kappa) J_i || O^(Auger, kind) || alpha_i J_i>  due to the interelectronic interaction for the given final and
            initial level. A newChannel::AutoIonization.Channel is returned.
    """
    function channelAmplitude(kind::String, channel::AutoIonization.Channel, energy::Float64, finalLevel::Level,
                              initialLevel::Level, grid::Radial.Grid)
        newiLevel = Basics.generateLevelWithSymmetryReducedBasis(initialLevel)
        newiLevel = Basics.generateLevelWithExtraSubshell(Subshell(101, channel.kappa), newiLevel)
        newfLevel = Basics.generateLevelWithSymmetryReducedBasis(finalLevel)
        cOrbital, phase  = Continuum.generateOrbital(energy, Subshell(101, channel.kappa), newfLevel, grid, contSettings)
        newcLevel  = Basics.generateLevelWithExtraElectron(cOrbital, channel.symmetry, newfLevel)
        newChannel = AutoIonization.Channel(channel.kappa, channel.symmetry, phase, 0.)
        amplitude = AutoIonization.amplitude(kind, channel, newcLevel, newiLevel, grid)

        newChannel = AutoIonization.Channel(newChannel.kappa, newChannel.symmetry, newChannel.phase, amplitude)
        return( newChannel )
    end

    """
    `AutoIonization.computeAmplitudesProperties(line::AutoIonization.Line, nm::Nuclear.Model, grid::Radial.Grid, nrContinuum::Int64,
                                                settings::AutoIonization.Settings; printout::Bool=true)`
        ... to compute all amplitudes and properties of the given line; a line::AutoIonization.Line is returned for which the amplitudes
            and properties are now evaluated.
    """
    function computeAmplitudesProperties(line::AutoIonization.Line, nm::Nuclear.Model, grid::Radial.Grid, nrContinuum::Int64, settings::AutoIonization.Settings;
                                         printout::Bool=true)
        newChannels = AutoIonization.Channel[];   contSettings = Continuum.Settings(false, nrContinuum);   rate = 0.
        for channel in line.channels
            newiLevel = Basics.generateLevelWithSymmetryReducedBasis(line.initialLevel, line.initialLevel.basis.subshells)
            newfLevel = Basics.generateLevelWithSymmetryReducedBasis(line.finalLevel, newiLevel.basis.subshells)
            newiLevel = Basics.generateLevelWithExtraSubshell(Subshell(101, channel.kappa), newiLevel)
            cOrbital, phase  = Continuum.generateOrbitalForLevel(line.electronEnergy, Subshell(101, channel.kappa), newfLevel, nm, grid, contSettings)
            newcLevel  = Basics.generateLevelWithExtraElectron(cOrbital, channel.symmetry, newfLevel)
            newChannel = AutoIonization.Channel(channel.kappa, channel.symmetry, phase, 0.)
            amplitude  = AutoIonization.amplitude(settings.operator, newChannel, newcLevel, newiLevel, grid, printout=printout)
            rate       = rate + conj(amplitude) * amplitude
            push!( newChannels, AutoIonization.Channel(newChannel.kappa, newChannel.symmetry, newChannel.phase, amplitude) )
        end
        totalRate = 2pi* rate;   angularAlpha = 0.
        newLine   = AutoIonization.Line(line.initialLevel, line.finalLevel, line.electronEnergy, totalRate, angularAlpha, true, newChannels)
        #
        if  settings.calcAnisotropy    angularAlpha = AutoIonization.computeIntrinsicAlpha(2, newLine)
            newLine   = AutoIonization.Line(line.initialLevel, line.finalLevel, line.electronEnergy, totalRate, real(angularAlpha), true, newChannels)
        end

        return( newLine )
    end

    """
    `AutoIonization.computeAmplitudesPropertiesPlasma(line::AutoIonization.Line, nm::Nuclear.Model, grid::Radial.Grid, nrContinuum::Int64,
                                                      settings::PlasmaShift.AugerSettings; printout::Bool=true)`
        ... to compute all amplitudes and properties of the given line but for the given plasma model; a line::AutoIonization.Line is returned
            for which the amplitudes and properties are now evaluated.
    """
    function computeAmplitudesPropertiesPlasma(line::AutoIonization.Line, nm::Nuclear.Model, grid::Radial.Grid, nrContinuum::Int64,
                                               settings::PlasmaShift.AugerSettings; printout::Bool=true)
        newChannels = AutoIonization.Channel[];   contSettings = Continuum.Settings(false, nrContinuum);   rate = 0.
        for channel in line.channels
            newiLevel = Basics.generateLevelWithSymmetryReducedBasis(line.initialLevel)
            newiLevel = Basics.generateLevelWithExtraSubshell(Subshell(101, channel.kappa), newiLevel)
            newfLevel = Basics.generateLevelWithSymmetryReducedBasis(line.finalLevel)
            @warn "Adapt a proper continuum orbital for the plasma potential"
            cOrbital, phase  = Continuum.generateOrbitalForLevel(line.electronEnergy, Subshell(101, channel.kappa), newfLevel, nm, grid, contSettings)
            newcLevel  = Basics.generateLevelWithExtraElectron(cOrbital, channel.symmetry, newfLevel)
            newChannel = AutoIonization.Channel(channel.kappa, channel.symmetry, phase, 0.)
            @warn "Adapt a proper Auger amplitude for the plasma e-e interaction"
            amplitude = 1.0
            # amplitude  = AutoIonization.amplitude(settings.operator, newChannel, newcLevel, newiLevel, grid)
            rate       = rate + conj(amplitude) * amplitude
            push!( newChannels, AutoIonization.Channel(newChannel.kappa, newChannel.symmetry, newChannel.phase, amplitude) )
        end
        totalRate = 2pi* rate;   angularAlpha = 0.
        newLine   = AutoIonization.Line(line.initialLevel, line.finalLevel, line.electronEnergy, totalRate, angularAlpha, true, newChannels)

        return( newLine )
    end

    """
    `AutoIonization.computeIntrinsicAlpha(k::Int64, line::AutoIonization.Line)`
        ... to compute the intrinsic alpha_k anisotropy parameter for the given line. A value::Float64 is returned.
    """
    function  computeIntrinsicAlpha(k::Int64, line::AutoIonization.Line)
        if  !line.hasChannels   error("No channels are defined for the given AutoIonization.line.")                   end
        wn = 0.;    for  channel in line.channels    wn = wn + conj(channel.amplitude) * channel.amplitude   end
        wa = 0.;    Ji = line.initialLevel.J;    Jf = line.finalLevel.J;
        for  cha  in line.channels
            j = AngularMomentum.kappa_j(cha.kappa);    l = AngularMomentum.kappa_l(cha.kappa)
            for  chp  in line.channels
                jp = AngularMomentum.kappa_j(chp.kappa);    lp = AngularMomentum.kappa_l(chp.kappa)
                wa = wa + sqrt( AngularMomentum.bracket([l, lp, j, jp]) ) *
                          AngularMomentum.ClebschGordan(l, AngularM64(0), lp, AngularM64(0), AngularJ64(k), AngularM64(0)) *
                          AngularMomentum.Wigner_6j(Ji, j, Jf, jp, Ji, AngularJ64(k)) *
                          AngularMomentum.Wigner_6j(l,  j, AngularJ64(1//2), jp, lp, AngularJ64(k)) *
                          cha.amplitude * conj(chp.amplitude)
            end
        end
        value = AngularMomentum.phaseFactor([Ji, +1, Jf, +1, AngularJ64(k), -1, AngularJ64(1//2)]) *
                sqrt(Basics.twice(Ji) + 1) * wa / wn

        if  false
            # Calculate the value given by M.H. Chen, PRA 47 (1993) 3733; this does not agree exactly so far.
            wa = 0.
            for  cha  in line.channels
                j = AngularMomentum.kappa_j(cha.kappa);    l = AngularMomentum.kappa_l(cha.kappa);    phase = cha.phase
                for  chp  in line.channels
                    jp = AngularMomentum.kappa_j(chp.kappa);    lp = AngularMomentum.kappa_l(chp.kappa);    phasep = chp.phase
                    wa = wa + AngularMomentum.phaseFactor([Ji, +1, Jf, -1, AngularJ64(1//2)]) * im^(Float64(l) - Float64(lp)) * cos(phase - phasep) *
                              sqrt( AngularMomentum.bracket([l, lp, j, jp, AngularJ64(k), Ji]) ) *
                              AngularMomentum.Wigner_3j(lp, l, AngularJ64(k), AngularJ64(0), AngularJ64(0), AngularJ64(0)) *
                              AngularMomentum.Wigner_6j(j, jp, AngularJ64(k), lp, l, AngularJ64(1//2)) *
                              AngularMomentum.Wigner_6j(Ji, Ji, AngularJ64(k), jp, j, Jf) *
                              cha.amplitude * conj(chp.amplitude)
                end
            end
            println("*** Comparison value = $value    Chen-value = $(wa/wn)")
        end

        return( value )
    end

    """
    `AutoIonization.computeLines(finalMultiplet::Multiplet, initialMultiplet::Multiplet, nm::Nuclear.Model, grid::Radial.Grid,
                                 settings::AutoIonization.Settings; output=true, printout::Bool=true)`
        ... to compute the Auger transition amplitudes and all properties as requested by the given settings. A list of
            lines::Array{AutoIonization.Lines} is returned.
    """
    function  computeLines(finalMultiplet::Multiplet, initialMultiplet::Multiplet, nm::Nuclear.Model, grid::Radial.Grid,
                           settings::AutoIonization.Settings; output=true, printout::Bool=true)
        println("")
        printstyled("AutoIonization.computeLines(): The computation of Auger rates and properties starts now ... \n", color=:light_green)
        printstyled("------------------------------------------------------------------------------------------- \n", color=:light_green)
        println("")
        lines = AutoIonization.determineLines(finalMultiplet, initialMultiplet, settings)
        # Display all selected lines before the computations start
        if  settings.printBefore    AutoIonization.displayLines(lines)    end
        # Determine maximum energy and check for consistency of the grid
        maxEnergy = 0.;   for  line in lines   maxEnergy = max(maxEnergy, line.electronEnergy)   end
        nrContinuum = Continuum.gridConsistency(maxEnergy, grid)
        # Calculate all amplitudes and requested properties
        newLines = AutoIonization.Line[]
        for  line in lines
            newLine = AutoIonization.computeAmplitudesProperties(line, nm, grid, nrContinuum, settings)
            push!( newLines, newLine)
        end
        # Print all results to screen
        AutoIonization.displayRates(stdout, newLines, settings)
        AutoIonization.displayLifetimes(stdout, newLines)
        printSummary, iostream = Defaults.getDefaults("summary flag/stream")
        if  printSummary   AutoIonization.displayRates(iostream, newLines, settings);   AutoIonization.displayLifetimes(iostream, newLines)     end
        #
        if    output    return( lines )
        else            return( nothing )
        end
    end

    """
    `AutoIonization.computeLinesCascade(finalMultiplet::Multiplet, initialMultiplet::Multiplet, nm::Nuclear.Model, grid::Radial.Grid,
                                        settings::AutoIonization.Settings; output::Bool=true, printout::Bool=true)`
        ... to compute the Auger transition amplitudes and all properties as requested by the given settings. The computations
            and printout is adapted for large cascade computations by including only lines with at least one channel and by sending
            all printout to a summary file only. A list of lines::Array{AutoIonization.Lines} is returned.
    """
    function  computeLinesCascade(finalMultiplet::Multiplet, initialMultiplet::Multiplet, nm::Nuclear.Model, grid::Radial.Grid,
                                  settings::AutoIonization.Settings; output::Bool=true, printout::Bool=true)

        lines = AutoIonization.determineLines(finalMultiplet, initialMultiplet, settings)
        # Display all selected lines before the computations start
        # if  settings.printBefore    AutoIonization.displayLines(lines)    end
        # Determine maximum energy and check for consistency of the grid
        maxEnergy = 0.;   for  line in lines   maxEnergy = max(maxEnergy, line.electronEnergy)   end
        nrContinuum = Continuum.gridConsistency(maxEnergy, grid)
        # Calculate all amplitudes and requested properties
        newLines = AutoIonization.Line[]
        for  (i,line)  in  enumerate(lines)
            if  rem(i,10) == 0    println("> Auger line $i:  ... calculated ")    end
            newLine = AutoIonization.computeAmplitudesProperties(line, nm, grid, nrContinuum, settings, printout=printout)
            push!( newLines, newLine)
        end
        # Print all results to a summary file, if requested
        printSummary, iostream = Defaults.getDefaults("summary flag/stream")
        if  printSummary   AutoIonization.displayRates(iostream, newLines, settings)     end
        #
        if    output    return( newLines )
        else            return( nothing )
        end
    end

    """
    `AutoIonization.computeLinesFromOrbitals(finalMultiplet::Multiplet, initialMultiplet::Multiplet, nm::Nuclear.Model, grid::Radial.Grid,
                                             settings::AutoIonization.Settings, contOrbitals::Dict{Subshell, Orbital}; output::Bool=true, printout::Bool=true)`
        ... to compute the Auger transition amplitudes and all properties as requested by the given settings but by using the given set of
            continuum orbitals. The computations and printout is adapted for large cascade computations by including only lines with at least
            one channel and by sending all printout to a summary file only. A list of lines::Array{AutoIonization.Lines} is returned.
    """
    function  computeLinesFromOrbitals(finalMultiplet::Multiplet, initialMultiplet::Multiplet, nm::Nuclear.Model, grid::Radial.Grid,
                                       settings::AutoIonization.Settings, contOrbitals::Dict{Subshell, Orbital}; output::Bool=true, printout::Bool=true)

        lines = AutoIonization.determineLines(finalMultiplet, initialMultiplet, settings)
        # Calculate all amplitudes and requested properties
        newLines = AutoIonization.Line[]
        for  (i,line)  in  enumerate(lines)
            if  rem(i,500) == 0    println("> Auger line $i:  ... calculated ")    end
            # Calculate the individual channels with the given orbitals
            newChannels = AutoIonization.Channel[];   rate = 0.
            for channel in line.channels
                newiLevel  = Basics.generateLevelWithSymmetryReducedBasis(line.initialLevel, line.initialLevel.basis.subshells)
                newfLevel  = Basics.generateLevelWithSymmetryReducedBasis(line.finalLevel, newiLevel.basis.subshells)
                sh         = Subshell(101, channel.kappa)
                if haskey(contOrbitals, sh)   cOrbital = contOrbitals[sh]      else    println(">>> skip Auger channel for $sh");   continue    end
                newiLevel  = Basics.generateLevelWithExtraSubshell(sh, newiLevel)
                newcLevel  = Basics.generateLevelWithExtraElectron(cOrbital, channel.symmetry, newfLevel)
                newChannel = AutoIonization.Channel(channel.kappa, channel.symmetry, 0., 0.)
                amplitude  = AutoIonization.amplitude(settings.operator, newChannel, newcLevel, newiLevel, grid, printout=printout)
                rate       = rate + conj(amplitude) * amplitude
                push!( newChannels, AutoIonization.Channel(newChannel.kappa, newChannel.symmetry, newChannel.phase, amplitude) )
            end
            totalRate = 2pi* rate;   angularAlpha = 0.
            newLine   = AutoIonization.Line(line.initialLevel, line.finalLevel, line.electronEnergy, totalRate, angularAlpha, true, newChannels)
            ##
            ## if  settings.calcAnisotropy    angularAlpha = AutoIonization.computeIntrinsicAlpha(2, newLine)
            ##  newLine   = AutoIonization.Line(line.initialLevel, line.finalLevel, line.electronEnergy, totalRate, real(angularAlpha), true, newChannels)
            ## end
            push!( newLines, newLine)
        end
        # Print all results to a summary file, if requested
        printSummary, iostream = Defaults.getDefaults("summary flag/stream")
        if  printSummary   AutoIonization.displayRates(iostream, newLines, settings)     end
        #
        if    output    return( newLines )
        else            return( nothing )
        end
    end

    """
    `AutoIonization.computeLinesPlasma(finalMultiplet::Multiplet, initialMultiplet::Multiplet, nm::Nuclear.Model, grid::Radial.Grid,
                                       settings::PlasmaShift.AugerSettings; output=true)`
        ... to compute the Auger transition amplitudes and all properties as requested by the given settings. A list of
            lines::Array{AutoIonization.Lines} is returned.
    """
    function  computeLinesPlasma(finalMultiplet::Multiplet, initialMultiplet::Multiplet, nm::Nuclear.Model, grid::Radial.Grid,
                                 settings::PlasmaShift.AugerSettings; output=true)
        println("")
        printstyled("AutoIonization.computeLinesPlasma(): The computation of Auger rates starts now ... \n", color=:light_green)
        printstyled("---------------------------------------------------------------------------------- \n", color=:light_green)
        println("")
        augerSettings = AutoIonization.Settings(false, settings.printBefore, settings.selectLines, settings.selectedLines, 0., 1.0e6, 100, CoulombInteraction())
        lines = AutoIonization.determineLines(finalMultiplet, initialMultiplet, augerSettings)
        # Display all selected lines before the computations start
        if  settings.printBefore    AutoIonization.displayLines(lines)    end
        # Determine maximum energy and check for consistency of the grid
        maxEnergy = 0.;   for  line in lines   maxEnergy = max(maxEnergy, line.electronEnergy)   end
        nrContinuum = Continuum.gridConsistency(maxEnergy, grid)
        # Calculate all amplitudes and requested properties
        newLines = AutoIonization.Line[]
        for  line in lines
            newLine = AutoIonization.computeAmplitudesPropertiesPlasma(line, nm, grid, nrContinuum, settings)
            push!( newLines, newLine)
        end
        # Print all results to screen
        AutoIonization.displayRates(stdout, newLines, augerSettings)
        AutoIonization.displayLifetimes(stdout, newLines)
        printSummary, iostream = Defaults.getDefaults("summary flag/stream")
        if  printSummary   AutoIonization.displayRates(iostream, newLines, augerSettings);   AutoIonization.displayLifetimes(iostream, newLines)     end
        #
        if    output    return( lines )
        else            return( nothing )
        end
    end

    """
    `AutoIonization.determineChannels(finalLevel::Level, initialLevel::Level, settings::AutoIonization.Settings)`
        ... to determine a list of Auger Channel for a transitions from the initial to final level and by taking into account the particular
            settings of for this computation; an Array{AutoIonization.Channel,1} is returned.
    """
    function determineChannels(finalLevel::Level, initialLevel::Level, settings::AutoIonization.Settings)
        channels  = AutoIonization.Channel[];
        symi      = LevelSymmetry(initialLevel.J, initialLevel.parity);    symf = LevelSymmetry(finalLevel.J, finalLevel.parity)
        kappaList = AngularMomentum.allowedKappaSymmetries(symi, symf)
        for  kappa in kappaList
            if  abs(kappa) > settings.maxKappa      continue    end
            push!(channels, AutoIonization.Channel(kappa, symi, 0., Complex(0.)) )
        end
        return( channels )
    end

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

    """
    `AutoIonization.displayLines(lines::Array{AutoIonization.Line,1})`
        ... to display a list of lines and channels that have been selected due to the prior settings. A neat table of all selected
            transitions and energies is printed but nothing is returned otherwise.
    """
    function  displayLines(lines::Array{AutoIonization.Line,1})
        nx = 150
        println(" ")
        println("  Selected Auger lines:")
        println(" ")
        println("  ", TableStrings.hLine(nx))
        sa = "  ";   sb = "  "
        sa = sa * TableStrings.center(18, "i-level-f"; na=2);                                sb = sb * TableStrings.hBlank(20)
        sa = sa * TableStrings.center(18, "i--J^P--f"; na=4);                                sb = sb * TableStrings.hBlank(22)
        sa = sa * TableStrings.center(14, "Energy"; na=4);
        sb = sb * TableStrings.center(14, TableStrings.inUnits("energy"); na=4)
        sa = sa * TableStrings.center(14, "Energy e_A"; na=4);
        sb = sb * TableStrings.center(14, TableStrings.inUnits("energy"); na=4)
        sa = sa * TableStrings.flushleft(37, "List of kappas and total symmetries"; na=4)
        sb = sb * TableStrings.flushleft(37, "partial (total J^P)                "; na=4)
        println(sa);    println(sb);    println("  ", TableStrings.hLine(nx))
        #
        for  line in lines
            sa  = "  ";    isym = LevelSymmetry( line.initialLevel.J, line.initialLevel.parity)
                           fsym = LevelSymmetry( line.finalLevel.J,   line.finalLevel.parity)
            sa = sa * TableStrings.center(18, TableStrings.levels_if(line.initialLevel.index, line.finalLevel.index); na=2)
            sa = sa * TableStrings.center(18, TableStrings.symmetries_if(isym, fsym); na=4)
            sa = sa * @sprintf("%.8e", Defaults.convertUnits("energy: from atomic", line.initialLevel.energy))  * "    "
            sa = sa * @sprintf("%.8e", Defaults.convertUnits("energy: from atomic", line.electronEnergy))       * "   "
            kappaSymmetryList = Tuple{Int64,LevelSymmetry}[]
            for  i in 1:length(line.channels)
                push!( kappaSymmetryList, (line.channels[i].kappa, line.channels[i].symmetry) )
            end
            sa = sa * TableStrings.kappaSymmetryTupels(80, kappaSymmetryList)
            println( sa )
        end
        println("  ", TableStrings.hLine(nx), "\n")
        #
        return( nothing )
    end

    """
    `AutoIonization.displayLifetimes(stream::IO, lines::Array{AutoIonization.Line,1})`
        ... to list all lifetimes as associated with the selected lines. A neat table is printed but nothing is returned otherwise.
    """
    function  displayLifetimes(stream::IO, lines::Array{AutoIonization.Line,1})
        nx = 104
        println(stream, " ")
        println(stream, "  Auger lifetimes, total rates and widths:")
        println(stream, " ")
        println(stream, "  ", TableStrings.hLine(nx))
        sa = "  ";   sb = "  "
        sa = sa * TableStrings.center(10, "Level";    na=2);                           sb = sb * TableStrings.hBlank(12)
        sa = sa * TableStrings.center( 8, "J^P";      na=4);                           sb = sb * TableStrings.hBlank(12)
        sa = sa * TableStrings.center(12, "Lifetime"; na=4);
        sb = sb * TableStrings.center(12,TableStrings.inUnits("time"); na=4)
        sa = sa * TableStrings.center(14, "Total rate"; na=6);
        sb = sb * TableStrings.center(14,TableStrings.inUnits("rate"); na=4)
        sa = sa * TableStrings.center(42, "Widths"; na=2);
        sb = sb * TableStrings.center(42, "  Hartrees         Kaysers           eV"; na=2)
        println(stream, sa);    println(stream, sb);    println(stream, "  ", TableStrings.hLine(nx))
        #
        notYetDone = trues(1000)
        for  line in lines
            if  notYetDone[line.initialLevel.index]
                notYetDone[line.initialLevel.index] = false
                totalRate = 0.
                for  ln in lines
                    if  ln.initialLevel.index == line.initialLevel.index   totalRate = totalRate + ln.totalRate    end
                end
                sa  = "  ";    isym = LevelSymmetry( line.initialLevel.J, line.initialLevel.parity)
                sa = sa * TableStrings.center(10, TableStrings.level(line.initialLevel.index); na=2)
                sa = sa * TableStrings.center( 8, string(isym); na=4)
                sa = sa * @sprintf("%.6e", Defaults.convertUnits("time: from atomic",  1/totalRate))            * "     "
                sa = sa * @sprintf("%.6e", Defaults.convertUnits("rate: from atomic",    totalRate))            * "      "
                sa = sa * @sprintf("%.6e", totalRate)                                                           * "    "
                sa = sa * @sprintf("%.6e", Defaults.convertUnits("energy: from atomic to Kayser",  totalRate))  * "    "
                sa = sa * @sprintf("%.6e", Defaults.convertUnits("energy: from atomic to eV",      totalRate))  * "    "
                println(stream, sa)
            end
        end
        println(stream, "  ", TableStrings.hLine(nx))
        #
        return( nothing )
    end

    """
    `AutoIonization.displayRates(stream::IO, lines::Array{AutoIonization.Line,1}, settings::AutoIonization.Settings)`
        ... to list all results, energies, rates, etc. of the selected lines. A neat table is printed but nothing is returned
            otherwise.
    """
    function  displayRates(stream::IO, lines::Array{AutoIonization.Line,1}, settings::AutoIonization.Settings)
        nx = 106
        println(stream, " ")
        if  settings.calcAnisotropy    println(stream, "  Auger rates and intrinsic angular parameters: \n")
        else                           println(stream, "  Auger rates (without angular parameters): \n")        end
        println(stream, "  ", TableStrings.hLine(nx))
        sa = "  ";   sb = "  "
        sa = sa * TableStrings.center(18, "i-level-f"; na=2);                         sb = sb * TableStrings.hBlank(20)
        sa = sa * TableStrings.center(18, "i--J^P--f"; na=4);                         sb = sb * TableStrings.hBlank(22)
        sa = sa * TableStrings.center(12, "Energy"   ; na=2);
        sb = sb * TableStrings.center(12,TableStrings.inUnits("energy"); na=4)
        sa = sa * TableStrings.center(14, "Electron energy"   ; na=2);
        sb = sb * TableStrings.center(14,TableStrings.inUnits("energy"); na=2)
        sa = sa * TableStrings.center(14, "Auger rate"; na=2);
        sb = sb * TableStrings.center(14, TableStrings.inUnits("rate"); na=2)
        sa = sa * TableStrings.center(15, "alpha_2"; na=2);                           sb = sb * TableStrings.hBlank(18)
        println(stream, sa);    println(stream, sb);    println(stream, "  ", TableStrings.hLine(nx))
        #
        for  line in lines
            sa  = "  ";    isym = LevelSymmetry( line.initialLevel.J, line.initialLevel.parity)
                           fsym = LevelSymmetry( line.finalLevel.J,   line.finalLevel.parity)
            sa = sa * TableStrings.center(18, TableStrings.levels_if(line.initialLevel.index, line.finalLevel.index); na=2)
            sa = sa * TableStrings.center(18, TableStrings.symmetries_if(isym, fsym); na=4)
            sa = sa * @sprintf("%.6e", Defaults.convertUnits("energy: from atomic", line.initialLevel.energy))  * "    "
            sa = sa * @sprintf("%.6e", Defaults.convertUnits("energy: from atomic", line.electronEnergy))       * "    "
            sa = sa * @sprintf("%.6e", Defaults.convertUnits("rate: from atomic", line.totalRate))              * "    "
            sa = sa * TableStrings.flushright(13, @sprintf("%.4e", line.angularAlpha))            * "    "
            println(stream, sa)
        end
        println(stream, "  ", TableStrings.hLine(nx))
        #
        return( nothing )
    end

end # module


using Gen
import Random

using FunctionalCollections
using ReverseDiff

function strip_lineinfo(expr::Expr)
    @assert !(expr.head == :line)
    new_args = []
    for arg in expr.args
        if (isa(arg, Expr) && arg.head == :line) || isa(arg, LineNumberNode)
        elseif isa(arg, Expr) && arg.head == :block
            stripped = strip_lineinfo(arg)
            append!(new_args, stripped.args)
        else
            push!(new_args, strip_lineinfo(arg))
        end
    end
    Expr(expr.head, new_args...)
end

function strip_lineinfo(expr)
    expr
end

############################
# reverse mode AD for fill #
############################

function Base.fill(x::ReverseDiff.TrackedReal{V,D,O}, n::Integer) where {V,D,O}
    tp = ReverseDiff.tape(x)
    out = ReverseDiff.track(fill(ReverseDiff.value(x), n), V, tp)
    ReverseDiff.record!(tp, ReverseDiff.SpecialInstruction, fill, (x, n), out)
    return out
end

@noinline function ReverseDiff.special_reverse_exec!(instruction::ReverseDiff.SpecialInstruction{typeof(fill)})
    x, n = instruction.input
    output = instruction.output
    ReverseDiff.istracked(x) && ReverseDiff.increment_deriv!(x, sum(ReverseDiff.deriv(output)))
    ReverseDiff.unseed!(output)
    return nothing
end

@noinline function ReverseDiff.special_forward_exec!(instruction::ReverseDiff.SpecialInstruction{typeof(fill)})
    x, n = instruction.input
    ReverseDiff.value!(instruction.output, fill(ReverseDiff.value(x), n))
    return nothing
end

#########
# model #
#########

@compiled @gen function datum(x::Float64, @ad(inlier_std::Float64), @ad(outlier_std::Float64),
                                @ad(slope::Float64), @ad(intercept::Float64))
    is_outlier::Bool = @addr(bernoulli(0.5), :z)
    std::Float64 = is_outlier ? inlier_std : outlier_std
    y::Float64 = @addr(normal(x * slope + intercept, std), :y)
    return y
end

data = plate(datum)

function compute_data_change(inlier_std_change, outlier_std_change, slope_change, intercept_change)
    if all([c !== nothing && (c == NoChange() || !c[1]) for c in [
            inlier_std_change, outlier_std_change, slope_change, intercept_change]])
        NoChange()
    else
        nothing
    end
end

@compiled @gen function model(xs::Vector{Float64})
    n::Int = length(xs)
    inlier_log_std::Float64 = @addr(normal(0, 2), :inlier_std)
    outlier_log_std::Float64 = @addr(normal(0, 2), :outlier_std)
    inlier_std::Float64 = exp(inlier_log_std)
    outlier_std::Float64 = exp(outlier_log_std)
    slope::Float64 = @addr(normal(0, 2), :slope)
    intercept::Float64 = @addr(normal(0, 2), :intercept)
    inlier_std_change::Union{Tuple{Bool,Float64},Nothing} = @change(:inlier_std)
    outlier_std_change::Union{Tuple{Bool,Float64},Nothing} = @change(:outlier_std)
    slope_change::Union{Tuple{Bool,Float64},Nothing} = @change(:slope)
    intercept_change::Union{Tuple{Bool,Float64},Nothing} = @change(:intercept)
    change::Union{NoChange,Nothing} = compute_data_change(
        inlier_std_change, outlier_std_change, slope_change, intercept_change)
    @addr(data(xs, fill(inlier_std, n), fill(outlier_std, n),
                                                   fill(slope, n), fill(intercept, n)),
                                          :data, change)
end

#######################
# inference operators #
#######################

@sel function slope_intercept_selector()
    @select(:slope)
    @select(:intercept)
end

@sel function std_selector()
    @select(:inlier_std)
    @select(:outlier_std)
end

@compiled @gen function flip_z(z::Bool)
    @addr(bernoulli(z ? 0.0 : 1.0), :z)
end

data_proposal = at_dynamic(flip_z, Int)

@compiled @gen function is_outlier_proposal(prev, i::Int)
    prev_z::Bool = get_assignment(prev)[:data => i => :z]
    # TODO introduce shorthand @addr(flip_z(zs[i]), :data => i)
    @addr(data_proposal(i, (prev_z,)), :data)
end

@compiled @gen function observe_datum(y::Float64)
    @addr(dirac(y), :y)
end

observe_data = plate(observe_datum)

@compiled @gen function observer(ys::Vector{Float64})
    @addr(observe_data(ys), :data)
end

Gen.load_generated_functions()

#####################
# generate data set #
#####################

Random.seed!(1)

prob_outlier = 0.5
true_inlier_noise = 0.5
true_outlier_noise = 5.0
true_slope = -1
true_intercept = 2
xs = collect(range(-5, stop=5, length=200))
ys = Float64[]
for (i, x) in enumerate(xs)
    if rand() < prob_outlier
        y = true_slope * x + true_intercept + randn() * true_inlier_noise
    else
        y = true_slope * x + true_intercept + randn() * true_outlier_noise
    end
    push!(ys, y)
end

##################
# run experiment #
##################

trace = simulate(model, (xs,))
datum_trace = simulate(datum, (1., 2., 3., 4., 5.))
(slope_intercept_selection,) = Gen.select(slope_intercept_selector, (), get_assignment(trace))
(std_selection,) = Gen.select(std_selector, (), get_assignment(trace))
slope_intercept_static_sel = StaticAddressSet(slope_intercept_selection)
std_static_sel = StaticAddressSet(std_selection)

#println("\n######################################################################\n")
#println("backprop_trace(model, ...)")
#println(strip_lineinfo(
    #Gen.codegen_backprop_trace(typeof(model), typeof(trace), typeof(slope_intercept_static_sel), Nothing)))
#println("\n######################################################################\n")

#println("\n######################################################################\n")
#println("backprop_trace(datum, ...)")
#println(strip_lineinfo(
    #Gen.codegen_backprop_trace(typeof(datum), typeof(datum_trace),
           #Gen.EmptyAddressSet, Nothing)))
#println("\n######################################################################\n")

function do_inference(n)
    observations = get_assignment(simulate(observer, (ys,)))

    # initial trace
    (trace, _) = generate(model, (xs,), observations)

    for i=1:n
        for j=1:5
            trace = map_optimize(model, slope_intercept_static_sel, trace, max_step_size=0.1, min_step_size=1e-10)
            trace = map_optimize(model, std_static_sel, trace, max_step_size=0.1, min_step_size=1e-10)
        end

        # step on the outliers
        for j=1:length(xs)
            trace = mh(model, is_outlier_proposal, (j,), trace)
        end

        score = get_call_record(trace).score

        # print
        assignment = get_assignment(trace)
        slope = assignment[:slope]
        intercept = assignment[:intercept]
        inlier_std = assignment[:inlier_std]
        outlier_std = assignment[:outlier_std]
        println("score: $score, slope: $slope, intercept: $intercept, inlier_std: $inlier_std, outlier_std: $outlier_std")
    end
end

@time do_inference(100)
@time do_inference(100)


struct LazyMulAdd{M,O,T<:NativeTypesV} <: Number
    data::T
    @inline LazyMulAdd{M,O,T}(data::T) where {M,O,T<:NativeTypesV} = new{M,O,T}(data)
end
# O for offset is kind of hard to read next to default of 0?
@inline LazyMulAdd{M,O}(data::T) where {M,O,T<:Union{Base.HWReal,AbstractSIMD}} = LazyMulAdd{M,O,T}(data)
@inline LazyMulAdd{M}(data::T) where {M,T<:NativeTypesV} = LazyMulAdd{M,0,T}(data)
@inline LazyMulAdd{0,O}(data::T) where {O,T<:Union{Base.HWReal,AbstractSIMD}} = @assert false#StaticInt{O}()
@inline LazyMulAdd{0}(data::T) where {T<:NativeTypesV} = @assert false#StaticInt{0}()
@inline LazyMulAdd(data::T, ::StaticInt{M}) where {M,T} = LazyMulAdd{M,0,T}(data)
@inline LazyMulAdd(data::T, ::StaticInt{M}, ::StaticInt{O}) where {M,O,T} = LazyMulAdd{M,O,T}(data)

@inline data(lm::LazyMulAdd) = data(getfield(lm, :data)) # calls data on inner for use with indexing (normally `data` only goes through one layer)

@inline _materialize(a::LazyMulAdd{M,O,I}) where {M,O,I} = vadd_fast(vmul_fast(StaticInt{M}(), getfield(a, :data)), StaticInt{O}())
@inline _materialize(x) = x
@inline Base.convert(::Type{T}, a::LazyMulAdd{M,O,I}) where {M,O,I,T<:Number} = convert(T, _materialize(a))
@inline Base.convert(::Type{LazyMulAdd{M,O,I}}, a::LazyMulAdd{M,O,I}) where {M,O,I} = a
# @inline Base.convert(::Type{LazyMulAdd{M,O,I}}, a::LazyMulAdd{M}) where {M,O,I} = a
# @inline Base.convert(::Type{LazyMulAdd{M,T,I}}, a::LazyMulAdd{M,StaticInt{O},I}) where {M,O,I,T} = a

Base.promote_rule(::Type{LazyMulAdd{M,O,I}}, ::Type{T}) where {M,O,I<:Number,T} = promote_type(I, T)
Base.promote_rule(::Type{LazyMulAdd{M,O,MM{W,X,I}}}, ::Type{T}) where {M,O,W,X,I,T} = promote_type(MM{W,X,I}, T)
Base.promote_rule(::Type{LazyMulAdd{M,O,Vec{W,I}}}, ::Type{T}) where {M,O,W,I,T} = promote_type(Vec{W,I}, T)

@inline lazymul(a, b) = vmul_fast(a, b)
@inline lazymul(::StaticInt{M}, b) where {M} = LazyMulAdd{M}(b)
@inline lazymul(::StaticInt{1}, b) = b
@inline lazymul(::StaticInt{0}, b) = StaticInt{0}()
@inline lazymul(a, ::StaticInt{M}) where {M} = LazyMulAdd{M}(a)
@inline lazymul(a, ::StaticInt{1}) = a
@inline lazymul(a, ::StaticInt{0}) = StaticInt{0}()
@inline lazymul(a::LazyMulAdd, ::StaticInt{1}) = a
@inline lazymul(::StaticInt{1}, a::LazyMulAdd) = a
@inline lazymul(a::LazyMulAdd, ::StaticInt{0}) = StaticInt{0}()
@inline lazymul(::StaticInt{0}, a::LazyMulAdd) = StaticInt{0}()
# @inline lazymul(a::LazyMulAdd, ::StaticInt{0}) = StaticInt{0}()
# @inline lazymul(::StaticInt{M}, b::MM{W,X}) where {W,M,X} = LazyMulAdd{M}(MM{W}(getfield(b, :data), StaticInt{M}()*StaticInt{X}()))
# @inline lazymul(a::MM{W,X}, ::StaticInt{M}) where {W,M,X} = LazyMulAdd{M}(MM{W}(getfield(a, :data), StaticInt{M}()*StaticInt{X}()))
@inline lazymul(a::MM{W,X}, ::StaticInt{M}) where {W,M,X} = MM{W}(vmul_fast(StaticInt{M}(), data(a)), StaticInt{X}() * StaticInt{M}())
@inline lazymul(::StaticInt{M}, b::MM{W,X}) where {W,M,X} = MM{W}(vmul_fast(StaticInt{M}(), data(b)), StaticInt{X}() * StaticInt{M}())
@inline lazymul(a::MM{W,X}, ::StaticInt{1}) where {W,X} = a
@inline lazymul(::StaticInt{1}, a::MM{W,X}) where {W,X} = a
@inline lazymul(a::MM{W,X}, ::StaticInt{0}) where {W,X} = StaticInt{0}()
@inline lazymul(::StaticInt{0}, a::MM{W,X}) where {W,X} = StaticInt{0}()

@inline lazymul(::StaticInt{M}, ::StaticInt{N}) where {M, N} = StaticInt{M}()*StaticInt{N}()

@inline lazymul(::StaticInt{M}, ::StaticInt{1}) where {M} = StaticInt{M}()
@inline lazymul(::StaticInt,    ::StaticInt{0}) = StaticInt{0}()

@inline lazymul(::StaticInt{1}, ::StaticInt{M}) where {M} = StaticInt{M}()
@inline lazymul(::StaticInt{0}, ::StaticInt   ) = StaticInt{0}()

@inline lazymul(::StaticInt{0}, ::StaticInt{0}) = StaticInt{0}()
@inline lazymul(::StaticInt{0}, ::StaticInt{1}) = StaticInt{0}()
@inline lazymul(::StaticInt{1}, ::StaticInt{0}) = StaticInt{0}()
@inline lazymul(::StaticInt{1}, ::StaticInt{1}) = StaticInt{1}()

@inline function lazymul(a::LazyMulAdd{M,O}, ::StaticInt{N}) where {M,O,N}
    LazyMulAdd(getfield(a, :data), StaticInt{M}() * StaticInt{N}(), StaticInt{O}() * StaticInt{N}())
end
@inline function lazymul(::StaticInt{M}, b::LazyMulAdd{N,O}) where {M,O,N}
    LazyMulAdd(getfield(b, :data), StaticInt{M}() * StaticInt{N}(), StaticInt{O}() * StaticInt{M}())
end

@inline function lazymul(a::LazyMulAdd{M,<:MM{W,X}}, ::StaticInt{N}) where {M,N,W,X}
    LazyMulAdd(MM{W}(getfield(a, :data), StaticInt{M}()*StaticInt{N}()*StaticInt{X}()), StaticInt{M}()*StaticInt{N}())
end
@inline function lazymul(::StaticInt{M}, b::LazyMulAdd{N,<:MM{W,X}}) where {M,N,W,X}
    LazyMulAdd(MM{W}(getfield(b, :data), StaticInt{M}()*StaticInt{N}()*StaticInt{X}()), StaticInt{M}()*StaticInt{N}())
end

@inline lazymul(a::LazyMulAdd{M,<:MM{W,X}}, ::StaticInt{0}) where {M,W,X} = StaticInt{0}()
@inline lazymul(::StaticInt{0}, b::LazyMulAdd{N,<:MM{W,X}}) where {N,W,X} = StaticInt{0}()
@inline lazymul(a::LazyMulAdd{M,<:MM{W,X}}, ::StaticInt{1}) where {M,W,X} = a
@inline lazymul(::StaticInt{1}, b::LazyMulAdd{N,<:MM{W,X}}) where {N,W,X} = b
@inline function lazymul(a::LazyMulAdd{M}, b::LazyMulAdd{N}) where {M,N}
    LazyMulAdd(vmul_fast(getfield(a, :data), getfield(b, :data)), StaticInt{M}()*StaticInt{N}())
end

vmul_fast(::LazyMulAdd{M,O}, ::StaticInt{0}) where {M,O} = Zero()
vmul_fast(::StaticInt{0}, ::LazyMulAdd{M,O}) where {M,O} = Zero()
vmul_fast(a::LazyMulAdd{M,O}, ::StaticInt{1}) where {M,O} = a
vmul_fast(::StaticInt{1}, a::LazyMulAdd{M,O}) where {M,O} = a
@inline function vmul_fast(a::LazyMulAdd{M,O}, ::StaticInt{N}) where {M,N,O}
    LazyMulAdd(getfield(a, :data), StaticInt{M}() * StaticInt{N}(), StaticInt{O}() * StaticInt{N}())
end
@inline function vmul_fast(::StaticInt{N}, a::LazyMulAdd{M,O}) where {M,N,O}
    LazyMulAdd(getfield(a, :data), StaticInt{M}() * StaticInt{N}(), StaticInt{O}() * StaticInt{N}())
end

@inline function Base.:(>>>)(a::LazyMulAdd{M,O}, ::StaticInt{N}) where {M,O,N}
    LazyMulAdd(getfield(a, :data), StaticInt{M}() >>> StaticInt{N}(), StaticInt{O}() >>> StaticInt{N}())
end

# The approach with `add_indices` is that we try and make `vadd_fast` behave well
# but for `i` and `j` type combinations where it's difficult,
# we can add specific `add_indices` methods that increment the pointer.
@inline function add_indices(p::Ptr, i, j) # generic fallback
    p, vadd_fast(i, j)
end
@inline vadd_fast(i::LazyMulAdd, ::Zero) = i
@inline vadd_fast(::Zero, i::LazyMulAdd) = i

# These following two definitions normally shouldn't be hit
@inline function vadd_fast(a::LazyMulAdd{M,O,MM{W,X,I}}, b::Integer) where {M,O,W,X,I}
    MM{W}(vadd_fast(vmul_fast(StaticInt{M}(), data(a)), b), StaticInt{X}() * StaticInt{M}())
end
@inline function vadd_fast(b::Integer, a::LazyMulAdd{M,O,MM{W,X,I}}) where {M,O,W,X,I}
    MM{W}(vadd_fast(vmul_fast(StaticInt{M}(), data(a)), b), StaticInt{X}() * StaticInt{M}())
end
@inline vadd_fast(a::LazyMulAdd{M,O,MM{W,X,I}}, ::Zero) where {M,O,W,X,I} = a
@inline vadd_fast(::Zero, a::LazyMulAdd{M,O,MM{W,X,I}}) where {M,O,W,X,I} = a

# because we should hit this instead:
@inline add_indices(p::Ptr, b::Integer, a::LazyMulAdd{M,O}) where {M,O} = (p + b, a)
@inline add_indices(p::Ptr, a::LazyMulAdd{M,O}, b::Integer) where {M,O} = (p + b, a)
# but in the case of `VecUnroll`s, which skip the `add_indices`, it's useful to still have the former two definitions.
# However, this also forces us to write:
@inline add_indices(p::Ptr, ::StaticInt{N}, a::LazyMulAdd{M,O}) where {M,O,N} = (p, vadd_fast(a, StaticInt{N}()))
@inline add_indices(p::Ptr, a::LazyMulAdd{M,O}, ::StaticInt{N}) where {M,O,N} = (p, vadd_fast(a, StaticInt{N}()))

@inline function vadd_fast(::StaticInt{N}, a::LazyMulAdd{M,O}) where {N,M,O}
    LazyMulAdd(getfield(a, :data), StaticInt{M}(), StaticInt{O}()+StaticInt{N}())
end
@inline function vadd_fast(a::LazyMulAdd{M,O}, ::StaticInt{N}) where {N,M,O}
    LazyMulAdd(getfield(a, :data), StaticInt{M}(), StaticInt{O}()+StaticInt{N}())
end
@inline function vadd_fast(::StaticInt{N}, a::LazyMulAdd{M,O,MM{W,X,I}}) where {N,M,O,W,X,I}
    LazyMulAdd(getfield(a, :data), StaticInt{M}(), StaticInt{O}()+StaticInt{N}())
end
@inline function vadd_fast(a::LazyMulAdd{M,O,MM{W,X,I}}, ::StaticInt{N}) where {N,M,O,W,X,I}
    LazyMulAdd(getfield(a, :data), StaticInt{M}(), StaticInt{O}()+StaticInt{N}())
end

@inline function vadd_fast(a::LazyMulAdd{M,O}, b::LazyMulAdd{M,A}) where {M,O,A}
    LazyMulAdd(vadd_fast(getfield(a, :data), getfield(b, :data)), StaticInt{M}(), StaticInt{O}()+StaticInt{A}())
end

@inline add_indices(p::Ptr, a::LazyMulAdd{M,O,V}, b::LazyMulAdd{N,P,J}) where {M,O,V<:AbstractSIMDVector,N,P,J<:IntegerTypes} = (gep(p, b), a)
@inline add_indices(p::Ptr, b::LazyMulAdd{N,P,J}, a::LazyMulAdd{M,O,V}) where {M,O,V<:AbstractSIMDVector,N,P,J<:IntegerTypes} = (gep(p, b), a)
@inline add_indices(p::Ptr, a::LazyMulAdd{M,O,V}, b::LazyMulAdd{M,P,J}) where {M,O,V<:AbstractSIMDVector,P,J<:IntegerTypes} = (p, vadd_fast(a, b))
@inline add_indices(p::Ptr, b::LazyMulAdd{M,P,J}, a::LazyMulAdd{M,O,V}) where {M,O,V<:AbstractSIMDVector,P,J<:IntegerTypes} = (p, vadd_fast(a, b))

@inline add_indices(p::Ptr, a::AbstractSIMDVector, b::LazyMulAdd{M,O,I}) where {M,O,I<:IntegerTypes} = (gep(p, b), a)
@inline add_indices(p::Ptr, b::LazyMulAdd{M,O,I}, a::AbstractSIMDVector) where {M,O,I<:IntegerTypes} = (gep(p, b), a)
@inline function add_indices(p::Ptr, ::MM{W,X,StaticInt{A}}, a::LazyMulAdd{M,O,T}) where {M,O,T<:IntegerTypes,A,W,X}
    gep(p, a), MM{W,X}(StaticInt{A}())
end
@inline function add_indices(p::Ptr, a::LazyMulAdd{M,O,T}, ::MM{W,X,StaticInt{A}}) where {M,O,T<:IntegerTypes,A,W,X}
    gep(p, a), MM{W,X}(StaticInt{A}())
end

@generated function add_indices(p::Ptr, a::LazyMulAdd{M,O,MM{W,X,StaticInt{I}}}, b::LazyMulAdd{N,P,J}) where {M,O,W,X,I,N,P,J<:IntegerTypes}
    d, r = divrem(M, N)
    if iszero(r)
        quote
            $(Expr(:meta,:inline))
            p, VectorizationBase.LazyMulAdd{$N,$(I*M)}(MM{$W,$d}(getfield(b, :data)))
        end
    else
        quote
            $(Expr(:meta,:inline))
            gep(p, b), a
        end
    end
end
@inline add_indices(p::Ptr, b::LazyMulAdd{N,P,J}, a::LazyMulAdd{M,O,MM{W,X,StaticInt{I}}}) where {M,O,W,X,I,N,P,J<:IntegerTypes} = add_indices(p, a, b)
@generated function vadd_fast(a::LazyMulAdd{M,O,MM{W,X,StaticInt{I}}}, b::LazyMulAdd{N,P,J}) where {M,O,W,X,I,N,P,J<:IntegerTypes}
    d, r = divrem(M, N)
    if iszero(r)
        quote
            $(Expr(:meta,:inline))
            VectorizationBase.LazyMulAdd{$N,$(I*M)}(MM{$W,$d}(getfield(b, :data)))
        end
    else
        quote
            $(Expr(:meta,:inline))
            vadd_fast(a, _materialize(b))
        end
    end
end
@inline vadd_fast(b::LazyMulAdd{N,P,J}, a::LazyMulAdd{M,O,MM{W,X,StaticInt{I}}}) where {M,O,W,X,I,N,P,J<:IntegerTypes} = vadd_fast(a, b)

@inline vadd_fast(a::VecUnroll, b::LazyMulAdd) = VecUnroll(fmap(vadd_fast, getfield(a, :data), b))
@inline vadd_fast(b::LazyMulAdd, a::VecUnroll) = VecUnroll(fmap(vadd_fast, b, getfield(a, :data)))