=begin rdoc
  
== SphereLab

Definition of Vector and Body objects used for n-body simulations.

=end

module RubyLabs
	
module SphereLab
  
=begin rdoc
  These class variables maintain the state of the display and other miscellaneous
  global values.
=end
  
  @@sphereDirectory = File.join(File.dirname(__FILE__), '..', 'data', 'spheres')

  @@viewerOptions = {
    :dotColor => '#000080',
    :dotRadius => 1.0,
    :origin => :center,
  }
  
  @@droppingOptions = {
    :canvasSize => 400,
    :mxmin => 100,
    :mymin => 50,
    :hmax => 100,
    :dash => 1,
  }
  
  @@robotOptions = {
    :canvasSize => 400,
    :polygon => [4,0,0,10,4,8,8,10],    
  }
  
  @@drawing = nil
  @@delay = 0.1
  
  NBodyView = Struct.new(:bodies, :origin, :scale, :options)
  TurtleView = Struct.new(:turtle, :options)
  MelonView = Struct.new(:bodies, :scale, :ground, :startloc, :options)

=begin rdoc
  The constant G is the universal gravitational constant, assuming mass is
  in units of kilograms, distances are in meters, and time is in seconds.
=end

  G = 6.67E-11

=begin rdoc
  A Vector is a 3-tuple of (x,y,z) coordinates.  Operators defined for
  vector objects (where v is a vector and a is a scalar):
     v == v
     v + v
     v - v
     v * a
     v.add(v)     (equivalent of v += v)
     v.sub(v)     (equivalent of v -= v)
     v.scale(a)   (equivalent of v *= a)
     v.norm
=end

  class Vector
    attr_accessor :x, :y, :z
  
=begin rdoc
  Make a new vector with the specified x, y, and z components.
=end
  
    def initialize(*args)
      @x, @y, @z = args
    end
  
    def inspect
      pos = [@x,@y,@z].map { |x| sprintf "%.5g", x }
      return "(" + pos.join(",") + ")"
    end
  
=begin rdoc
  v1 == v2 if the three components are the same
=end

    def ==(v)
      return (@x == v.x) && (@y == v.y) && (@z == v.z)
    end

=begin rdoc
  Arithmetic methods are invoked for a vector v1 when Ruby evaluates an expression of the
  form v1 <op> v2 where <op> is +, -, or *.  They create a a new vector containing the 
  result of the operation.  + and - do element-wise addition or and subtraction, *
  is a vector-scalar multiplication.
=end
  
    def +(v)
      Vector.new(@x + v.x, @y + v.y, @z + v.z)
    end
  
    def -(v)
      Vector.new(@x - v.x, @y - v.y, @z - v.z)
    end

    def *(a)
      Vector.new(@x * a, @y * a, @z * a)
    end

=begin rdoc
  v1.add(v2) adds the components of v2 to v1 -- would be v1 += v2 if Ruby
  allowed us to overload +=
=end

    def add(v)
      @x += v.x
      @y += v.y
      @z += v.z
      self
    end
  
=begin rdoc
  v1.sub(v2) subtracts the components of v2 from v1 -- would be v1 -= v2 if Ruby
  allowed us to overload -=
=end

    def sub(v)
      @x -= v.x
      @y -= v.y
      @z -= v.z
      self
    end
    
=begin rdoc
  v.scale(a) -- multiply each element in v by scalar a
=end

    def scale(a)
      @x *= a
      @y *= a
      @z *= a
    end

=begin rdoc
  The magnitude of a vector is the Euclidean norm: sqrt(x**2 + y**2 + z**2)
=end

    def norm
      sqrt(@x*@x + @y*@y + @z*@z)
    end
    
=begin rdoc
  Compute the angle between this vector and v -- used when drawing 2D vectors, so
  the z dimension is ignored
=end

    def angle(v)
      acos((@x * v.x + @y * v.y) / (norm * v.norm))
    end
    
=begin rdoc
  Accessor functions for the coordinates as a vector of three floats
=end

    def coords
      [@x, @y, @z]
    end
    
    def coords=(a)
      @x = a[0]
      @y = a[1]
      @z = a[2]
    end
  
  end # Vector


=begin rdoc
  A Body object represents the state of a celestial body.  A body has mass (a scalar),
  position (a vector), and velocity (a vector).  A third vector, named force, is used
  when calculating forces acting on a body.  The size and color attributes are used by
  the visualization methods.
=end

  class Body
    attr_accessor :mass, :position, :velocity, :force
    attr_accessor :name, :size, :color, :prevx, :prevy, :graphic
  
    def initialize(*args)
      if args[1].class == Vector
        @mass, @position, @velocity, @name = args
      # elsif args[0] == :random
      #   @mass = rand(10000) * 1e6
      #   @position = Vector.new(rand(300)-150, rand(300)-150, 0)
      #   @velocity = Vector.new(rand(20)-10, rand(10)-5, 0)
      #   @name = "b" + self.object_id.to_s
      # elsif args[0].is_a?(Numeric) && args[1].is_a?(Numeric)
      #   @mass = args[0]
      #   @position = Vector.new(args[1], 0.0, 0.0)
      #   @velocity = Vector.new(0.0, 0.0, 0.0)
      #   @name = args[2]
      #   @linear = true
      else
        @mass = 0.0
        @position = Vector.new(0.0, 0.0, 0.0)
        @velocity = Vector.new(0.0, 0.0, 0.0)
        @name = nil
      end
      @force = Vector.new(0.0, 0.0, 0.0)
    end
    
    def clone
      copy = super
      if graphic
        copy.position = position.clone
        copy.velocity = velocity.clone
        copy.force = force.clone
        copy.graphic = Canvas.circle(prevx, prevy, size, :fill => color)
      end
      return copy
    end
  
    def inspect
      name = @name ? @name : ""
      return sprintf "%s: %.3g kg %s %s", name, @mass, @position.inspect, @velocity.inspect
    end

=begin rdoc
  Reset the force vector to 0 (to get ready for a new round of interaction calculations).
=end

    def clear_force
      @force.x = 0.0
      @force.y = 0.0
      @force.z = 0.0
    end
  
=begin rdoc
  Compute force acting on this body wrt body b
=end

    def add_force(b)
      r = @position - b.position
      nr = r.norm ** 3
      mr = b.mass / nr
      @force.add(r * mr)
    end
  
=begin rdoc
  Move this body by applying current force vector for dt seconds
=end

    def move(dt)
      acc = @force * G * -1.0
      @velocity.add( acc * dt ) 
      @position.add( @velocity * dt )
    end
  
=begin rdoc
  Class method to compute the interaction between bodies b1 and b2 and update
  their force vectors.
=end

    def Body.interaction(b1, b2)
      r = b1.position - b2.position
      a = r.norm ** 3
      b1.force.add(r * (b2.mass / a))
      b2.force.add(r * (-b1.mass / a))
    end
  
  end # Body
  
=begin rdoc
  A class for rudimentary Turtle graphics.  All we need is enough functionality to
  implement a robot that draws a circle by moving forward, then correcting its direction.
  Attributes:  location and velocity vectors.
=end

  class Turtle
    
    attr_accessor :position, :velocity, :reference, :graphic
    
    @@turtleOptions = {
      :x => 20,
      :y => 100,
      :heading => 0,
      :speed => 10,
      :track => :off,
    }
    
    @@north = Vector.new(0, 10, 0)

    def initialize(args = nil)
      args = { } if args.nil?
      raise "Turtle: initialize: args must be a Hash" unless args.class == Hash
      options = @@turtleOptions.merge(args)
      alpha = Canvas.radians(options[:heading] + 90.0)
      @position = Vector.new(options[:x], options[:y], 0.0)
      @velocity = Vector.new(-1.0 * options[:speed] * cos(alpha), options[:speed] * sin(alpha), 0.0)
      @graphic = nil
      @reference = nil
      @tracking = options[:track]
    end
    
    def inspect
      sprintf "#<SphereLab::Turtle x: %d y: %d heading: %d speed: %.2f>", @position.x, @position.y, heading, speed 
    end
    
    def location
      return [@position.x, @position.y]
    end
    
    def heading
      d = Canvas.degrees( @velocity.angle(@@north) )
      return (@velocity.x < 0) ? (360 - d) : d
    end
    
    def speed
      @velocity.norm
    end
    
    def turn(alpha)
      theta = -1.0 * Canvas.radians(alpha)
      x = @velocity.x * cos(theta) - @velocity.y * sin(theta)
      y = @velocity.x * sin(theta) + @velocity.y * cos(theta)
      @velocity.x = x
      @velocity.y = y
      if @graphic
        Canvas.rotate(@graphic,alpha)
        Canvas.sync
      end
      return alpha
    end
    
    def advance(dt)
      prevx = @position.x
      prevy = @position.y
      @position.add( @velocity * dt )
      if @graphic
        Canvas.move(@graphic, @position.x-prevx, prevy-@position.y, @tracking)
        Canvas.sync
      end
      return dt
    end  
    
    def orient
      if @reference.nil?
        puts "no reference point"
        return nil
      end
      mx, my = @reference
      tx, ty = @position.x, @position.y
      mid = Vector.new(tx-mx, ty-my, 0.0)
      theta = Canvas.degrees(mid.angle(@velocity))
      alpha = 2 * (90.0 - theta)
      turn(alpha)
      Canvas.sync if @graphic
      return alpha      
    end
    
    def track(val)
      case val
      when :off  : @tracking = :off
      when :on   : @tracking = :track
      else
        puts "tracking is either :on or :off"
        return nil
      end
      return val
    end
    
  end # class Turtle
  
=begin rdoc
  Set up an experiment with a robot/turtle.  Initialize the canvas, make the turtle,
  assign it a graphic.  These methods are all defined with respect to a canvas, and
  they will print an error message if the experiment has not been set up with a call
  to init_robot (which sets @@drawing).
=end

  def view_robot(userOptions = {})
    options = @@robotOptions.merge(userOptions)
    poly = options[:polygon].clone
    edge = options[:canvasSize]
    (0...poly.length).step(2) do |i| 
      poly[i] += edge/10
      poly[i+1] += edge/2
    end
    Canvas.init(edge, edge, "SphereLab")
    turtle = Turtle.new( :x => edge/10, :y => edge/2, :heading => 0, :speed => 10 )
    turtle.graphic = Canvas.polygon(poly, :outline => 'black', :fill => '#00ff88')
    if options[:flag]
      turtle.set_flag( *options[:flag] )
    end
    if options[:track]
      turtle.track( options[:track] )
    end
    class <<turtle
      def plant_flag
        set_flag( @position.x, @position.y )
        Canvas.sync
      end
    end
    @@drawing = TurtleView.new( turtle, options )
    Canvas.sync
    return true
  end
  
  def set_flag(fx, fy)
    r = 3.0
    Canvas.circle( fx + r/2, fy + r/2, r, :fill => 'darkblue' )
    @reference = [ fx, fy ]          
  end
    
  def robot
    if @@drawing.nil?
      puts "No robot; call view_robot to initialize"
      return nil
    else
      return @@drawing.turtle
    end
  end
      
  # :begin :make_circle
  def make_circle(nseg, time, angle)
    view_robot
    robot.turn( angle/2 )
    nseg.times { robot.advance(time); robot.turn(angle) }
    robot.turn( -angle/2 )
    return true
  end
  # :end :make_circle


=begin rdoc
  These methods create a set of bodies with random size, location, and initial
  velocity.
=end

  def random_vectors(r, i, n)
    theta = (2 * PI / n) * i + (PI * rand / n)
    # radius = r + (r/3)*(rand-0.5)
    radius = r + r * (rand-0.5)
    x = radius * cos(theta)
    y = radius * sin(theta)
    vtheta = (PI - theta) * -1 + PI * (rand-0.5)
    vx = radius/20 * cos(vtheta)
    vy = radius/20 * sin(vtheta)
    return Vector.new(x, y, 0), Vector.new(vx, vy, 0)
  end

  def random_velocity(v, r)
    res = Vector.new(-r * cos)
  end
  
  # todo -- put mm, mr in global options

  def random_bodies(n, big)
    big = 1 if big.nil?
    moving = true if moving.nil?
    res = []
    mm = 1e12          # average mass
    mr = 150          # average distance from origin
    bigm = (mm * 100) / big
    big.times do |i|
      r, v = random_vectors(mr/2, i, big)
      ms = (1 + (rand/2 - 0.25) )
      b = Body.new(bigm*ms, r, v)
      b.name = "b#{i}"
      b.color = '#0080ff'
      b.size = 10*ms
      res << b
    end
    (n-big).times do |i|
      r, v = random_vectors(mr, i, n-big)
      b = Body.new(mm, r, v)
      b.name = "b#{i+big}"
      b.color = '#0080ff'
      b.size = 5
      res << b
    end
    return res
  end
  
  def falling_bodies(n)
    raise "n must be 5 or more" unless n >= 5
    a = random_bodies(n-1, n-1)
    b = Body.new(1e13, (a[0].position + a[1].position), Vector.new(0,0,0))
    # b = Body.new(1e13, (a[0].position + a[1].position)*0.85, Vector.new(0,0,0))
    # pos = a[0].position
    # (1..(n-2)).each { |i| pos.add( a[i].position ) }
    # b = Body.new(1e14, pos * (1.0 / n), Vector.new(0,0,0))
    b.name = "falling"
    b.size = 5
    b.color = 'red'
    a.insert(0, b)
    return a
  end
  
=begin rdoc
  Initialize a new n-body system, returning an array of body objects.  The
  parameter is the name of a predefined system, the name of a file, or the
  symbol :random.
=end

	def make_system(*args)
    bodies = []
		begin
		  raise "usage: make_system(id)" unless args.length > 0
		  if args[0] == :random
		    raise "usage:  make_system(:random, n, m)" unless args.length >= 2 && args[1].class == Fixnum
		    return random_bodies(args[1], args[2])
	    elsif args[0] == :falling
		    raise "usage:  make_system(:falling, n)" unless args.length >= 1 && args[1].class == Fixnum
	      return falling_bodies(args[1])
	    end
	    filename = args[0]
  	  if filename.class == Symbol
  	    filename = File.join(@@sphereDirectory, filename.to_s + ".txt")
      end
      File.open(filename).each do |line|
        line.strip!
        next if line.length == 0
        next if line[0] == ?#
        a = line.chomp.split
        for i in 1..7
          a[i] = a[i].to_f
        end
        b = Body.new( a[1], Vector.new(a[2],a[3],a[4]), Vector.new(a[5],a[6],a[7]), a[0] )
        b.size = a[-2].to_i
        b.color = a[-1]
        bodies << b
      end
      if args[0] == :melon
        class <<bodies[0]
          def height
            return 0 if prevy.nil?
            return position.y - prevy
          end
          # def height=(x)
          # end
        end       
      end
		rescue
			puts "error: #{$!}"
			return nil
		end
		return bodies
	end
	
=begin rdoc
  Write the mass, position, and velocity for each body to a file.  Not intended to
  be used by students, but used to save interesting data sets they can load and use.
=end

# d 4.5e16 300 -75 0 -5 2 0 6 #ff6666

	
	def save_system(b, fn)
	  raise "file exists" if File.exists?(fn)
	  File.open(fn, "w") do |f|
	    b.each do |x|
	      f.printf "%s %g %g %g %g %g %g %g %d %s\n", 
	        x.name, x.mass, 
	        x.position.x, x.position.y, x.position.z, 
	        x.velocity.x, x.velocity.y, x.velocity.z, 
	        x.size, x.color
	    end
	  end
	end
	
=begin rdoc
  Initialize the drawing canvas by drawing a circle for each body in list b.
=end

  def view_system(blist, userOptions = {})
    Canvas.init(700, 700, "SphereLab")
    if prev = @@drawing
      prev.bodies.each { |x| x.graphic = nil }
    end
    options = @@viewerOptions.merge(userOptions)
    origin = setOrigin(options[:origin])
    sf = setScale(blist, options[:origin], options[:scale])
    blist.each do |b|
      x, y = scale(b.position, origin, sf)
      b.graphic = Canvas.circle(x, y, b.size, :fill => b.color)
      b.prevx = x
      b.prevy = y
    end
    @@drawing = NBodyView.new(blist, origin, sf, options)
    if options.has_key?(:dash)
      options[:dashcount] = 0
      options[:pendown] = :track
    elsif options.has_key?(:pendown)
      options[:pendown] = :track
    end
    Canvas.sync
    return true  
  end

=begin rdoc
  Demonstrate adding force vectors by moving only one body
=end

  def update_one(falling, stationary, time)
    if falling.graphic.nil?
      puts "display the system with view_system"
      return nil
    end
    stationary.each do |x|
      Body.interaction( falling, x )
    end
    falling.move(time)
    if @@drawing.options.has_key?(:dash)
      @@drawing.options[:dashcount] = (@@drawing.options[:dashcount] + 1) % @@drawing.options[:dash]
      if @@drawing.options[:dashcount] == 0
        @@drawing.options[:pendown] = @@drawing.options[:pendown].nil? ? :track : nil
      end        
    end
    newx, newy = scale(falling.position, @@drawing.origin, @@drawing.scale)
    Canvas.move(falling.graphic, newx-falling.prevx, newy-falling.prevy, @@drawing.options[:pendown])
    falling.prevx = newx
    falling.prevy = newy
    falling.clear_force
    Canvas.sync
    return true
  end

=begin rdoc
  Call SphereLab::step(bodies, time) to calculate the pairwise interactions between all
  Body objects in the array +bodies+ and then compute their new positions after +time+ seconds.
=end

# :begin :step_system
  def step_system(bodies, dt)
    nb = bodies.length
  
    for i in 0..(nb-1)         # compute all pairwise interactions
      for j in (i+1)..(nb-1)
        Body.interaction( bodies[i], bodies[j] )
      end
    end
  
    bodies.each do |b|
      b.move(dt)          # apply the accumulated forces
      b.clear_force         # reset force to 0 for next round
    end
  end
# :end :step_system
    
=begin rdoc
  After making a canvas with view_system, call this method to move the bodies on the canvas.
=end

  def update_system(bodies, dt)
    return false unless @@drawing
    step_system(bodies, dt)
    if @@drawing.options.has_key?(:dash)
      @@drawing.options[:dashcount] = (@@drawing.options[:dashcount] + 1) % @@drawing.options[:dash]
      if @@drawing.options[:dashcount] == 0
        @@drawing.options[:pendown] = @@drawing.options[:pendown].nil? ? :track : nil
      end        
    end
    bodies.each do |b|
      next unless b.graphic
      newx, newy = scale(b.position, @@drawing.origin, @@drawing.scale)
      Canvas.move(b.graphic, newx-b.prevx, newy-b.prevy, @@drawing.options[:pendown])
      b.prevx = newx
      b.prevy = newy
    end
    Canvas.sync
    return true
  end

=begin rdoc
  Map a simulation's (x,y) coordinates to screen coordinates using origin
  and scale factor
=end

  def scale(vec, origin, sf)
    loc = vec.clone
    loc.scale(sf)
    loc.add(origin)
    return loc.x, loc.y
  end
    
=begin rdoc
  Make a vector that defines the location of the origin (in pixels).
=end

  def setOrigin(type)
    case type
    when :center
      Vector.new( Canvas.width/2, Canvas.height/2, 0)
    else
      Vector.new(0,0,0)
    end
  end
  
=begin rdoc
  Set the scale factor.  Use the parameter passed by the user, or find the
  largest coordinate in a list of bodies.  Add 20% for a margin
=end

  def setScale(blist, origin, scale)
    if scale == nil
      dmax = 0.0
      blist.each do |b|
        b.position.coords.each { |val| dmax = val.abs if val.abs > dmax } 
      end
    else
      dmax = scale
    end
    sf = (origin == :center) ? (Canvas.width / 2.0) : Canvas.width 
    return (sf / dmax) * 0.8
  end
  
=begin rdoc
  Initialize the drawing for a "2-body system" with a small object and the
  earth.  The parameters are the list of Body objects and viewing options.
=end
  
  def view_melon(blist, userOptions = {})
    options = @@droppingOptions.merge(userOptions)
    options[:dashcount] = 0
    options[:pendown] = :track
    edge = options[:canvasSize]
    mxmin = options[:mxmin]
    mymin = options[:mymin]
    hmax = options[:hmax]
    Canvas.init(edge, edge, "SphereLab")
    earth = Canvas.circle(200, 2150, 1800, :fill => blist[1].color)
    blist[1].graphic = earth
    mymax = earth.coords[1]
    melon = Canvas.circle(mxmin, mymax, 5, :fill => blist[0].color)
    blist[0].graphic = melon
    scale = (mymax-mymin) / hmax.to_f
    @@drawing = MelonView.new(blist, scale, blist[0].position.y, melon.coords, options)
    Canvas.sync
    return true
  end
  
=begin rdoc
  Position the melon (body 0) at a specified height.  If no drawing, save the current
  position in prevy, but if drawing, get the earth's surface from the drawing (this allows
  students to move the melon up or down in the middle of an experiment)
=end
  
  def position_melon(blist, height)
    melon = blist[0]
    if @@drawing && blist[0].graphic
      if height < 0 || height > @@drawing.options[:hmax]
        puts "Height must be between 0 and #{@@drawing.options[:hmax]} meters"
        return false
      end
      melon.prevy = @@drawing.ground
      melon.position.y = @@drawing.ground + height
      a = @@drawing.startloc.clone
      a[1] -= height * @@drawing.scale
      a[3] -= height * @@drawing.scale
      melon.graphic.coords = a
      Canvas.sync
    else
      melon.prevy = melon.position.y
      melon.position.y += height
    end
    melon.velocity.y = 0.0
    return height
  end
  
  def update_melon(blist, dt)
    melon = blist[0]
    mx = melon.position.x
    my = melon.position.y
    return "splat" if melon.prevy.nil? || my < melon.prevy
    step_system(blist, dt)
    if @@drawing && blist[0].graphic
      if @@drawing.options[:dash] > 0
        @@drawing.options[:dashcount] = (@@drawing.options[:dashcount] + 1) % @@drawing.options[:dash]
        if @@drawing.options[:dashcount] == 0
          @@drawing.options[:pendown] = @@drawing.options[:pendown].nil? ? :track : nil
        end
      end     
      dx = (melon.position.x - mx) * @@drawing.scale
      dy = (my - melon.position.y) * @@drawing.scale
      Canvas.move(melon.graphic, dx, dy, @@drawing.options[:pendown])
      Canvas.sync
    end
    return blist[0].height
  end
  
  # :begin :drop_melon
  def drop_melon(blist, dt)
    count = 0
    loop do
      res = update_melon(blist, dt)
      break if res == "splat"
      count += 1
    end
    Canvas.sync if @@drawing
    return count * dt
  end
  # :end :drop_melon

  
=begin rdoc
  Methods used in IRB sessions.
=end

# :begin :dist
  def dist(r, t)
    return r * t
  end
# :end :dist

# :begin :falling
  def falling(t)
    return 0.5 * 9.8 * t**2
  end
# :end :falling

end # SphereLab

end # RubyLabs