require 'rubygems'
require 'gsl'

include GSL

class NodeAccessor
  
  attr_reader :cpm
  
  def initialize(cpm)
    @cpm = cpm
  end
  
  def is(*labels)
    ChildAccessor.new(cpm, *labels)
  end
  
  def is_not(*labels)
    ChildAccessor.new(cpm, *(cpm.child.labels - labels))
  end
end

class ChildAccessor
  
  attr_reader :cpm, :labels
  def initialize(cpm, *labels)
    @cpm, @labels = cpm, labels
  end
  
  def given(parent_name)
    ParentAccessor.new(cpm, labels)
  end
  
end

class ParentAccessor
  
  attr_reader :cpm, :node, :child_labels, :child_indices
  def initialize(cpm, child_labels)
    @cpm = cpm
    @child_labels = child_labels
    @node = cpm.parent
    @child_indices = child_labels.map {|label| @cpm.child.labels.index(label)}
  end
  
  def is(*labels)
    indices = labels.map {|label| get_index(label) }
    sum_probabilities(indices)
  end
  
  def is_not(*labels)
    not_indices = labels.map {|label| get_index(label) }
    indices = (0..node.labels.size).to_a - not_indices
    sum_probabilities(indices)
  end
  
  protected
  
    # TODO: Not right...
    def sum_probabilities(indices)
      first_child = child_indices.first
      cpm.matrix[indices.first, first_child]
      # indices.inject(0.0) do |s, i|
      #   s += cpm.matrix[i, first_child]
      # end
    end
    
    def get_index(label)
      node.labels.index(label)
    end
    
end


class ConditionalProbabilityMatrix
  
  class << self
    def define_node_accessor(node, cpm)
      define_method(node.name.to_sym) do
        NodeAccessor.new(cpm)
      end
    end
  end
  
  attr_reader :parent, :child, :matrix
  
  def initialize(parent, child)
    @parent, @child = parent, child
    @matrix = @parent.probabilities.col * @child.probabilities
    assert_name_access
  end
  
  def probability(opts={})
    child_label = opts[:child]
    parent_label = opts[:parent]
    raise ArgumentError, "Must provide a child and parent label.  E.g., probability(:child => true, :parent => false)" unless child_label and parent_label
    child_label_index = get_index(child, child_label)
    parent_label_index = get_index(parent, parent_label)
    self.matrix[parent_label_index, child_label_index]
  end
  alias :p :probability
  
  def inspect
    "ConditionalProbabilityMatrix: #{matrix.to_a.inspect}"
  end
  
  protected
  
    def assert_name_access
      ConditionalProbabilityMatrix.define_node_accessor(child, self)
    end
    
    def get_index(node, label)
      node.labels.index(label)
    end
end

class CPM < ConditionalProbabilityMatrix
end