#!/usr/bin/ruby -w

=begin
/***************************************************************************
 *   Copyright (C) 2006, Paul Lutus                                      *
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.,                                       *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
 ***************************************************************************/
=end

require 'gravityui_ui'

PROGRAM_VERSION = "1.4"

# subclass of QFrame to control painting updates

class MyFrame < Qt::Frame
   def initialize(parent,a,b)
      super(a,b)
      @parent = parent
      setPaletteBackgroundColor(Qt::black)
      Qt::ToolTip.add(self, trUtf8("Drag mouse to rotate, zoom with wheel") )
   end
   # redraw image on these events
   def paintEvent(*e)
      @parent.draw_image
   end
   def resizeEvent(*e)
      @parent.draw_image
   end
end

# a class for routines and constants of common utility

class CommonCode
   CommonCode::ToRad = Math::PI / 180.0
   CommonCode::ToDeg = 180.0 / Math::PI
   def CommonCode::fmt_num(n)
      sprintf("%.2e",n)
   end
end

# solar system data class, all units mks

class SolarSystem
   SolarSystem::Data=<<-EOF
   "Name","OrbitRad","BodyRad","Mass","OrbitVel"
   "Sun",0,695000000,1.989E+030,0
   "Mercury",57900000000,2440000,3.33E+023,47900
   "Venus",108000000000,6050000,4.869E+024,35000
   "Earth",150000000000,6378140,5.976E+024,29800
   "Mars",227940000000,3397200,6.421E+023,24100
   "Jupiter",778330000000,71492000,1.9E+027,13100
   "Saturn",1429400000000,60268000,5.688E+026,9640
   "Uranus",2870990000000,25559000,8.686E+025,6810
   "Neptune",4504300000000,24746000,1.024E+026,5430
   "Pluto",5913520000000,1137000,1.27E+022,4740
   EOF
end

# a Cartesian 3D vector class
# with a number of important operator overrides

class Cart3
   attr_accessor :x,:y,:z
   def initialize(x = 0,y = 0,z = 0)
      if(x.class == self.class)
         @x = x.x
         @y = x.y
         @z = x.z
      else
         @x = x
         @y = y
         @z = z
      end
   end

   def -(e)
      Cart3.new(@x - e.x,@y - e.y,@z - e.z)
   end

   def +(e)
      Cart3.new(@x + e.x,@y + e.y,@z + e.z)
   end

   def *(e)
      if(e.class != self.class)
         # multiply by scalar
         Cart3.new(@x * e,@y * e,@z * e)
      else
         # multiply by vector
         Cart3.new(@x * e.x,@y * e.y,@z * e.z)
      end
   end

   def /(e)
      if(e.class != self.class)
         # divide by scalar
         Cart3.new(@x / e,@y / e,@z / e)
      else
         # divide by vector
         Cart3.new(@x / e.x,@y / e.y,@z / e.z)
      end
   end

   # sum of squares
   def sumsq
      @x*@x+@y*@y+@z*@z
   end

   def to_s
      "[#{CommonCode::fmt_num(@x)},#{CommonCode::fmt_num(@y)},#{CommonCode::fmt_num(@z)}]"
   end
end # class Cart3

=begin

Planet, a simple data class

name = string
radius = float
pos = Cart3
vel = Cart3
mass = float

=end

class Planet
   attr_accessor :name,:radius,:pos,:vel,:mass
   def initialize(name,radius,pos,vel,mass = 0)
      @name = name.gsub(/"/,"")
      @radius = radius
      @pos = pos
      @vel = vel
      @mass = mass
   end
   def to_s
      "#{@name},#{CommonCode::fmt_num(@radius)},#{@pos},#{@vel},#{CommonCode::fmt_num(@mass)}"
   end
end

# RotationMatrix performs 3D rotations and perspective

class RotationMatrix

   # perspective depth cue for 3D -> 2D transformation
   PerspectiveDepth = 16
   # empirical constant for anaglyphic rotation
   AnaglyphScale = 0.03

   RotationMatrix::ToRad = Math::PI / 180.0
   RotationMatrix::ToDeg = 180.0 / Math::PI

   # populate 3D matrix with values for x,y,z rotations

   def populate_matrix(xa,ya)

      # create trig values
      @sy = Math.sin(xa * RotationMatrix::ToRad);
      @cy = Math.cos(xa * RotationMatrix::ToRad);
      @sx = Math.sin(ya * RotationMatrix::ToRad);
      @cx = Math.cos(ya * RotationMatrix::ToRad);
   end

   # 3D -> 2D, add perspective cue,
   # perform anaglyph position change if specified

   def convert_3d_to_2d(v,anaglyph_flag = 0)
      v.x = (v.x * (PerspectiveDepth + v.z))/PerspectiveDepth
      v.x += v.z * anaglyph_flag * AnaglyphScale if anaglyph_flag
      v.y = (v.y * (PerspectiveDepth + v.z))/PerspectiveDepth
   end

   # rotate a 3D point using matrix values

   def rotate(v)
      # borrowed from my "Apple World" 1979
      hf = (v.x * @sx - v.z * @cx)
      py = v.y * @cy + @sy * hf
      px = v.x * @cx + v.z * @sx
      pz = -v.y * @sy + @cy * hf
      v.x = px; v.y = py; v.z = pz
   end
end # class RotationMatrix

=begin

Gravitational force f, Newtons, between two masses M and m:

f = G M m
   ------
    r^2

G = universal gravitational constant, 6.6742e-11 N m^2 / kg^2

M,m = masses of the two bodies, kilograms

r = radius (distance) between M and m, meters

acceleration, m/s, a for a force f and a mass m:

a = f/m

The shorthand version drops one of the masses
for a slight speed improvement, and goes
directly to acceleration a:

a = G M
   -----
    r^2

BUT the shorthand form assumes one
of the masses is infinite

In a numerical simulation using time slice dt:

velocity v += a * dt

position p += v * dt

All mks units

=end

class OrbitalPhysics

   # The all-important force of gravity

   OrbitalPhysics::G = 6.6742e-11 # N m^2 / kg^-2

   def OrbitalPhysics::compute_acceleration(pa,pb,dt)
      # don't compute self-gravitation
      if(pa != pb)
         radius = pa.pos - pb.pos
         sumsq = radius.sumsq
         hypot = Math.sqrt(sumsq)
         acc = -G * pb.mass / sumsq
         # this line assigns the acceleration to
         # the x,y,z velocity components along the
         # radius pa - pb without using trig functions
         pa.vel += radius / hypot * acc * dt
      end
   end

   def OrbitalPhysics::process_planets(planet_list,dt,symmetric = false)
      if(symmetric)
         # compute gravitation interactively for all bodies
         # very slow ... requires p^2 time
         planet_list.each do |p1|
            planet_list.each do |p2|
               compute_acceleration(p1,p2,dt)
            end
            p1.pos += p1.vel * dt
         end
      else
         # compute gravitation only wrt the sun
         # which is assumed to be first member of array
         sun = planet_list.first
         planet_list.each do |planet|
            compute_acceleration(planet,sun,dt)
            planet.pos += planet.vel * dt
         end
      end
   end

end # class OrbitalPhysics

# main class

class Gravity < GravityUI

   slots 'perform_orbit_calc()' # for timer assignment

   Gravity::PlanetColors = [ Qt::white,Qt::yellow,Qt::cyan,Qt::Color.new(128,128,255),
   Qt::red,Qt::green,Qt::magenta,Qt::blue ]

   Gravity::AnimStrings = [ "1 hour","2 hours","4 hours","8 hours","16 hours",
   "1 day","2 days","4 days","8 days","16 days","32 days","64 days","128 days","256 days" ]

   Gravity::CometStrings = [ "1","2","4","8","16","32","64","128","256","512" ]

   def initialize(app)
      super()
      title = self.class.name + " " + PROGRAM_VERSION
      setCaption(title)
      @app = app
      @total_time_hours = 0
      # this sets the timer delay
      @anim_time = 1
      # initial drawing scale, change with mouse wheel
      @drawing_scale = 6e-12
      @rotx = -20
      @roty = 0
      @planet_list = []
      @time_step = nil
      @pixmap = nil
      @erase = true
      @anaglyph_mode = false
      @symmetric = false
      @timer = nil
      @rotator = RotationMatrix.new
      # replace automatically created QFrame with my own
      @GravityUILayout.remove(@graphicPane)
      @graphicPane.setHidden(true)
      @graphicPane = MyFrame.new(self,centralWidget(), "graphicPane")
      @GravityUILayout.addMultiCellWidget(@graphicPane, 0, 0, 0, 10)
      @timeStepComboBox.insertStringList(AnimStrings)
      @timeStepComboBox.setCurrentText("16 hours")
      @cometComboBox.insertStringList(CometStrings)
      @cometComboBox.setCurrentText("16")
      set_time_step
      @solarSystemCheckBox.setChecked(true)
      @pixelsSpinBox.setValue(5)
      @min_draw_radius = 5
      load_objects
   end

   def set_time_step
      s = @timeStepComboBox.currentText()
      v,units = s.split(" ")
      v = v.to_i
      v *= 3600 if units =~ /hour/i
      v *= 86400 if units =~ /day/i
      @time_step = v
   end

   def show_status
      y = @total_time_hours
      h = y % 24
      y /= 24
      y *= 100
      d = (y % 36525) / 100
      y /= 36525
      str = sprintf("%6dy %03dd %02dh",y,d,h)
      statusBar.message(str)
   end

   def draw_planets(xsize,ysize,pixpainter,anaglyph_flag = nil,td_color = nil)
      if(anaglyph_flag)
         pixpainter.setBrush(td_color)
         pixpainter.setPen(td_color)
      end
      i = 0
      @planet_list.each do |planet|
         v = planet.pos * @drawing_scale
         @rotator.rotate(v)
         @rotator.convert_3d_to_2d(v,anaglyph_flag)
         sxa = @x_screen_center + (v.x * @screen_scale)
         sya = @y_screen_center - (v.y * @screen_scale)
         if(sxa >= 0 && sxa < xsize && sya >= 0 && sya < ysize)
            # fake the sun's radius for aesthetics
            sr = (i == 0)?4e7:(planet.radius)
            s = Cart3.new(sr * @drawing_scale,0,-planet.pos.z * @drawing_scale);
            @rotator.convert_3d_to_2d(s)
            s.x *= @screen_scale * 100
            s.x = (s.x < @min_draw_radius)?@min_draw_radius:s.x
            sc = s.x/2
            unless(anaglyph_flag)
               col = PlanetColors[i % PlanetColors.size]
               pixpainter.setBrush(col)
               pixpainter.setPen(col)
            end
            pixpainter.drawEllipse(sxa-sc,sya-sc,s.x,s.x)
         end
         i += 1
      end
   end

   def draw_image(erase = false)
      if(@planet_list)
         show_status
         @rotator.populate_matrix(@rotx,@roty)
         xsize,ysize = @graphicPane.width,@graphicPane.height
         # create an off-screen pixmap for image drawing
         # whenever a change requires it
         if(@pixmap == nil || xsize != @old_xsize || ysize != @old_ysize)
            @pixmap = Qt::Pixmap.new(xsize,ysize)
            @old_xsize = xsize
            @old_ysize = ysize
            @x_screen_center = xsize / 2
            @y_screen_center = ysize / 2
            @screen_scale = (@x_screen_center > @y_screen_center)?@y_screen_center:@x_screen_center
            @pixmap.fill(Qt::black)
         end
         # set image background to black
         @pixmap.fill(Qt::black) if @erase || erase
         pixpainter = Qt::Painter.new
         pixpainter.begin(@pixmap)
         if(@anaglyph_mode)
            # In 3D mode, let overlapping red & cyan lines appear white
            pixpainter.setRasterOp(Qt::OrROP)
            # draw complete, rotated right-hand and left-hand
            # images in cyan and red for anaglyphic glasses
            draw_planets(xsize,ysize,pixpainter,-1,Qt::cyan) # right eye image
            draw_planets(xsize,ysize,pixpainter,1,Qt::red) # left eye image
         else
            # draw image once
            draw_planets(xsize,ysize,pixpainter)
         end
         pixpainter.end
         # move the completed image to the screen
         bitBlt(@graphicPane,0,0,@pixmap)
         @total_time_hours += @time_step / 3600
      end
   end

   def perform_orbit_calc
      OrbitalPhysics::process_planets(@planet_list,@time_step,@symmetric)
      draw_image
   end

   def toggle_animation
      if(@timer != nil)
         @timer.stop
         @timer = nil
      else
         @total_time_hours = 0
         @timer = Qt::Timer.new(self)
         Qt::Object.connect(@timer, SIGNAL('timeout()'), self, SLOT('perform_orbit_calc()'))
         @timer.start(@anim_time,false)
         draw_image(true)
      end
      @runStopButton.setText((@timer == nil)?"Run":"Stop")
   end

   def load_comets()
      n = @cometComboBox.currentText.to_i
      1.upto(n) do |i|
         name = "comet#{i}"
         ca = rand(360) # angle in x-z plane
         cr = rand(4e11) + 4e11 # distance from sun
         pos = Cart3.new(cr * Math.sin(ca * CommonCode::ToRad),0,cr * Math.cos(ca * CommonCode::ToRad))
         # comet initial velocity
         v = (rand(200) + 100) * 50.0
         v = (i % 2 == 1)?-v:v
         vel = Cart3.new(0,v,0)
         comet = Planet.new(name,1e3,pos,vel,1e9)
         @planet_list << comet
      end
   end

   def load_orbital_data(data,sun_only = false)
      data = data.split("\n")
      data.shift # drop header line
      data.each do |line|
         fields = line.split(",")
         pos = Cart3.new(-fields[1].to_f,0,0)
         vel = Cart3.new(0,0,fields[4].to_f)
         planet = Planet.new(fields[0],fields[2].to_f,pos,vel,fields[3].to_f)
         @planet_list << planet
         break if sun_only
      end
   end

   def load_objects
      @planet_list = []
      sun_only = !@solarSystemCheckBox.isChecked
      load_orbital_data(SolarSystem::Data,sun_only)
      load_comets if @cometsCheckBox.isChecked
      draw_image(true)
   end

   # data can be easily read from a file
   # not used in this version

   def load_file(sun_only = false)
      fn = Qt::FileDialog.getOpenFileName( nil, nil, self)
      if(fn != nil)
         data = File.read(fn)
         load_orbital_data(data,sun_only)
      end
   end

   # close application cleanly

   def close(*x)
      @timer.stop if @timer != nil
      @app.exit
   end

   # mouse events

   def mouseMoveEvent (e)
      # rotate image by dragging mouse
      dx = (e.y - @mouse_press_x) / 2
      dy = (e.x - @mouse_press_y) / 2
      @rotx = @mouse_press_rx - dx
      @roty = @mouse_press_ry - dy
      draw_image(true)
   end

   def mousePressEvent (e)
      # set up to control rotation
      # by dragging mouse
      @mouse_press_rx = @rotx
      @mouse_press_ry = @roty
      @mouse_press_x = e.y
      @mouse_press_y = e.x
      draw_image(true)
   end

   def wheelEvent (e)
      # change drawing scale using mouse wheel
      v = e.delta.to_f
      v = 1.0 - (v/500.0)
      @drawing_scale *= v
      draw_image(true)
   end

   # control actions

   def stepButton_clicked(*k)
      toggle_animation if @timer
      perform_orbit_calc
   end

   def runStopButton_clicked(*k)
      toggle_animation
   end

   def solarSystemCheckBox_clicked(*k)
      load_objects
   end

   def cometsCheckBox_clicked(*k)
      load_objects
   end

   def timeStepComboBox_activated(*k)
      set_time_step
   end

   def anaglyphicCheckBox_clicked(*k)
      @anaglyph_mode = !@anaglyph_mode
      draw_image(true)
   end

   def pixelsSpinBox_valueChanged(*k)
      @min_draw_radius = pixelsSpinBox.text.to_i
      draw_image(true)
   end

   def cometComboBox_activated(*k)
      cometsCheckBox.setChecked(true)
      load_objects
      draw_image(true)
   end

   def trailsCheckBox_clicked(*k)
      @erase = !trailsCheckBox.isChecked
   end

   def symmCheckBox_clicked(*k)
      @symmetric = symmCheckBox.isChecked
   end

end # class Gravity

# launch Qt application

if $0 == __FILE__
   app = Qt::Application.new(ARGV)
   w = Gravity.new(app)
   app.mainWidget = w
   w.show
   app.exec
end