using ScikitLearn
@sk_import linear_model: LogisticRegression

struct Example
  featureVector::Array{Float64,1}
  outputVector::Array{Float64,1}
end

Example(x1::Float64, x2::Float64, y::String) = Example(
  [1.0, (x1 - 50.0) / 50.0, (x2 - 50.0) / 50.0, ((x1 - 50.0) / 50.0)^2, ((x2 - 50.0) / 50.0)^2],
  if y == "red"
    [1.0, 0.0, 0.0]
  elseif y == "green"
    [0.0, 1.0, 0.0]
  elseif y == "blue"
    [0.0, 0.0, 1.0]
  end
)

Example(x1::Int64, x2::Int64, y::String) = Example(Float64(x1), Float64(x2), y)

training_examples = [
    Example(6,  48, "red"),
    Example(8,  11, "red"),
    Example(15,  1, "red"),
    Example(16, 19, "red"),
    Example(22, 40, "red"),
    Example(25, 33, "red"),
    Example(30, 49, "red"),
    Example(31, 44, "red"),
    Example(32, 29, "red"),
    Example(38,  2, "red"),
    Example(53, 63, "green"),
    Example(50, 55, "green"),
    Example(90, 70, "green"),
    Example(69, 59, "green"),
    Example(76, 28, "green"),
    Example(86, 33, "green"),
    Example(51, 55, "green"),
    Example(88, 57, "green"),
    Example(85, 54, "green"),
    Example(60, 62, "green"),
    Example(14, 66, "blue"),
    Example(13, 76, "blue"),
    Example(40, 82, "blue"),
    Example(10, 98, "blue"),
    Example( 6, 51, "blue"),
    Example(47, 62, "blue"),
    Example(40, 85, "blue"),
    Example(11, 59, "blue"),
    Example(26, 95, "blue"),
    Example(20, 88, "blue")
]
red = LogisticRegression()
fit!(red, [ex.featureVector for ex in training_examples], [1 - ex.outputVector[1] for ex in training_examples])
println("Red: $(join(["theta[$(i - 1)] = $(@sprintf("%0.2f", red[:coef_][i]))" for i = 1:length(red[:coef_])], ", "))")
green = LogisticRegression()
fit!(green, [ex.featureVector for ex in training_examples], [1 - ex.outputVector[2] for ex in training_examples])
println("Green: $(join(["theta[$(i - 1)] = $(@sprintf("%0.2f", green[:coef_][i]))" for i = 1:length(green[:coef_])], ", "))")
blue = LogisticRegression()
fit!(blue, [ex.featureVector for ex in training_examples], [1 - ex.outputVector[3] for ex in training_examples])
println("Blue: $(join(["theta[$(i - 1)] = $(@sprintf("%0.2f", blue[:coef_][i]))" for i = 1:length(blue[:coef_])], ", "))")