struct Example
  featureVector::Array{Array{Float64,1},1}
  outputVector::Array{Array{Float64,1},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)

function Model(theta, example::Example)
  [1.0 ./ s for s in 1 .+ [e .^ p for p in theta' * example.featureVector]]
end

function Cost(theta, examples::Array{Example, 1})
  function CostOfExample(theta, example::Example)
    m = Model(theta, example)
    results = []
    for (x_i, y_i) in zip(m, example.outputVector)
      if y_i[1] == 1.0
        result = log(x_i[1])
      else
        result = log(1 - x_i[1])
      end
      if result == -Inf
        result = -1e20
      end
      push!(results, result)
    end
    results
  end

  function SumCostOverExamples(theta, examples::Array{Example, 1})
    sum([CostOfExample(theta, ex) for ex in examples])
  end

  -1.0 / length(examples) * SumCostOverExamples(theta, examples)
end

function DCostDTheta(theta, featureIndex::Int64, classIndex::Int64, examples::Array{Example, 1})
  terms = [(Model(theta, ex)[classIndex] - ex.outputVector[classIndex])[1] * ex.featureVector[featureIndex][1] for ex in examples]
  sum(terms)[1]
end

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")
]

# theta has a row for each feature and a column for each class
theta_shape = (
  length(training_examples[1].featureVector),
  length(training_examples[1].outputVector)
)
theta = zeros(theta_shape)

# currentCost is a row vector with a column for each class
currentCost = Cost(theta, training_examples)

println("Starting cost: $(sum(currentCost))")
learningRate = 0.5

adjustments = ones(theta_shape)
iterCount = 0

while sum(sum(abs.(adjustments))) > 1e-2
  for featureIndex = 1:length(training_examples[1].featureVector)
    for classIndex = 1:length(training_examples[1].outputVector)
      gradient = DCostDTheta(theta, featureIndex, classIndex, training_examples)
      adjustments[(classIndex - 1) * length(training_examples[1].featureVector) + featureIndex] = learningRate * gradient
    end
  end
  theta += adjustments
  iterCount += 1
  currentCost = Cost(theta, training_examples)
  if (0 == iterCount % 1000)
    println("Cost after $(iterCount) iterations: $(sum(currentCost))")
  end
end
println("Red: $(join(["theta[$(i-1)] = $(@sprintf("%.2f", theta[i]))" for i = 1:length(training_examples[1].featureVector)], ", "))")
println("Green: $(join(["theta[$(i-1)] = $(@sprintf("%.2f", theta[5 + i]))" for i = 1:length(training_examples[1].featureVector)], ", "))")
println("Blue: $(join(["theta[$(i-1)] = $(@sprintf("%.2f", theta[10 + i]))" for i = 1:length(training_examples[1].featureVector)], ", "))")
println("in $(iterCount) iterations")
println("Final cost = $(sum(Cost(theta, training_examples)))")