import Pkg
try
    using Plots
    using Flux
    using Flux: train!
catch
    Pkg.add("Plots")
    Pkg.add("Flux")

    using Flux
    using Flux: train!
    using Plots
end

using Statistics
using Random
using Base.Iterators: flatten

gr()

"""
Das Programm kann mit
julia -i WitjesHahnZettel5.5.jl
gestartet werden. Dazu muss https://julialang.org/ (am besten v1.10) installiert sein.

Der interaktive Teil wird mit
julia
include("WitjesHahnZettel5.5.jl")
inter()
gestartet. Fuer den ML Teil kann f() ausgefuehrt werden.

Beispielhafte Programmausfuehrung:
```Bitte gebe Punkte ein, wobei die Komponenten mit der Leertaste getrennt sind. Dann bestaetige mit Enter. Dabei durfen keine zwei Punkte den selben ersten Wert haben. Mit Shift und Enter kann ein weiterer Punkt hinzugefuegt werden.
3 5
4 7
5 -8
10 3
12 0

┌ Info: Eingelesene Punkte
│   points =
│    5-element Vector{Tuple{Float32, Float32}}:
│     (3.0, 5.0)
│     (4.0, 7.0)
│     (5.0, -8.0)
│     (10.0, 3.0)
└     (12.0, 0.0)
┌ Info: Berechnete Koeffizienten des Polynoms
│   coeffs =
│    5-element Vector{Float32}:
│     -336.9761904761904
│      258.4361111111111
│      -65.99146825396825
│        6.630555555555556
└       -0.22757936507936508```

Dabei wird dann auch ein Bild der Grafik erstellt und als svg gespeichert.
"""


# Helping types for safety
const MaybePoint = Union{Tuple{Float32,Float32},Nothing}
const PointStack = Vector{Tuple{Float32,Float32}}

"""
Read and save the inputs from the user.
When entering nothing the input is finished.
Otherwise it continues forever.

This does use stdin and stdout.
"""
function get_points()::PointStack
    points::PointStack = Vector()
    while true
        point::MaybePoint = read_point()
        if point == nothing
            return points
        end
        push!(points, point)
    end
end

function read_point()::MaybePoint
    line = strip(readline())
    parts = split(line)
    length(parts) == 2 || return nothing
    vals = try
        parse.(Float32, parts)
    catch
        return nothing
    end
    return (vals[1], vals[2])
end


# Simple right shift of any Vector
shift_right(v::Vector) = [0; v[1:end-1]]

# https://en.wikipedia.org/wiki/Neville%27s_algorithm
# https://math.stackexchange.com/questions/997348/calculating-coefficients-of-interpolating-polynomial-using-nevilles-algorithm
function get_coefficients(points, i, j)::Vector
    i > 0 && j > 0 || error("Cannot evaluate at negative index!")
    j >= i || error("The second index must be the bigger one!")
    coeffs = zeros(length(points))
    if i == j
        yi = points[i][2]
        coeffs[1] = yi
        return coeffs
    end
    xi = points[i][1]
    xj = points[j][1]
    return (xj * get_coefficients(points, i, j - 1) - shift_right(get_coefficients(points, i, j - 1)) + shift_right(get_coefficients(points, i + 1, j)) - xi * get_coefficients(points, i + 1, j)) / (xj - xi)
end

get_coefficients(points) = get_coefficients(points, 1, length(points))

function evaluate_coefficients(x, coeffs)::Float32
    res::Float32 = 0
    for (i, coeff) in enumerate(coeffs)
        res += coeff * x^(i - 1)
    end
    return res
end

function plot_points_simple(points, coeffs)
    OFFSET = 0

    poly(x) = evaluate_coefficients(x, coeffs)

    # Plotten des Polynoms und der Punkte
    graph_plot = scatter(points, title="Polynom")
    plot!(poly, minimum(first, points) - OFFSET, maximum(first, points) + OFFSET)
    bar_plot = bar(coeffs, title="Koeffizienten")
    plot(graph_plot, bar_plot, legend=false)

    savefig("OutputPlot.svg")
    gui()
end

function plot_points_ml(points, exact_coeffs, ml_coeffs, loss_history)
    exact_poly(x) = evaluate_coefficients(x, exact_coeffs)
    ml_poly(x) = evaluate_coefficients(x, ml_coeffs)

    # Plotten des Polynoms und der Punkte
    graph_plot = scatter(points, title="Polynom")
    plot!(exact_poly, minimum(first, points), maximum(first, points))
    plot!(ml_poly, minimum(first, points), maximum(first, points))

    loss_plot = plot(loss_history, title="Training Loss")
    exactbar_plot = bar(exact_coeffs, title="Algo Koeffizienten")
    mlbar_plot = bar(ml_coeffs, title="Ml Koeffizienten")

    plot(graph_plot, loss_plot, exactbar_plot, mlbar_plot, legend=false)
    # plot(graph_plot, loss_plot, legend=false)

    savefig("OutputPlotMl.svg")
    gui()
end


"""
Main entrypoint to run the program as a interactive script
"""
function inter()
    println("Bitte gebe Punkte ein, wobei die Komponenten mit der Leertaste getrennt sind. Dann bestaetige mit Enter. Dabei durfen keine zwei Punkte den selben ersten Wert haben. Mit Shift und Enter kann ein weiterer Punkt hinzugefuegt werden.")

    points::PointStack = get_points()
    @info "Eingelesene Punkte" points
    coeffs = get_coefficients(points)
    @info "Berechnete Koeffizienten" coeffs

    # Here first, points is short for map(first, points)
    length(points) > 0 || error("Must enter at least one point!")
    length(unique(first, points)) == length(points) || error("Points not valid!")

    plot_points_simple(points, coeffs)
end


function generate_training_data()
    data = []
    for i in 1:NSAMPLES
        points = [(rand(), rand()) for _ in 1:NPOINTS]
        coeffs = get_coefficients(points) # N Vector
        push!(data, (collect(flatten(points)), coeffs)) # Tuple of 2N input and N solution
    end

    X = hcat(first.(data)...)
    Y = hcat(last.(data)...)

    Y_mean = mean(Y, dims=2)
    Y_std = std(Y, dims=2) .+ eps()

    Y = (Y .- Y_mean) ./ Y_std

    (X, Y, Y_mean, Y_std)
end


p = 2  # power to emphasize larger coefficients

loss(model, x, y) = mean(abs2.(model(x) .- y))
# loss(model, x, y) = mean(((model(x) .- y) .* (abs.(y) .^ p)).^2)
# function loss(model, x, y)
#     y_pred = model(x)
#     N, B = size(y)                       
#     powers = reshape(1:N, N, 1)          
#     errors = abs.(y_pred .- y) .^ powers     
#     return mean(errors)                  
# end
# function loss(model, x, y)
#     yht = model(x)
#     N = length(y[:, 1])
#     weights = collect(map(x -> x^2, 1:N)) ./ N^2
#     return mean(weights .* abs2.(yht .- y))
# end

function train_model(Xt, Yt, Xv, Yv)
    N = size(Yt, 1)

    model = Chain(
        Dense(2N, 16, relu),  # hidden layer
        Dense(16, 64, relu),  # hidden layer
        Dense(64, 32, relu),  # hidden layer
        Dense(32, 16, relu),  # hidden layer
        Dense(16, N)         # output layer
    )

    opti = NAdam(0.001, (0.9, 0.999), 1e-8)
    opt = Flux.setup(opti, model)

    loss_history = zeros(EPOCHS)
    for epoch in 1:EPOCHS
        train!(loss, model, [(Xt, Yt)], opt)  # Pass as a tuple of matrices
        losss = loss(model, Xv, Yv)
        loss_history[epoch] = losss
        if losss < BREAKUP
            break
        end
        # println("=> ", round(epoch/EPOCHS, digits=2))
        epoch % 100 == 0 && print("\rEpoch $(epoch/EPOCHS), Loss: $(round(losss, digits=6))")
    end
    println()

    loss_history = filter(x -> x != 0, loss_history)

    return (model, loss_history)
end

# TODO: Implement batching of the training data
const NPOINTS = 4
const NSAMPLES = 8
const EPOCHS = 100
const CHEAT = true
const SIMPLE = false
const BREAKUP = 0.01

# TODO: Check functionality
function f()
    println("Starting ML. Be patient.")

    X, Y, Y_mean, Y_std = generate_training_data()

    N = size(X, 1)
    idx = Random.shuffle(1:N)
    n_train = Int(floor(0.8N))
    train_idx = idx[1:n_train]
    val_idx = idx[n_train+1:end]

    Xt = X[:, train_idx]
    Yt = Y[:, train_idx]
    Xv = X[:, val_idx]
    Yv = Y[:, val_idx]

    model, history = train_model(Xt, Yt, Xv, Yv)
    @info model

    points::PointStack = []
    index = rand(1:length(X[1, :]))
    if CHEAT
        for (i, x) in enumerate(X[:, 1])
            if i % 2 == 1
                push!(points, (x, X[:, 1][i+1]))
            end
        end
    else
        if SIMPLE
            points = [(0.2, 0.8), (0.5, 0.2), (0.8, 0.8)]

        else
            points = [(rand(), rand()) for _ in 1:NPOINTS]
        end
    end

    xtest = vcat(collect(flatten(points))...)
    exact_coeffs = get_coefficients(points)
    ml_coeffs = model(xtest) .* Y_std .+ Y_mean

    plot_points_ml(points, exact_coeffs, ml_coeffs, history)
end

"""
Erkenntnisse:
- Manchmal invertiert das Modell die Koeffizienten
- Wie kann der Fehler so klein sein, wenn das Modell nicht auswendig lernt und auch nicht gut ist
- Kleine Fehler in hohen Koeffizienten machen riesige Fehler im Plot
- Wie kann die Fehlerfunktion besser gemacht werden?
- Manchmal stimmt die Form aber die Koeffizienten sind um ein verschoben
"""

# Start the ML part
# f()

# Start the interactive part
# inter()

println("Use f() to start ml and inter() to start interactive mode.")