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 LinearRegression(X, Y)
    n = length(X)
    A = ones(n, 2)
    for i in 1:n
        A[i,2] = X[i]
    end
    println(A)
    w = inv(A'*A)*A'*Y
    return w
end

function PredictLinear(X, w)
    Yhat = [(w[1] + w[2]*x) for x in X]
    return Yhat
end


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

X0, X, Y = GenerateSamples(10)

w = LinearRegression(X, Y)
Yhat = PredictLinear(X0, w)
PlotData(X0, X, Y, Yhat)