using Plots, Random, Distributions

Random.seed!(13)

function GenerateSamples(n)
    d = Normal(0, 0.1)

    X = [(1.0/n*i) for i in 0:n]
    Y = [(sin(2*pi*x) + rand(d)) for x in X]

    X0 = [i*1/1000 for i in 0:1000]

    return X0, X, Y
end

function PlotData(X0, X, Y, Yhat=[])
    plot(X, Y, seriestype=:scatter, title="Curve Fitting", label="Training (x, y)")

    plot!(X0, [sin(2*pi*x) for x in X0], lw=2, label="True model")

    if length(Yhat) > 0
        plot!(X0, Yhat, seriestype=:line, lw=2, label="Fitted (x, y)")
    end
end

function PolyRegression(X, Y, m)
    n = length(X)
    A = ones(n, 1+m)
    for i in 1:n
        for p in 2:(m+1)
            A[i,p] = X[i]^(p-1)
        end
    end
    w = inv(A'*A)*A'*Y
    return w
end

function PredictPoly(X, w, m)
    Yhat = [sum([w[p]*x^(p-1) for p in 2:(m+1)]) for x in X]
    return Yhat
end

println("--------- START ---------")

X0, X, Y = GenerateSamples(10)

w = PolyRegression(X, Y, 9)
Yhat = PredictPoly(X0, w, 9)
PlotData(X0, X, Y, Yhat)