Pkg.add(["Statistics", "Flux", "Symbolics"])

# Pkg.add( "ChainPlots" ) # Использовать осторожно, от версии к версии бывают нестыковки

using Flux, Statistics, Random
Random.seed!( 2 )           # Обеспечим управляемость процесса обучения

# Архитектура модели: два полносвязанных слоя с небольшим количеством нейронов в каждом.
model = Chain(
      Dense(2 => 3, tanh),
      Dense(3 => 2),        # Сколько классов, столько и выходов нейросети
      softmax )

# Генерация входных данных
inputs = rand( Float32, 2, 1000 );                                   # 2×1000 Matrix{Float32}
truth = [ xor(col[1]>0.5, col[2]>0.5) for col in eachcol(inputs) ];  # Vector{Bool} из 1000 элементов

# Сохраним на будущее прогноз "необученной" модели
probs1 = model( inputs );

# Подготовка данных и обучение
targets = Flux.onehotbatch( truth, [true, false] );          # Раскладываем выходную переменную на логиты и создаем загрузчик данных
data = Flux.DataLoader( (inputs, targets), batchsize=64, shuffle=true );
opt_state = Flux.setup( Adam( 0.01 ), model );               # Настройки процедуры оптимизации и конкретная функция потерь
loss(ỹ, y) = Flux.crossentropy( ỹ, y )
accuracy(ỹ, y) = mean( Flux.onecold( ỹ ) .== Flux.onecold( y ))
loss_history, accuracy_history = [], []                      # Производим обучение, записывая результаты
for i in 1:5000
    Flux.train!( model, data, opt_state) do m, x, y
        loss( m(x), y ) # Функция потерь - ошибка на каждом элементе датасета
    end
    push!( loss_history, loss( model(inputs), targets ) ) # Запомним значение функции потерь и точности прогноза
    push!( accuracy_history, accuracy( model(inputs), targets ) )
end

# Прогноз модели после обучения
probs2 = model( inputs );

# Выведем график, по которому можно оценить качество обучения
gr()
plot( [ loss_history, accuracy_history], size=(300,200), label=["loss" "accuracy"], leg=:right )

println( "Точность прогноза перед обучением: ", 100 * mean( (probs1[1,:] .> 0.5) .== truth ), "%" )
println( "Точность прогноза после обучения: ", 100 * mean( (probs2[1,:] .> 0.5) .== truth ), "%" )

# Выведем график исходных данных
p_true = scatter( inputs[1,:], inputs[2,:], zcolor=truth, title="Исходные данные" );
p_raw = scatter( inputs[1,:], inputs[2,:], zcolor=probs1[1,:], title="Прогноз до обучения" );
p_done = scatter( inputs[1,:], inputs[2,:], zcolor=probs2[1,:], title="После обучения" );

plot(p_true, p_raw, p_done, layout=(1,3), size=(700,200), titlefont=font(9), ms=3.5, legend=false, cbar=false )

# Сгенерируем новое изображение для размещения на лицевой стороне блока
# (успех зависит от стабильности текущей версии ChainPlots)

# using ChainPlots
# p = plot( model,
#     titlefontsize=10, size=(300,300),
#     xticks=:none, series_annotations="", markersize=8,
#     markercolor="white", markerstrokewidth=4, linewidth=1 )
# savefig( p, "$(@__DIR__)/neural_net_block_mask.png");

# Создадим код нейросети

using Symbolics
@variables x1 x2
s = model( [x1, x2] );

# Загрузим модель, если она еще не открыта на холсте
if "neural_classification" ∉ getfield.(engee.get_all_models(), :name)
    engee.load( "$(@__DIR__)/neural_classification.engee");
end

# Шаблон кода, который мы поместим в блок Engee Function
code_strings = """
struct Block <: AbstractCausalComponent; end

# У нейросети два выхода: s[1] и s[2]
nn(x1, x2) = ($(s[1]), $(s[2]))

# Вычислим выходы нейросети и вернем результат классификации: 0 или 1
function (c::Block)(t::Real, x1, x2)
    # "Вероятность" каждого из классов
    c1, c2 = nn(x1, x2)
    # Вычисляем выходное значение по результатам классификации
    # - если больше вероятность c1, то будет выбран класс true (1)
    # - если больше – вероятность c2, то будет выбран класс false (2)
    return (c1 > c2) ? 1 : 0
end
"""

# Какой блок модели должен содержать код нейросети?
block_address = "neural_classification/Engee Function"
engee.set_param!( block_address, "StepMethodCode" => code_strings)

# Сохраним модель после изменения 
engee.save( "neural_classification", "$(@__DIR__)/neural_classification.engee"; force = true )

model_data = engee.run( "neural_classification" );

# Подготовим выходные переменные
model_x1 = model_data["X1"].value;
model_x2 = model_data["X2"].value;
model_y = vec( hcat( model_data["Y"].value... ));

# Построим график
scatter( model_x1, model_x2, model_y, ms=2.5, msw=.5, leg=false, zcolor=model_y, c=:viridis,
         xlimits=(0,1), ylimits=(0,1), title="Прогноз от блока Engee Function", titlefont=font(10) )

# plot( model, size=(600, 350) )

data = Flux.DataLoader( (inputs, targets), batchsize=1 );
first( data )