##  Lattice Boltzmann sample, written in Ruby
##
##  Main author: Jean-Luc Falcone
##  Co-author: Jonas Latt
##  Copyright (C) 2006 University of Geneva
##  Address: Jean-Luc Falcone, Rue General Dufour 24,
##           1211 Geneva 4, Switzerland 
##
##  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., 51 Franklin Street, Fifth Floor,
##  Boston, MA  02110-1301, USA.
##

require 'lb_util.rb'

module LB

    # A lattice site (9 distribution functions + a collision object)
    class Site
      
        def initialize(initial_f, collision_operator)
          @f = initial_f.dup
          @collision_operator = collision_operator
        end #def initialize
        
        def [](i)
          @f[i]
        end #def [](f_i)

        def []=(i,rhs)
          @f[i] = rhs
        end #def []=(i,rhs)
        
        def collide
          @f = @collision_operator.update(@f)
        end #collide
        
        def rho()
          @collision_operator.rho(@f)    
        end #rho()
        
        def u()
          @collision_operator.u(@f)
        end #u()
            
    end #class Site


    # The full lattice, i.e. an array of lattice sites
    class Lattice

	include D2Q9

        # Constructor: builds up the array of lattice sites
        def initialize(l_x,l_y,bulk_collision_operator)
            @l_x = l_x
            @l_y = l_y
            @lattice = Array.new
            @temp_lattice = Array.new
            1.upto( @l_x ) {
                column = Array.new
                temp_column = Array.new
                1.upto( @l_y) {
                    column  <<  Site.new(T, bulk_collision_operator)
                    temp_column  <<  Site.new(T, bulk_collision_operator)
                }
                @lattice << column
                @temp_lattice << temp_column
            }
        end #initialize

        # Initialize a lattice site to an equilibrium distribution
        def ini_equilibrium(x,y, rho, u)
            Q.times { |i|
                @lattice[x][y][i] = equilibrium(i, rho, u)
            }
        end #ini_equilibrium

        # Attribute a collision (e.g. LBGK) to a lattice site
        def attribute_collision(x,y,collision_operator)
            @lattice[x][y] = Site.new(T,collision_operator) 
            @temp_lattice[x][y] = Site.new(T,collision_operator)         
        end #attribute_collision

        # Full iteration (collision+propagation) step
        def step
            @lattice.each { |column|
                column.each { |site|
                    site.collide
                }
            }

            propagation()
        end #step
        
        def propagation
            @lattice.each_with_index { |column,x|
                column.each_with_index { |site,y|
                    Q.times { |i|
                        prev_x = x-C[i][0]
                        prev_y = y-C[i][1]
                        if prev_x >= 0 && prev_y >= 0 && prev_x < @l_x && prev_y < @l_y then
                            @temp_lattice[x][y][i] = @lattice[prev_x][prev_y][i]
                        end
                    }
                }
            }
            @lattice,@temp_lattice = @temp_lattice,@lattice
        end
        
        # Extract a matrix with the velocity norm
        def norm_u
            u = Array.new
            @lattice.each_with_index { |column,x| 
                u_column = Array.new
                column.each_with_index { |site,y|
                    u_column << norm(site.u)
                }
                u << u_column
            }
            u
        end  

        # Extract a matrix with the density
        def rho
            rho = Array.new
            @lattice.each_with_index { |column,x| 
                rho_column = Array.new
                column.each_with_index { |site,y|
                    rho_column << site.rho
                }
                rho << rho_column
            }
            rho
        end

      end #class Lattice

end #module LB
