# This file is a part of TaCo: a morphosyntactic tagset converter
#
# License: GPL v. 3
# Author: Bartosz Zaborowski
#

require "#{File.dirname(__FILE__)}/tab_maker"
require "#{File.dirname(__FILE__)}/helpers"
require "#{File.dirname(__FILE__)}/corpus"
require "#{File.dirname(__FILE__)}/configuration"
require 'thread'

module ConverterLogic
NAME="c50 + morphoanalyser_correction"

class Leaf
  attr :klass, :confidence, :entries
  def initialize klass, confidence, entries
    @klass = klass
    @confidence = confidence
    @entries = entries
  end

  def classify t, colnam
    self
  end

  def to_s
    "#{@klass} (#{@confidence} / #{@entries})"
  end

  def pretty_print_tree depth, specials
    "#{@klass} (#{@entries.to_i}/#{((1.0-@confidence)*@entries).to_i})"
  end

  def most_frequent_leaf
    self
  end

  def depth
    0
  end
end

class Node
  attr :entries

  def initialize lines, strees, indent, i2val
    truelines = lines[lines.index{|x| !x.empty?}..lines.rindex{|x| !x.empty?}].map{|x| x[indent..-1]}

    subtrees = []
    sti = -1
    truelines.each do |l|
      if l[0] =~ /^[a-z]/
        sti += 1
        subtrees[sti] = []
      end
      subtrees[sti] << l
    end

    subtrees.first.first =~ /^([a-z_0-9-]+)/
    @attr = $~[1]

    @subnodes = {}
    subtrees.each do |st|
      i = 0
      vals = []
      if st.first[@attr.length+1] == "="
        st.first[@attr.length+3..-1] =~ /[0-9]+/
        vals = [$~[0]]
      elsif st.first[@attr.length+1..-1][0..3] == "in {"
        vals = []
        while st[i][@attr.length+5..-1].index("}").nil?
          vals += st[i][@attr.length+5..-2].split(",")
          i += 1
        end
        vals += st[i][@attr.length+5..st[i].index("}")-1].split(",")        
      else
        puts st.inspect
        raise
      end
      rest = st[i][st[i].rindex(":")..-1]

      if rest == ":"
        # real subtree
        nod = Node.new(st[i+1..-1], strees, 4, i2val)

      elsif rest =~ /: \[S([0-9]+)\]/
        # subtree delayed (e.g. [S1234])
        nod = strees[$~[1].to_i]

      else # rest == 'class (0/2)' or similar
        rest[1..-1] =~ /[0-9]+/
        klass = i2val[$~[0].to_i]
        rest =~ /\(([0-9\.]+)(\/[0-9\.]+)?\)/
        conf = 0
        entries = $~[1].to_f
        if $~[2].nil?
          conf = entries > 0.0 ? 1.0 : 0.0
        else
          conf = 1.0 - $~[2].to_f/entries.to_f
        end
        nod = Leaf.new(klass, conf, entries)

      end
      vals.each do |v|
        @subnodes[i2val[v.to_i]] = nod
      end
    end

    begin
	    @entries = @subnodes.values.reduce(0){|a,x| x.entries+a}
    rescue
	    puts "blad!"
	    puts lines
	    puts
	    puts @subnodes.values.map{|x| x.class.to_s}.inspect
	    puts subtrees.map{|x| x.inspect}
	    puts strees.map{|x| x.class.to_s}.inspect
	    raise
    end
    @most_frequent = @subnodes.values.max{|a,b| a.entries <=> b.entries}
  end

  def classify t, colnam
    return @most_frequent.classify t, colnam if @subnodes[t[colnam[@attr]]].nil?
    @subnodes[t[colnam[@attr]]].classify t, colnam
  end

  def to_s
    ret = [""]
    @subnodes.each_pair do |k,v|
      ret << "#{@attr} = #{k}: #{v.to_s.lines.to_a.map{|x| ":   "+x}.join("\n")[4..-1]}"
    end
    ret.join("\n")
  end

  def confidence
    @subnodes.values.reduce(0){|a, sn| a += sn.confidence * sn.entries}/@entries
  end

  def most_frequent_leaf
    @most_frequent.most_frequent_leaf
  end

private
  def print_subnode maxlen, k, v, specials, depth
    vlines = v.pretty_print_tree(depth-1, specials).lines.to_a.map{|x| ":   "+x}
    if v.is_a? Leaf
      vlines = "\t"+vlines.first[4..-1]
    else
      vlines = "\n"+vlines.join
    end
    "@#{@attr} = #{sprintf "% #{maxlen}s", k}:#{vlines}"
  end
public

  def pretty_print_tree depth, specials
    if depth == 0
      return "[subtree @#{@attr}, (#{@entries.to_i}/#{((1.0-confidence)*@entries).to_i}), most frequent: #{most_frequent_leaf.pretty_print_tree depth, specials}]"
    else
      ret = []
      defaults = []
      if @attr.start_with? "special"
        maxlen = ([9]+@subnodes.keys.map{|a| TabMaker.decode_special(@attr, specials, a).length}).max
        @subnodes.each_pair do |k,v|
          if v.entries == 0
            defaults << v
          else
            ret << print_subnode(maxlen, TabMaker.decode_special(@attr, specials, k), v, specials, depth)
          end
        end
      else
        maxlen = ([9]+@subnodes.keys.map{|a| a.length}).max
        @subnodes.each_pair do |k,v|
          if v.entries == 0
            defaults << v
          else
            ret << print_subnode(maxlen, k, v, specials, depth)
          end
        end
      end
      unless defaults.empty?
        ret << "@#{@attr} = #{sprintf "% #{maxlen}s", "<default>"}:\t#{defaults.first.most_frequent_leaf.klass}"
      end
      return ret.join("\n")
    end
  end

  def depth
    @subnodes.values.map{|x| x.depth+1}.max
  end
end

class DecisionTree
  attr :root
  def initialize lines, i2val
    idx = (lines.to_a.index{|x| x.start_with? "SubTree"} || 0) -1

    treelines = lines.to_a[0..idx].map{|x| x.rstrip}

    subtrees = []
    sti = 0

    if idx > 0
      lines.to_a[idx..-1].each do |l|
        next if l.strip.empty?
        if l.start_with? "SubTree"
          sti = l[l.index("[S")+2..l.index("]")-1].to_i
          subtrees[sti] ||= []
          next
        end
        subtrees[sti] << l.rstrip
      end
    end
    stcomp = []
    (0..subtrees.length-1).to_a.reverse.each do |i|
      next if subtrees[i].nil?
      stcomp[i] = Node.new(subtrees[i], stcomp, 0, i2val)
    end


    @root = Node.new(treelines, stcomp, 0, i2val)
  end

  def apply test, columnnames
    results = []
    colnam = {}
    columnnames.each_with_index{|c,i| colnam[c] = i}
    test.each do |t|
      nd = @root
      leaf = nd.classify t, colnam
      results << leaf.klass
    end
    results
  end

  def pretty_print_tree depth, specials
    @root.pretty_print_tree depth, specials
  end

  def depth
    @root.depth
  end
end

class BoostingTree
  def initialize treelines, i2val, tree_errors
    @trees = treelines.map{|tl| DecisionTree.new(tl, i2val)}
    if treelines.length == 1
      @weights = [1]
    else
      @weights = tree_errors.map{|x| 1}
    end
  end

  def apply test, columnnames
    results = @trees.map{|t| t.apply test, columnnames}
    ret = []
    test.each_with_index do |t,i|
      r = {}
      results.each_with_index{|x, ri| r[x[i]] ||= 0; r[x[i]] += @weights[ri] }
      ret << r.to_a.max{|a,b| a[1] <=> b[1]}[0]
    end
    return ret
  end

  def pretty_print_tree depth, specials
    @trees.map{|t| t.pretty_print_tree depth, specials}.join("\n\n")
  end

  def depth
    @trees.map{|x| x.depth}.max
  end
end

class Model
  attr :trainmem

  def initialize options = {}
    @val2i = {}
    @i2val = []
    @options = options
  end

  def train_c50 tempname_path, options
    c50 = IO.popen "#{Configuration.c50_path} #{options.to_a.map{|a,b| " -#{a}"+(b ? " #{b}" : "")}.join} -f #{tempname_path}", "r"
    treelines = [[]]
    tree_errors = []
    cases_number = 0
    treeno = 0
    intree = false
    ineval = false
    alllines = []
    while line = c50.gets
      alllines << line
      if intree && line.strip.start_with?("Evaluation on training data")
        line.strip =~ /\(([0-9]+) cases\)/
        cases_number = $~[1].to_f
        intree = false
        ineval = true if treeno > 0
        next
      end
      if ineval && line.strip.start_with?("boost")
        ineval = false
        next
      end
      if ineval && (line.strip =~ /^([0-9]+)[ \t]+[0-9]+ +([0-9]+)\(/)
        tree_errors[$~[1].to_i] = $~[2].to_f/cases_number
        next
      end
      if line.start_with?("-----  Trial ")
        treeno += 1
        treelines << []
        intree = false
        next
      end
      if intree
        treelines[treeno] << line.rstrip
      end
      if line.strip == "Decision tree:" 
        trainmem ||= Helpers::get_mem_stats c50.pid
        intree = true
      end
    end

    treelines = treelines[1..-1] if treelines.length > 1

    if treelines.find{|x| x.empty?}
      puts alllines.join("\n")
      raise
    end

    c50.close

    Helpers::delete [tempname_path+".names", tempname_path+".data", tempname_path+".tree", tempname_path+".tmp"]

    return BoostingTree.new(treelines, @i2val, tree_errors), trainmem
  end
  
  def write_c50 instances, desc, filename, comment = ""
    f = File.open filename+".names", "w"
    f.puts comment.lines.to_a.map{|x| "| #{x.strip}"}
    f.puts
    desc.last[1] << "-"
    desc.last[1].uniq!
    if desc.last[1].length > 1
      f.puts desc.last[1].map{|x| val_to_i x}.join(", ")+"."
    else
      f.puts (desc.last[1]+["--"]).map{|x| val_to_i x}.join(", ")+"."
    end
    f.puts
    desc[0..-2].each do |a|
      a[1] << "-"
      a[1].uniq!
      if a[1].length > 1
        f.puts "#{a[0]}: #{a[1].map{|x| val_to_i x}.join(",")}."
      else
        f.puts "#{a[0]}: #{(a[1]+["--"]).map{|x| val_to_i x}.join(",")}."
      end
    end
    f.close
    f = File.open filename+".data", "w"
    instances.each do |i|
      f.puts i.map{|x| val_to_i x}.join(",")
    end
    f.close
  end

  def val_to_i str, add = true
    if !@val2i[str] && add
      @val2i[str] = @i2val.length
      @i2val << str
    end
    return @val2i[str] || "-"
  end

  def i_to_val i
    @i2val[i]
  end

  def debug x
    Helpers::DEBUGSTREAM.puts x, Helpers::DEBUG
  end

  def info x
    Helpers::DEBUGSTREAM.puts x, Helpers::INFO
  end

  def special_inner_crossvalidation sents, tempname_path
    debug "   -> inner crossvalidation..."

    debug "        preparing training..."

    inner_opts = {:left_context_size => @options[:left_context_size], 
      :right_context_size => @options[:right_context_size],
      :orth_prefix_size => @options[:orth_prefix_size],
      :orth_suffix_size => @options[:orth_suffix_size]}

    columnnames, train = TabMaker.make_tab sents, inner_opts

    debug "        preparing desc..."

    inner_desc = []
    columnnames.each_with_index do |a, i|
      inner_desc << [a, train.map{|x| x[i]}.sort.uniq]
    end

    buckets = (0..@options[:folds]-1).to_a.map{|x| train[x*train.length/@options[:folds]..(x+1)*train.length/@options[:folds]-1]}

    slcs = @options[:special_left_context_size]
    srcs = @options[:special_right_context_size]
    sbps = @options[:special_base_prefix_size]
    sops = @options[:special_orth_prefix_size]
    sbss = @options[:special_base_suffix_size]
    soss = @options[:special_orth_suffix_size]

    toks_o = sents.map{|s| s.aligned_with.toks}.flatten
    toks = toks_o[-slcs..-1]+toks_o+[0..srcs-1]

    golden_buckets = (0..@options[:folds]-1).to_a.map{|x| toks[x*train.length/@options[:folds]..(x+1)*train.length/@options[:folds]+slcs+srcs-1]}
  
    error_orths = (0..slcs+srcs).to_a.map{|x| {}}
    error_bases = error_orths.map{|x| {}}
    error_cntr = 0
    error_base_prefixes = (0..sbps).to_a.map{|x| {}}
    error_orth_prefixes = (0..sops).to_a.map{|x| {}}
    error_base_suffixes = (0..sbss).to_a.map{|x| {}}
    error_orth_suffixes = (0..soss).to_a.map{|x| {}}
    tok_cntr = 0

  # inner crossvalidation
    @options[:folds].times do |k|
      debug "         inner fold #{k+1}/#{@options[:folds]}..."
      subtrain = []
      ((0..@options[:folds]-1).to_a-[k]).each{|x| subtrain += buckets[x]}
      subtest = buckets[k]

      write_c50 subtrain, inner_desc, tempname_path

      tree, mem = train_c50 tempname_path, @options[:classifier_options]

      result = tree.apply subtest, columnnames

      goldentest = golden_buckets[k]

      result.each_with_index do |r,i|
        next if goldentest[i+slcs].is_a? Ns
        tok_cntr += 1
        unless r == goldentest[i+slcs].tag
          error_cntr += 1
          (0..slcs+srcs).each do |ii|
            error_orths[ii][goldentest[i+ii].orth.downcase] ||= 0
            error_orths[ii][goldentest[i+ii].orth.downcase] += 1
            error_bases[ii][goldentest[i+ii].base.downcase] ||= 0
            error_bases[ii][goldentest[i+ii].base.downcase] += 1
          end
          (1..sbss).each do |ss|
            error_base_suffixes[ss][goldentest[i+slcs].base.downcase.suffix(ss)] ||= 0
            error_base_suffixes[ss][goldentest[i+slcs].base.downcase.suffix(ss)] += 1
          end
          (1..soss).each do |ss|
            error_orth_suffixes[ss][goldentest[i+slcs].orth.downcase.suffix(ss)] ||= 0
            error_orth_suffixes[ss][goldentest[i+slcs].orth.downcase.suffix(ss)] += 1
          end
          (1..sbps).each do |ss|
            error_base_prefixes[ss][goldentest[i+slcs].base.downcase.prefix(ss)] ||= 0
            error_base_prefixes[ss][goldentest[i+slcs].base.downcase.prefix(ss)] += 1
          end
          (1..sops).each do |ss|
            error_orth_prefixes[ss][goldentest[i+slcs].orth.downcase.prefix(ss)] ||= 0
            error_orth_prefixes[ss][goldentest[i+slcs].orth.downcase.prefix(ss)] += 1
          end
        end
      end
    end

    debug "        got #{error_cntr} errors (#{100.0*(tok_cntr.to_f-error_cntr.to_f)/tok_cntr.to_f}%), using #{@options[:threshold]} threshold"

    @special_orths = []
    error_orths.each_with_index do |orths, i|
      @special_orths[i] = orths.to_a.sort{|a,b| b[1] <=> a[1]}[0..orths.length*@options[:threshold]].map{|x| x[0]}
    end

    @special_bases = []
    error_bases.each_with_index do |bases, i|
      @special_bases[i] = bases.to_a.sort{|a,b| b[1] <=> a[1]}[0..bases.length*@options[:threshold]].map{|x| x[0]}
    end

    @special_orth_prefixes = []
    error_orth_prefixes.each_with_index do |orths, i|
      @special_orth_prefixes[i] = orths.to_a.sort{|a,b| b[1] <=> a[1]}[0..orths.length*@options[:threshold]].map{|x| x[0]}
    end

    @special_base_prefixes = []
    error_base_prefixes.each_with_index do |bases, i|
      @special_base_prefixes[i] = bases.to_a.sort{|a,b| b[1] <=> a[1]}[0..bases.length*@options[:threshold]].map{|x| x[0]}
    end

    @special_orth_suffixes = []
    error_orth_suffixes.each_with_index do |orths, i|
      @special_orth_suffixes[i] = orths.to_a.sort{|a,b| b[1] <=> a[1]}[0..orths.length*@options[:threshold]].map{|x| x[0]}
    end

    @special_base_suffixes = []
    error_base_suffixes.each_with_index do |bases, i|
      @special_base_suffixes[i] = bases.to_a.sort{|a,b| b[1] <=> a[1]}[0..bases.length*@options[:threshold]].map{|x| x[0]}
    end
  end

  def train sents, tempdir
    FileUtils.mkdir_p tempdir

    tempname = Helpers::mk_temp_name("train_c50_temp")
    tempname_path = File.join(tempdir, tempname)

    info "   -> training..."

    special_inner_crossvalidation sents, tempname_path

    # actual classification
    debug "      -> preparing training..."
    columnnames, train = TabMaker.make_tab sents, @options, @special_orths, @special_bases, @special_orth_prefixes, @special_base_prefixes, @special_orth_suffixes, @special_base_suffixes

    debug "      -> preparing desc..."

    desc = []
    columnnames.each_with_index do |a, i|
      desc << [a, train.map{|x| x[i]}.uniq]
    end


    write_c50 train, desc, tempname_path

    tt = Time.now

    debug "      -> training unigram..."
    @unigram = {}
    train.each{|t| @unigram[t[-1]] ||= 0; @unigram[t[-1]] += 1}

    train = nil
    GC.start

    debug "      -> training tree..."

    @tree, @trainmem = train_c50 tempname_path, @options[:classifier_options]

    return Time.now-tt
  end

  def apply sents
    return if @options.nil?

    info "   -> tagging..."

    columnnames, arr = TabMaker.make_tab sents, @options, @special_orths, @special_bases, @special_orth_prefixes, @special_base_prefixes, @special_orth_suffixes, @special_base_suffixes

    result = @tree.apply arr, columnnames

    asents = sents.map{|s| s.aligned_with}

    if Configuration.use_morphoanalyser
      asents.each {|s| s.analyse}
    end
    toks = asents.map{|s| s.toks}.flatten


    if Configuration.use_morphoanalyser
      unigram_fired = 0
      toks.each_with_index do |t, i|
        next if t.is_a? Ns
        if t.lexes.find{|l| l.tag == result[i]} || t.lexes.empty?
          t.selected_tag = result[i]
        else
          freqs = t.lexes.map{|x| [@unigram[x.tag] || 0, x.tag]}
          t.selected_tag = freqs.max{|a,b| a[0] <=> b[0]}[1]
          unigram_fired += 1
        end
      end
      debug "   -> unigram fired #{unigram_fired} times"
    else
      toks.each_with_index do |t, i|
        next if t.is_a? Ns
        t.selected_tag = result[i]
      end
    end

  end

  def pretty_print_tree depth
    @tree.pretty_print_tree depth, {"special_orth" => @special_orths, "special_base" => @special_bases, "special_base_prefix" => @special_base_prefixes, "special_orth_prefix" => @special_orth_prefixes, "special_base_suffix" => @special_base_suffixes, "special_orth_suffix" =>  @special_orth_suffixes, :opts => @options}
  end

  def tree_depth
    @tree.depth
  end
end

def self.prepare sents
end

def self.do_a_try sents_train, sents_test, options, tempdir
  return 0, 0, 0 if sents_test.empty?
  tt = Time.now
  m = Model.new options
  m.train sents_train, tempdir
  traintime = Time.now - tt

  tt = Time.now
  m.apply sents_test
  testtime = Time.now - tt

  return traintime, testtime, m.trainmem
end

def self.train_model sents, options, tempdir
  m = Model.new options
  m.train sents, tempdir
  m
end

def self.apply_model model, sents
  model.apply sents
end

end
