using PyPlot




function simulate(N::Int64, D::Int64, simCounter::Int64)

  E1 = zeros(Float64, N)
  Eq1 = zeros(Float64, N)
  D1 = zeros(Float64, N)

  E2 = zeros(Float64, N)
  Eq2 = zeros(Float64, N)
  D2 = zeros(Float64, N)

  for n in 2:N
    for i in 1:simCounter
      urns = zeros(Int64, n)
      # First process
      for i in 1:n
        ballToUrn = rand(1:n)
        urns[ballToUrn] += 1
      end
      E1[n] += maximum(urns)
      Eq1[n] += (maximum(urns)^2)

      # Second process
      urns = zeros(Int64, n)
      for i in 1:n
        idx = getUrnIdx(urns, D, n)
        urns[idx] += 1
      end
      E2[n] += maximum(urns)
      Eq2[n] += (maximum(urns)^2)
    end
    E1[n] /= simCounter
    Eq1[n] /= simCounter
    D1[n] = sqrt(Eq1[n] - (E1[n]^2))

    E2[n] /= simCounter
    Eq2[n] /= simCounter
    D2[n] = sqrt(Eq2[n] - (E2[n]^2))
  end
  #printDiagram(E1,D1,1)
  printDiagram(E2,D2,2)
end


function printDiagram(E::Array, Std::Array, case::Int64)
  name = ""
  title = ""
  D1 = copy(Std)
  D2 = copy(Std)
  for i in 1:length(Std)
    D1[i] = E[i] + sqrt(2)*Std[i]
    D2[i] = E[i] - sqrt(2)*Std[i]
  end
  if case == 1
    title = "First task"
    name = "1"
  elseif case == 2
    title = "Second task"
    name = "2"
  elseif case == 3
    title = "Third task"
    name = "3"
  end
  x = 1:N

  fig, ax = subplots()
  suptitle(title)
  ax[:plot](x, E, linewidth=2, alpha=0.6, label="Expected value")
  ax[:plot](x, D1, linewidth=2, alpha=0.6, label="(+)Deviation", linestyle="--", color="red")
  ax[:plot](x, D2, linewidth=2, alpha=0.6, label="(-)Deviation", linestyle="--", color="red")
  ax[:legend]()
  fileName = string("lab5_$name.png")
  savefig(fileName)
end

function getUrnIdx(urns::Array, D::Int64, N::Int64)
  group = zeros(Int64, 0)
  values = zeros(Int64, D)

  while length(group) != D
    idx = rand(1:N)
    if !in(idx, group)
      push!(group, idx)
    end
  end

  for i in 1:length(group)
    values[i] = urns[group[i]]
  end

  minIdx = indmin(values)
  min = values[minIdx]

  indices = zeros(Int64, 0)
  push!(indices, group[minIdx])
  splice!(values, minIdx)
  splice!(group, minIdx)

  while length(values) > 0 && min == values[indmin(values)]
    push!(indices, group[indmin(values)])
    splice!(group, indmin(values))
    splice!(values, indmin(values))
  end

  urnIdx = indices[1]
  if length(indices) > 1
    urnIdx = indices[rand(1:length(indices))]
  end

  return urnIdx
end

println("START")
D = 3
N = 100
numberOfSimulations = 1000
simulate(N, D, numberOfSimulations)
println("END")