#
#  rquantum.rb
#    Date: 2008/07/13
#    by Tomasz Rekawek <newton@martel.pl>
#    http://newton.net.pl/
#
# == Overview
#
# It's Ruby implementation of Quantum Computer simulator. With rquantum
# you can easily create programs for quantum computers and test it. For example,
# program:
#
#  require 'rquantum'
#  q = QuantumReg.new 2
#  q.reset
#  q.not(0)
#  q.hadamard(1)
#  q.controlled_not(1,0)
#  q.pr
#
# will return:
#
#  0.707106781186547+0.0i	|01>
#  0.707106781186547+0.0i	|10>
#
# == Download
# 
# rquantum.rb[link:../rquantum.rb]
#

require 'complex'
require 'matrix'
require 'set'

class String #:nodoc: all
  #
  # Converts the CamelFormattedString to underscore_formated_string.
  # Copied from Ruby on Rails.
  #
  def underscore
    gsub(/::/, '/').
      gsub(/([A-Z]+)([A-Z][a-z])/,'\1_\2').
      gsub(/([a-z\d])([A-Z])/,'\1_\2').
      tr("-", "_").
      downcase
  end
end

class Vector #:nodoc: all
  #
  # Return the tensor product of the vector and the argument.
  #
  def tensor_product v
    n = size
    m = v.size
    w = Array.new(n*m)
    0.upto(n-1) { |i|
      0.upto(m-1) { |j|
	w[n*i+j] = self[i] * v[j]
      }
    }
    Vector.elements(w)
  end
end

#
# Module contains classes that simulates basic quantum gates and the matrices
# that representing these gates. Each class which simulate quantum gate has to
# have <tt>compute</tt> method which takes one parameter - vector with two complex
# numbers - two qubit amplitudes. If some gate needs additional parameters,
# then the initialize method takes them, and PARAMS constant is set to true.
# Parameters has to been passed in array.
#
# These classes are used to generating methods operating on Qubit class and
# QuantumReg class - methods have the same names as classes, but underscored
# (so not PauliY but pauli_y).
#
module Gates
  # 
  # Reverse amplitudes of qubit.
  # 
  #  q = Qubit.new
  #  q.reset
  #  g = Gates::Not.new
  #  q.to_vector		# -> Vector[Complex(1, 0), Complex(0, 0)]
  #  g.compute q.to_vector	# -> Vector[Complex(0, 0), Complex(1, 0)]
  #
  class Not
    PARAMS = false # :nodoc:
    # 
    # Takes vector with two complex numbers (amplitudes), perform Not operation 
    # and returns new vector.
    #
    def compute v
      Matrices::NOT * v
    end
  end

  #
  # Hadamarad gate - puts qubit into superposition state.
  #
  #  q = Qubit.new
  #  q.reset
  #  g = Gates::Hadamard.new
  #  q.to_vector		# -> Vector[Complex(1, 0), Complex(0, 0)]
  #  g.compute q.to_vector	# -> Vector[Complex(0.707106781186547, 0.0), Complex(0.707106781186547, 0.0)]
  #
  class Hadamard
   PARAMS = false # :nodoc:
    # 
    # Takes vector with two complex numbers (amplitudes), perform Hadamard gate operation
    # and returns new vector.
    #
    def compute vector
      Matrices::H * vector
    end
  end

  #
  # PauliX gate - the same function as Not gate.
  #
  #  q = Qubit.new
  #  q.reset
  #  g = Gates::PauliX.new
  #  q.to_vector		# -> Vector[Complex(1, 0), Complex(0, 0)]
  #  g.compute q.to_vector	# -> Vector[Complex(0, 0), Complex(1, 0)]
  #
  class PauliX
    PARAMS = false # :nodoc:
    # 
    # Takes vector with two complex numbers (amplitudes), perform Not operation
    # and returns new vector.
    #
    def compute vector
      Matrices::NOT * vector
    end
  end

  #
  # PauliY gate - composition of PauliZ and Not gates. It swaps the amplitudes
  # and change the sign of the second one (beta).
  #
  #  q = Qubit.new
  #  q.reset
  #  g = Gates::PauliY.new
  #  q.to_vector		# -> Vector[Complex(1, 0), Complex(0, 0)]
  #  g.compute q.to_vector	# -> Vector[Complex(0, 0), Complex(-1, 0)]
  #
  class PauliY
    PARAMS = false # :nodoc:
    # 
    # Takes vector with two complex numbers (amplitudes), perform PauliY operation
    # and returns new vector.
    #
    def compute vector
      Matrices::Y * vector
    end
  end

  # 
  # PauliZ gate - changes the sign of the second amplitude.
  #
  #  q = Qubit.new
  #  q.reset
  #  q.not
  #  g = Gates::PauliY.new
  #  q.to_vector		# -> Vector[Complex(0, 0), Complex(1, 0)]
  #  g.compute q.to_vector	# -> Vector[Complex(0, 0), Complex(-1, 0)]
  #
  class PauliZ
    PARAMS = false # :nodoc:
    # 
    # Takes vector with two complex numbers (amplitudes), perform PauliZ operation
    # and returns new vector.
    #
    def compute vector
      Matrices::Z * vector
    end
  end

  #
  # The phase shifter gate. It takes one parameter - theta - which is phase
  # shift.
  #
  #  q = Qubit.new
  #  q.not
  #  g = Gates::PhaseShifter 0.2
  #  q.to_vector		# -> Vector[Complex(0, 0), Complex(1, 0)]
  #  g.compute q.to_vector	# -> Vector[Complex(0, 0), Complex(1.0, -4.89842541528951e-16)]
  #
  class PhaseShifter
    PARAMS = true # :nodoc:
    #
    # Creates new phase shifter gate, using given value as theta.
    #
    def initialize theta
      theta = theta[0] if theta.class == Array
      @theta = theta
    end
    # 
    # Takes vector with two complex numbers (amplitudes), perform Phase Shift
    # operation and returns new vector.
    #
    def compute v
      Matrices.R(@theta) * v
    end
  end

  #
  # Phase rotation - it's phase shifter gate with phase shift = arccos(3/5)
  #
  #  q = Qubit.new
  #  q.not
  #  g = Gates::PhaseRotation
  #  q.to_vector		# -> Vector[Complex(0, 0), Complex(1, 0)]
  #  g.compute q.to_vector	# -> Vector[Complex(0, 0), Complex(-0.902685361933071, -0.430301217000092)]
  #
  class PhaseRotation
    PARAMS = false # :nodoc:
    def initialize
      @phase_shifter = PhaseShifter.new(Math.acos(3/5))
    end
    # 
    # Takes vector with two complex numbers (amplitudes), perform phase rotation operation
    # and returns new vector.
    #
    def compute vector
      @phase_shifter.compute vector
    end
  end

  #
  # Universal control gate - it takes two parameters: 2x2 matrix and control
  # qubit.
  #
  #  control = Qubit.new
  #  target = Qubit.new
  #  control.reset
  #  target.reset
  #
  #  gate = Gates::ControlledU.new [Matrices::NOT, control]
  #  gate.compute target.to_vector   # -> Vector[1+0i, 0i]
  #  control.not
  #  gate.compute target.to_vector   # -> Vector[0i, 1+0i]
  #  control.hadamard
  #  gate.compute target.to_vector   # -> Vector[0.707106781186547+0.0i, -0.707106781186547+0.0i]
  #
  class ControlledU
    PARAMS = true # :nodoc:
    #
    # Creates any controlled gate, using given Matrix and control qubit.
    #
    #  q = Qubit.new
    #  m = Matrix[[0, 1], [1, 0]]
    #  g = ControlledU [m, q]
    #
    def initialize param
      @u = param[0]
      @control = param[1]
    end
    # 
    # Takes vector with two complex numbers (amplitudes), perform gate operation
    # and returns new vector.
    #
    def compute vector
      w = (Matrices.C(@u) * (@control.to_vector.tensor_product vector))
      Vector.elements([w[0] + w[2], w[1] + w[3]])
    end
  end

  #
  # Controlled not gate - if the control qubit (passed as the parameter) equals
  # 1, then it's works as the simple Not gate. Otherwise nothing is done.
  #
  #  control = Qubit.new
  #  target = Qubit.new
  #  control.reset
  #  target.reset
  #
  #  gate = Gates::ControlledNot.new control
  #  gate.compute target.to_vector   # -> Vector[1+0i, 0i]
  #  control.not
  #  gate.compute target.to_vector   # -> Vector[0i, 1+0i]
  #  control.hadamard
  #  gate.compute target.to_vector   # -> Vector[0.707106781186547+0.0i, -0.707106781186547+0.0i]
  #
  class ControlledNot
    PARAMS = true # :nodoc:
    #
    # Creates new Controlled Not gate using given control qubit.
    #
    def initialize control
      @controlled_u = ControlledU.new([Matrices::NOT, control])
    end
    # 
    # Takes vector with two complex numbers (amplitudes), perform gate operation
    # and returns new vector.
    #
    def compute vector
      @controlled_u.compute vector
    end
  end

  private
  #
  # Matrices use to simulate gates.
  #
  class Matrices
    #
    # Simple NOT gate
    #
    NOT = Matrix[
      [0, 1],
      [1, 0]
    ]

    #
    # Hadamard gate
    #
    H = 1/Math.sqrt(2) * Matrix[
      [1,  1],
      [1, -1]
    ]

    #
    # PauliZ gate
    #
    Z = Matrix[
      [1,  0],
      [0, -1]
    ]

    #
    # PauliY gate
    #
    Y = Z * NOT

    # 
    # Phase shifter gate
    #
    #  Gates::Matrices.R 0.2	# -> Matrix[[1, 0], [0, Complex(0.309016994374947, 0.951056516295154)]]
    #
    def Matrices.R theta
      Matrix[
        [1, 0],
        [0, Math.exp(2*Math::PI*Complex::I*theta)]
      ]
    end

    #
    # Controlled-U gate
    #
    #  Gates::Matrices.C Matrix[[1.5, 2.5], [3.5, 4.5]]	# -> Matrix[[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1.5, 2.5], [0, 0, 3.5, 4.5]]
    def Matrices.C matrix
      Matrix[
        [1, 0, 0, 0],
        [0, 1, 0, 0],
        [0, 0] + matrix.row(0).to_a,
        [0, 0] + matrix.row(1).to_a
      ]
    end
  end
end

# 
# A quantum bit, or qubit is a unit of quantum information. That information is
# described by a state vector in a two-level quantum mechanical system which is
# formally equivalent to a two-dimensional vector space over the complex
# numbers. (Wikipedia)
#
# That class represents one qubit.
#
#  q = Qubit.new
#  q.reset
#  q.to_vector	# -> Vector[Complex(1, 0), Complex(0, 0)]
#  q.alpha	# -> Complex(1, 0)
#  q.beta	# -> Complex(1, 0)
#
# Qubit class has also methods generated from classes with <tt>compute</tt> method from
# the Gates module.
#
#  q = Qubit.new
#  q.reset
#  q.not
#  q.to_vector	# -> Vector[Complex(0, 0), Complex(1, 0)]
#  q.phase_shifter 0.2
#  q.to_vector	# -> Vector[Complex(0, 0), Complex(0.309016994374947, 0.951056516295154)]
#
class Qubit
  # 
  # Generating methods for quantum gates from classes in the Gates module.
  #
  Gates.constants.each { |name|
    c = eval("Gates::" + name)
    next unless c.method_defined? :compute
    if c::PARAMS
      define_method(name.underscore) {|param|
        @sequence += [c.new param]
      }
    else
      define_method(name.underscore) {
        @sequence += [c.new]
      }
    end
  }

  #
  # At the begining qubit has random amplitudes.
  # 
  #  q = Qubit.new
  #  q.to_vector	# -> Vector[Complex(-0.273340195540226, -0.459092674794958), Complex(-0.530742833147416, 0.657899003278084)]
  #
  def initialize
    x1 = rand*2 - 1
    range = Math.sqrt(1-x1**2)
    y1 = rand*2*range - range
    @alpha = Complex.new(x1, y1)

    range = Math.sqrt(1-@alpha.abs**2)
    x2 = rand*2*range - range
    y2 = Math.sqrt((Math.sqrt(1-@alpha.abs**2)**2)-x2**2)
    @beta = Complex.new(x2, y2)
    @sequence = Array.new
  end

  #
  # Prints the state of the qubit.
  # 
  # Warning: that operations is possible only in the simulator environment. In
  # the "real world" it's impossible to check the qubit state without changing
  # it. The "real" check can be done with measure method.
  #
  #  irb(main):021:0> q.pr
  #  -0.273340195540226-0.459092674794958i	|0>
  #  -0.530742833147416+0.657899003278084i	|1>
  #
  def pr
    v = to_vector
    puts v[0].to_s + "\t|0>\n" + v[1].to_s + "\t|1>\n"
  end

  #
  # Returns state of qubit as vector with two complex numbers.
  #
  #  q = Qubit.new
  #  q.reset
  #  q.to_vector	# -> Vector[Complex(1, 0), Complex(0, 0)]
  #
  def to_vector
    case @tmp
    when 0
      return Vector.elements([Complex.new(1,0),Complex.new(0,0)])
    when 1
      return Vector.elements([Complex.new(0,0),Complex.new(1,0)])
    end

    v = Vector.elements([@alpha, @beta])
    @sequence.each {|g| v = g.compute v}
    v
  end

  #
  # Resets state of the qubit to the val and deletes all quantum gates.
  # 
  #  q = Qubit.new
  #  q.hadamard
  #  q.reset
  #  q.to_vector	# -> Vector[Complex(1, 0), Complex(0, 0)]
  #
  def reset val=0
    @alpha = Complex.new(1,0)
    @beta = Complex.new(0,0)
    @sequence = Array.new

    self.not if val==1
  end

  #
  # Returns first amplitude.
  #
  #  q = Qubit.new
  #  q.reset
  #  q.alpha	# -> Complex(1, 0)
  #
  def alpha
    to_vector[0]
  end

  #
  # Returns second amplitude.
  #
  #  q = Qubit.new
  #  q.reset
  #  q.beta	# -> Complex(0, 0)
  #
  def beta
    to_vector[1]
  end

  #
  # Measure the qubit state. If the qubit is in the superposition, then that
  # method get random state (according to the probability computed using
  # amplitudes), set the qubit to it, and returns it - so the method is
  # equivalent to the "real world" check qubit state.
  #
  #  q = Qubit.new
  #  q.reset
  #  q.hadamard
  #  q.to_vector	# -> Vector[Complex(0.707106781186547, 0.0), Complex(0.707106781186547, 0.0)]
  #  q.measure	# -> 1
  #  q.to_vector	# -> Vector[Complex(0, 0), Complex(1, 0)]
  #
  def measure
    if rand < alpha.abs**2
      reset
      return 0
    else
      reset 1
      return 1
    end
  end

  # 
  # Sets the qubit set to 0 or 1 for a while. If parameter what is nil, then
  # qubit backs to the normal state.
  #
  def set_tmp what # :nodoc:
    @tmp = what
  end
end

#
# QuantumReg represents register of several qubits. It's also checks the
# consitency of Universal control and Controlled not gates (searching the
# loops - when state of one qubit depends on second and vice versa).
#
#  r = QuantumReg.new 2
#  r.reset
#  r.hadamard 0
#  r.not 1
#  r.pr
#
#  0.707106781186547+0.0i	|01>
#  0.707106781186547+0.0i	|11>
#
# QuantumReg class (like Qubit) has also methods generated from the classes with
# <tt>compute</tt> method from Gates module. It always takes one parameter - number
# of registry qubit.
#
class QuantumReg
  #
  # It's generating new registry with given length.
  #
  def initialize length
    @reg = Array.new length
    0.upto(length-1) { |i| @reg[i] = Qubit.new }
    @control_hash = Hash.new
  end

  #
  # Prints register state. See the warning in the pr method in Qubit class.
  #
  #  q = QuantumReg.new 2
  #  q.reset
  #  q.pr
  #
  #  1+0i	|00>
  #
  def pr
    result = evaluate
    result.each_index do |i|
      next if(result[i]==0)
      printf "%s\t|%0#{@reg.size}b>\n", result[i].to_s, i
    end
  end

  #
  # Measure state of range of the qubits. With no parameters it's measure
  # state of all qubits. With one parameter - measure state of the qubit
  # with given index. With two parameters - measure state of qubit with
  # indexes begining with first and ending with last parameter.
  #
  def measure from=nil, to=nil
    if !from && !to
      from=0
      to=@reg.size-1
    elsif !to
      to=from
    end
    result = 0
    w = 1
    to.downto from do |i|
      result += w if @reg[i].measure == 1
      w *= 2
    end
    result
  end

  # 
  # Resets state of the registry.
  #
  def reset
    @reg.each {|q| q.reset}
    @control_hash = Hash.new
  end

  #
  # Controlled universal gate. Takes the control and the target qubit indexes
  # and the 2x2 matrix as parameters.
  #
  def controlled_u control, target, u
    return unless controlled control, target
    @reg[target].controlled_u u, @reg[control]
  end

  #
  # Controlled not gate. Takes the control and the target qubit indexes as
  # parameters.
  #
  def controlled_not control, target
    return unless controlled control, target
    @reg[target].controlled_not @reg[control]
  end

  #
  # Generating methods for quantum gates (see the Gates module). Each method
  # takes at last one parameter - the target qubit index. It's always the first
  # parameter. Methods can also take more parameters - it depends on the
  # operation.
  #
  Gates.constants.each { |name|
    c = eval("Gates::" + name)
    next unless c.method_defined? :compute
    next if QuantumReg.method_defined? name.underscore
    if c::PARAMS
      define_method(name.underscore) {|i, *param|
        @reg[i].send(name.underscore, [i] + param)
      }
    else
      define_method(name.underscore) { |i|
        @reg[i].send(name.underscore)
      }
    end
  }

  private
  def check_control_hash
    @reg.each_index do |i|
      @tmp_set = Set.new
      return false unless check_control_index i
    end
    true
  end

  def check_control_index i
    return true unless @control_hash[i]

    @tmp_set += [i]
    @control_hash[i].each do |v|
      return false if @tmp_set.include?(v)
      return check_control_index v
    end
    true
  end

  def evaluate
    result = Array.new(2**@reg.size)
    order = Array.new
    while order.size != @reg.size
      @reg.each_index do |i|
        next if @control_hash[i] && !@control_hash[i].subset?(Set.new(order))
	next if order.index(i)
        order += [i]
      end
    end
    0.upto(2**@reg.size-1) do |c|
      prob = Complex.new(1,0)
      order.each do |i|
        b = c & 2**(@reg.size - i - 1)
	if b == 0
          prob *= @reg[i].alpha
	  @reg[i].set_tmp 0
	else
          prob *= @reg[i].beta
	  @reg[i].set_tmp 1
	end
      end
      result[c] = prob
      @reg.each {|i| i.set_tmp nil}
    end
    result
  end

  def controlled control, target
    @control_hash[target] = Set.new unless @control_hash[target]
    @control_hash[target] += [control]

    unless check_control_hash
      puts "Inifinite dependency in register detected!"
      puts "Controlled NOT hasn't been added"
      @control_hash[target].delete(control)
      return false
    end
    return true
  end
end