import Flux
using Flux
using Flux: params, flatten
import Zygote
using Statistics
import LinearAlgebra
import BSON
using BSON: @save
using BSON: @load
using Dates
using Serialization

function ResNetBlock(depth; filter_size=(3,3), momentum=0.1f0)
    pad = filter_size .÷ 2
    layers = Chain(
        Conv(filter_size, depth=>depth, pad=pad),
        BatchNorm(depth, relu; momentum=momentum),
        Conv(filter_size, depth=>depth, pad=pad),
        BatchNorm(depth; momentum=momentum))
    Chain(
        SkipConnection(layers, +),
        x -> relu.(x))
end

mutable struct Incremental <: Network
    config::MuZeroConfig
    saved_at::Union{Nothing, String}
    step::Integer
    test_mode::Bool
    force_cpu::Bool
    representation_network::Chain
    prediction_common_network::Chain
    prediction_policy_network::Chain
    prediction_value_network::Chain
    dynamics_hidden_state_network::Chain
    dynamics_reward_network::Chain
    function Incremental(config; test_mode=false, force_cpu=false)
        game_type = config.game_type
        hidden_size = config.hidden_state_size
        observation_resolution = config.observation_shape
        initial_filter_size = (3,3)
        input_depth = config.observation_shape[3]
        pad = initial_filter_size .÷ 2
        num_filters = config.num_filters
        flat_size = observation_resolution[1] * observation_resolution[2] * num_filters

        rep_layers = config.representation_resnet_layers
        representation_net = Chain(
            Conv(initial_filter_size, input_depth=>num_filters, pad=pad),
            BatchNorm(num_filters, relu, momentum=0.1f0),
            [ResNetBlock(num_filters) for _ in 1:rep_layers]...,
        )
        
        dynamics_layers = config.dynamics_resnet_layers
        dynamics_hidden_state_net = Chain(
            Conv(initial_filter_size, (num_filters+1)=>num_filters, pad=pad),
            BatchNorm(num_filters, relu, momentum=0.1f0),
            [ResNetBlock(num_filters) for _ in 1:dynamics_layers]...,
        )

        reward_layers = config.reward_fc_layers
        reward_net = Chain(
            Conv((1,1), num_filters=>num_filters),
            BatchNorm(num_filters, relu, momentum=0.1f0),
            flatten,
            [Dense(flat_size, flat_size, relu) for _ in 1:reward_layers]...,
            Dense(flat_size, full_support_size(config), relu),
            Dense(full_support_size(config), 1, tanh))

        prediction_common = Chain(
            ResNetBlock(num_filters))

        policy_layers = config.policy_fc_layers
        prediction_policy = Chain(
            Conv((1,1), num_filters=>num_filters),
            BatchNorm(num_filters, relu, momentum=0.1f0),
            flatten,
            [Dense(flat_size, flat_size, relu) for _ in 1:policy_layers]...,
            Dense(flat_size, action_space_size(game_type)))

        value_layers = config.value_fc_layers
        prediction_value = Chain(
            Conv((1,1), num_filters=>num_filters),
            BatchNorm(num_filters, relu, momentum=0.1f0),
            flatten,
            [Dense(flat_size, flat_size, relu) for _ in 1:value_layers]...,
            Dense(flat_size, full_support_size(config), relu),
            Dense(full_support_size(config), 1, tanh))

        device(net) = force_cpu ? cpu(net) : gpu(net)
        x = new(
            config, 
            nothing,
            0,
            test_mode,
            force_cpu, 
            device(representation_net),
            device(prediction_common),
            device(prediction_policy),
            device(prediction_value),
            device(dynamics_hidden_state_net),
            device(reward_net),
        )
        enable_test_mode_if_needed!(x)
        x
    end
end

function get_nets(x::Incremental)
    [getfield(x, n) for n in get_net_names(x)]
end

function enable_test_mode_if_needed!(n::Incremental)
    testmode!.(get_nets(n), n.test_mode)
end

zeropadder(padto) = (x) -> Matrix{Float32}(I, padto, size(x, 1)) * x
device(n::Incremental, x) = n.force_cpu ? cpu(x) : gpu(x)

function even_1_odd_n1(x)
    res = x / Float32(2)
    floor(res) ≈ res ? Float32(1) : Float32(-1)
end

function prediction(n::Incremental, hidden_state::AbstractArray{Float32})
    common = n.prediction_common_network(hidden_state)
    policy_logits = n.prediction_policy_network(common)
    state_values = n.prediction_value_network(common)
    policy_logits, state_values
end

function initial_inference(n::Incremental, games::AbstractArray{MuZeroGame})
    initial_inference(
        n, 
        current_image.(getgame.(games))
    )
end

initial_inference(n::Incremental, game::MuZeroGame) = initial_inference(n, game.game)
initial_inference(n::Incremental, game::AbstractEnv) = initial_inference(n, current_image(game))
initial_inference(n::Incremental, game::Vector{<:AbstractEnv}) = initial_inference(n, current_image.(game))
initial_inference(n::Incremental, image::AbstractImage) = initial_inference(n, tensorize(image))
initial_inference(n::Incremental, image::Vector{<:AbstractImage}) = initial_inference(n, tensorize(image))

function initial_inference(n::Incremental, s::AbstractEnvState)
    initial_inference(
        n, 
        get_image(s),
    )
end

#t = TTT(0)
#out = initial_inference(pttt, current_image(t)) |> SingleNetworkOutput

function initial_inference(n::Incremental, images::AbstractArray{Float32, 4})
    hidden_state = n.representation_network(device(n, images))
    policy_logits, value = prediction(n, hidden_state)
    reward = n.dynamics_reward_network(hidden_state)
    value, reward, policy_logits, hidden_state  
end   

function dynamics(n::Incremental, hidden_state::AbstractArray{Float32, 4}, actions::AbstractArray{Float32, 4})
    #example: size(actions) = (3,3,1,9)
    hidden_state_plus_actions = cat(hidden_state, actions; dims=3)
    next_hidden_state = n.dynamics_hidden_state_network(hidden_state_plus_actions)
    rewards = n.dynamics_reward_network(next_hidden_state)
    next_hidden_state, rewards
end

# out = initial_inference(pttt, current_image(t)) |> SingleNetworkOutput
# next_states, rewards = dynamics(pttt, hidden_state(out), [2])
# dynamics(pttt, images, [2,2])

function recurrent_inference(n::Incremental, hidden_states::Vector{<:AbstractArray{Float32, 3}}, actions::Vector{<:Integer})
    recurrent_inference(
        n,
        reduce((x,y)->cat(x,y; dims=4), Flux.unsqueeze.(hidden_states,4)),
        reduce(
            (x,y)->cat(x,y; dims=4), 
            Flux.unsqueeze.(
                map(a->get_action_image(n.config.game_type, a).channels, actions),
                4
            )
        ),
    )
end

function recurrent_inference(n::Incremental, hidden_state::AbstractArray{Float32, 3}, action::Integer)
    action_tensor = Flux.unsqueeze(get_action_image(n.config.game_type, action).channels, 4)
    recurrent_inference(n, Flux.unsqueeze(hidden_state, 4), action_tensor)
end

function recurrent_inference(n::Incremental, hidden_state::AbstractArray{Float32, 4}, action::AbstractArray{Float32, 4})
    next_hidden_state, reward = dynamics(n, hidden_state, action)
    policy_logits, value = prediction(n, next_hidden_state)
    value, reward, policy_logits, next_hidden_state  
end

get_weights(x::Incremental) = params(get_nets(x)...)

# function save!(n::Incremental, optimizer, path::String)
#     ensure_pathto_exists(path)
#     saved_at = string(now())
#     n.saved_at = saved_at
#     opt                     = cpu(optimizer)
#     representation          = cpu(n.representation_network)
#     prediction_common       = cpu(n.prediction_common_network)
#     prediction_policy       = cpu(n.prediction_policy_network)
#     prediction_value        = cpu(n.prediction_value_network)
#     dynamics_hidden_state   = cpu(n.dynamics_hidden_state_network)
#     dynamics_reward         = cpu(n.dynamics_reward_network)
#     @save expanduser(path) opt saved_at representation prediction_common prediction_policy prediction_value dynamics_hidden_state dynamics_reward
# end

function copy!(from::Incremental, to::Incremental)
    to.saved_at                         = from.saved_at
    to.step                             = from.step
    # @todo - need to make sure device actually *copies* for both cpu and gpu
    to.representation_network           = device(to, from.representation_network)
    to.prediction_common_network        = device(to, from.prediction_common_network)
    to.prediction_policy_network        = device(to, from.prediction_policy_network)
    to.prediction_value_network         = device(to, from.prediction_value_network)
    to.dynamics_hidden_state_network    = device(to, from.dynamics_hidden_state_network)
    to.dynamics_reward_network          = device(to, from.dynamics_reward_network)
    enable_test_mode_if_needed!(to)
    return
end

function save!(n::Incremental, path::String)
    ensure_pathto_exists(path)
    #     d = Dict()
    #     for x in get_net_names(n)
    #         d[x] = params(getfield(n, x))
    #     end
    #     @show keys(d)
    #     BSON.bson(expanduser(path), d)
    saved_at = string(now())
    n.saved_at = saved_at
    step = n.step
    representation          = cpu(n.representation_network)
    prediction_common       = cpu(n.prediction_common_network)
    prediction_policy       = cpu(n.prediction_policy_network)
    prediction_value        = cpu(n.prediction_value_network)
    dynamics_hidden_state   = cpu(n.dynamics_hidden_state_network)
    dynamics_reward         = cpu(n.dynamics_reward_network)
    @save expanduser(path) saved_at step representation prediction_common prediction_policy prediction_value dynamics_hidden_state dynamics_reward
    dir = join(splitpath(expanduser(path))[1:(end-1)], "/")
    serialize("$dir/step.bin", step)
    serialize("$dir/saved_at.bin", saved_at)
    serialize("$dir/representation.bin", representation)
    serialize("$dir/prediction_common.bin", prediction_common)
    serialize("$dir/prediction_policy.bin", prediction_policy)
    serialize("$dir/prediction_value.bin", prediction_value)
    serialize("$dir/dynamics_hidden_state.bin", dynamics_hidden_state)
    serialize("$dir/dynamics_reward.bin", dynamics_reward)
end

get_net_names(n::Incremental) = get_net_names(typeof(n))

function get_net_names(::Type{Incremental})
    [
        :representation_network,
        :prediction_common_network,
        :prediction_policy_network,
        :prediction_value_network,
        :dynamics_reward_network,
        :dynamics_hidden_state_network,
    ]
end

# function load!(n::Incremental, d::Dict)
#     n.saved_at                      = d[:saved_at]
#     n.representation_network        = device(n, d[:representation])
#     n.prediction_common_network     = device(n, d[:prediction_common])
#     n.prediction_policy_network     = device(n, d[:prediction_policy])
#     n.prediction_value_network      = device(n, d[:prediction_value])
#     n.dynamics_reward_network       = device(n, d[:dynamics_reward])
#     n.dynamics_hidden_state_network = device(n, d[:dynamics_hidden_state])
#     enable_test_mode_if_needed!(n)
#     nothing
# end

function load!(n::Incremental, path::String)
    # load!(
    #     n,
    #     BSON.load(expanduser(path), @__MODULE__)
    # )
    dir = join(splitpath(expanduser(path))[1:(end-1)], "/")
    n.step                          = device(n, deserialize("$dir/step.bin"))
    n.saved_at                      = device(n, deserialize("$dir/saved_at.bin"))
    n.representation_network        = device(n, deserialize("$dir/representation.bin"))
    n.prediction_common_network     = device(n, deserialize("$dir/prediction_common.bin"))
    n.prediction_policy_network     = device(n, deserialize("$dir/prediction_policy.bin"))
    n.prediction_value_network      = device(n, deserialize("$dir/prediction_value.bin"))
    n.dynamics_hidden_state_network = device(n, deserialize("$dir/dynamics_hidden_state.bin"))
    n.dynamics_reward_network       = device(n, deserialize("$dir/dynamics_reward.bin"))
    enable_test_mode_if_needed!(n)
end

se(ŷ, y) = (ŷ .- y) .^ 2

function sum_se_scale_gradient(ŷ, y, n, scale)
    squared_error = se(ŷ, y)
    @assert size(squared_error) == size(scale)
    sum(scale_gradient.(squared_error, scale)) / Float32(n)
end

logitcrossentropy(ŷ, y; dims = 1) = .-sum(y .* logsoftmax(ŷ; dims = dims); dims = dims)
function sum_logitcrossentropy_scale_gradient(ŷ, y, n, scale)
    logit_cross_entropy = reshape(logitcrossentropy(ŷ, y), :)
    sum(scale_gradient.(logit_cross_entropy, scale)) / Float32(n)
end

function my_mse(ŷ, y, n, gradient_scale)
    sum(
        #scale_gradient.(
            (mse(ŷ, y; agg=collect) / Float32(n))
            #gradient_scale
        #)
    )
end
function my_crossentropy(ŷ, y, n, gradient_scale)
    sum(
        #scale_gradient.(
            (crossentropy(ŷ, y; agg=collect) / Float32(n))#,
            #gradient_scale
        #)
    )
end
#     scale_gradient.(
#         crossentropy(preds, targets; agg=collect) ./ Float32(size(targets,2)),
#         gradient_scale
#     ) |> sum
# end

function calc_target_loss(preds::Tuple, target, num_targets, gradient_scale=gpu(ones(Float32, 1, size(target, 2))))
    value_preds, reward_preds, policy_logits_preds, _ = preds

    target_values = view(target, 1:1, :)
    value_loss = sum_se_scale_gradient(value_preds, target_values, num_targets, gradient_scale)

    target_rewards = view(target, 2:2, :)
    reward_loss = sum_se_scale_gradient(reward_preds, target_rewards, num_targets, gradient_scale)

    target_policies = view(target, 3:size(target, 1), :)
    #policy_preds = softmax(policy_logits_preds)
    policy_loss = sum_logitcrossentropy_scale_gradient(policy_logits_preds, target_policies, num_targets, gradient_scale) #my_crossentropy

    #policy_preds = softmax(policy_logits_preds)
    #policy_loss = my_crossentropy(policy_preds, target_policies, num_targets, gradient_scale)

    TargetLoss(policy_loss, reward_loss, value_loss)
end

function calc_initial_loss(network::Incremental, batch::TrainingBatchTensor)
    num_targets = calc_num_initial_inference_targets(batch)
    net_output = initial_inference(network, images(batch))
    loss = calc_target_loss(net_output, image_targets(batch), num_targets)
    num_games = size(image_targets(batch), 2)
    (loss, net_output[4], 1:num_games)
end

function getloss(network::Incremental, batch::TrainingBatchTensor)
    #target_losses = calc_initial_loss(network, batch)[1]
    calc_recurrent_loss = recurrent_loss_fn(network, batch)
    target_loss, _, _ = reduce(calc_recurrent_loss, action_indices(batch); init=calc_initial_loss(network, batch))  
    #regularization_loss = weight_decay * sum(l2_loss(p) for p in all_params)
    Loss(target_loss, Float32(0))
end

# @todo make action space size a config param
function recurrent_loss_fn(network::Incremental, batch::TrainingBatchTensor)
    action_targets = batch.action_targets
    actions = batch.actions
    num_targets = calc_num_recurrent_inference_targets(batch)

    return (prev_output, action_idx) -> begin

        prev_loss_sum, prev_hidden_states, prev_game_indices = prev_output

        action_idx_filter = action_targets[2,:] .== Float32(action_idx)
        
        game_indices = Zygote.dropgrad(view(action_targets, 1, action_idx_filter))

        remaining_games_filter = index_in(cpu(game_indices), cpu(prev_game_indices))
        remaining_hidden_states = view(prev_hidden_states, :, :, :, remaining_games_filter)

        @assert size(actions, 4) == size(action_targets, 2)
        next_actions = view(actions, :, :, :, action_idx_filter)
        
        # @todo make this more robust. Perhaps
        # put policy indices at the end of the vector.
        gradient_idx = size(action_targets, 1)
        curr_targets = view(action_targets, 4:(gradient_idx-1), action_idx_filter)
        
        net_output = recurrent_inference(network, remaining_hidden_states, next_actions)

        gradient_scale = view(action_targets, gradient_idx:gradient_idx, action_idx_filter)
        curr_loss = calc_target_loss(net_output, curr_targets, num_targets, gradient_scale)
        hidden_states = scale_gradient(net_output[4], Float32(0.5))
        
        (prev_loss_sum + curr_loss, hidden_states, game_indices)
    end
end
