#=
This code reads a FOO.jld generated by isingED.jl, generates N random
spin chain measurements in the computational basis, and saves them to 
a text file FOO.csv.
=#

using ArgParse, HDF5, JLD

function parse_commandline()
    s = ArgParseSettings()
    @add_arg_table s begin
      "--infilename", "-i"
        help =  "Path to the .jld where the wavefunction is stored."
        arg_type = AbstractString
        required = true
      "-N"
        help = "Number of configurations to generate."
        arg_type = Int
        default = 20
      "--overwrite"
        help = "If infilename.out already exists and this flag is set it will
                be overwritten (by default it will be appended to)."
        action = :store_true

    end
    return parse_args(s)
end

"""
Generates a sample spin chain measurement from the wavefunction wf.
The idea is to simulate measuring
each site in sequence. After each measurement, the wavefunction is projected
onto the appropriate Hilbert subspace according to the Born rule.
"""
function measure(wf::AbstractArray)
  #rng = MersenneTwister()
  Nsites = convert(Int64, log2(length(wf)))
  measurement = zeros(Int, Nsites)

  #shape = tuple(collect(repeated(2,Nsites))...)
  for n = 1:Nsites
    shape=(2^(Nsites-n), 2)
    wf = reshape(wf, shape) 
    #println(abs2(wf[:,1]))
    P0 = sum(abs2.(wf[:,1]))
    @assert P0>=0 && P0<(1.0+1e-12) "P0 was $P0 which is not a probability"
    roll = rand()
    if roll < P0
      outcome = 0
      wf = wf[:,1] / sqrt(P0) #this is the Born rule
    else
      outcome = 1
      wf = wf[:,2] / sqrt(1-P0)
    end
    measurement[Nsites-n+1] = outcome
  end
  return measurement
end

"""
***UNIMPLEMENTED***
Transforms the wavefunction so that computational-basis measurements upon
it are equivalent to measurements of the given string of Pauli operators
(0 -> I, 1->X, 2->Y, 3->Z).

For now this is just a null-op.
"""
function transformtopauli(wavefunction::AbstractArray, paulis::AbstractVector)
    return wavefunction
end

function writetofile(f, array)
    N = length(array)
    for i = 1:N
      write(f,  string(array[i]))
      if i!=N
        write(f, ",")
      else
        write(f, ";")
      end
    end
end


"""
Generates N measurements from the wavefunction and saves them to disk.
The overwrite flag determines whether the output file is overwritten or
appended to.
"""
function savemeasurements(wavefunction::AbstractArray, 
                          outfilename::AbstractString, N::Int, overwrite::Bool)
  if overwrite
    mode = "w"
  else
    mode = "a"
  end
  open(outfilename, mode) do f
    for n = 1:N
      Nsites = convert(Int64, log2(length(wavefunction)))
      paulis = [rand(0:3) for j in 1:Nsites] 
      pauliwf = transformtopauli(wavefunction, paulis)
      measurement = measure(pauliwf)
      println("Pauli ops: ", paulis)
      writetofile(f, paulis)
      println("Outcome: ", measurement)
      writetofile(f, measurement)
      write(f, "\n")
    end
  end
end


function main()
  parsed_args = parse_commandline()
  println("Hi!! Generating field configurations from:")
  for (arg, val) in parsed_args
    println("   $arg    =>    $val")
  end
  infilename = parsed_args["infilename"]
  wavefunction = load(infilename, "wavefunction")
  outfilename = infilename[1:end-3]*string("csv")
  savemeasurements(wavefunction, outfilename, parsed_args["N"], 
                   parsed_args["overwrite"])
  println("Done.")

end
main()
