# coding: utf-8

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

# tagset converter v0.1 by Bartosz Zaborowski
#

require 'rexml/document'
require "#{File.dirname(__FILE__)}/helpers"
require "#{File.dirname(__FILE__)}/morpho_analyser"
require "#{File.dirname(__FILE__)}/configuration"

# Corpus implementation (with XCES reader)

class Sentence
  attr_accessor :toks
  attr :file
  attr_accessor :aligned_with

  @@mut = Mutex.new

  def initialize file, toks = []
    @toks = toks
    @toks.pop if @toks.last.is_a? Ns
    @toks.shift if @toks.first.is_a? Ns
    @file = file
  end

  def self.from_array arr, file
    self.new file, arr
  end

  def self.from_xml xml, file
    self.new(file).from_xml xml
  end

  def self.from_dom root, file
    self.new(file).from_dom root
  end

  def clone_untagged file
    s = Sentence.new file
    s.toks = @toks.map{|t| t.clone_untagged s}
    s.aligned_with = self
    @aligned_with = s
    s
  end

  def from_xml xml
    xml.scan(/(<tok[^>]*>(?:[^<]|<[^\/]|<\/[^t])+<\/tok>|<ns\/>)/) do |e|
      ent = e.first
      if ent == "<ns/>"
        @toks << Ns.new(self)
      else
        @toks << Tok.from_xml(ent, self)
      end
    end
    self
  end

  def from_dom root
    root.elements.each do |el|
      case el.name
      when "tok"
        @toks << Tok.from_dom(el, self)
      when "ns"
        @toks << Ns.new(self)
      end
    end
    return self
  end
  
  def to_xml
    "<chunk type=\"s\">\n"+toks.map{|x| x.to_xml}.join("\n")+"\n</chunk>"
  end

  def get_orths
    orths = []
    last_space = false
    @toks.each do |t|
      if t.is_a? Ns
        last_space = false
      else
        orths << " " if last_space
        orths << t.orth
        last_space = true
      end
    end
    orths
  end

  def analyse
    @original_toks = @toks
    orths = get_orths
    @toks = []
    toks = nil
    @@mut.synchronize {
      toks = MorphoAnalyser.analyse orths, @file.corpus.tagset
    }
    last_space = true
    toks.each do |t|
      if t[0] == :space
        last_space = true
        next
      end
      @toks << Ns.new(self) unless last_space
      @toks << Tok.from_bt_pairs(t[0], t[1], self)
      last_space = false
    end

    @toks.each_with_index do |t, i|
      if t.orth != @original_toks[i].orth
        $stderr.puts "po analizie nieuzgodnione:"
        $stderr.puts @original_toks.map{|x| x.orth || ""}.inspect
        $stderr.puts @toks.map{|x| x.orth || ""}.inspect
        $stderr.puts
        raise
      end
      t.golden_tag = @original_toks[i].tag
    end
  end

  def restore
    @toks = @original_toks if @original_toks
  end

  def split_interps
    ntoks = []
    @toks.each do |t|
      if t.tag == "interp" && t.orth =~ /^[-.+_,?()\[\]!)”]+$/
        t.orth.split(//).each do |o|
          ntoks << Tok.from_bt_pairs(o, [[o, "interp"]], self)
        end
      else
        ntoks << t
      end
    end
    @toks = ntoks
  end
end

class Entity
  attr_accessor :eos, :golden_tag
  attr_reader :selected_tag

  def select_all_interps
  end

  def selected_tag= x
    @selected_tag = x
  end
end
  
class Ns < Entity
  attr :sentence

  def initialize sentence
    @sentence = sentence
  end

  def clone_untagged s
    Ns.new s
  end

  def orth
    nil
  end

  def the_same other
    other.is_a? Ns
  end
  
  def ns_or_interp
    true
  end

  def to_s
    "#<Ns>"
  end

  def to_xml
    "<ns/>"
  end

  def base
    "ns"
  end
  
  def orth
    ""
  end

  def positional_tag
    sentence.file.corpus.tagset.make_positional "ns"
  end

  def tag
    "ns"
  end
end

class Tok < Entity
  attr_accessor :orth
  attr :lexes, :sentence

  def clone_untagged s
    t = Tok.new s
    t.orth = @orth
    t
  end

  def initialize sentence
    @lexes = []
    @sentence = sentence
  end

  def self.from_dom root, sent
    self.new(sent).from_dom root
  end

  def from_dom root
    @orth = REXML::XPath.first(root, "orth").text
    REXML::XPath.match(root, "lex").each do |lex|
      gold = lex.attributes["disamb"] == "1"
      base = lex.elements["base"].text
      tag = lex.elements["ctag"].text
      @lexes << Lex.new(base, tag, gold, self)
    end
    return self
  end

  def self.from_xml xml, sent
    self.new(sent).from_xml xml
  end

  def from_xml xml
    xml =~ /<orth>([^>]*)<\/orth>/
    @orth = $~[1].dexml_entity
    reg = /<lex ?(disamb="1"|disamb_sh="0")? ?(disamb="1"|disamb_sh="0")?><base>([^<]*)<\/base><ctag>([^<]+)<\/ctag>/
    xml.scan(reg) do |l|
      gold = !l[0..1].index("disamb=\"1\"").nil?
      base = l[2].dexml_entity
      tag = l[3]
      @lexes << Lex.new(base, tag, gold, self)
    end
    @golden_tag = tag
    self
  end

  def self.from_bt_pairs orth, btpairs, sent
    self.new(sent).from_bt_pairs orth, btpairs
  end

  def from_bt_pairs orth, btpairs
    @orth = orth
    btpairs.each do |b,t|
      @lexes << Lex.new(b, t, false, self)
    end
    self
  end

  def ns_or_interp
    @lexes.find{|x| x.tag != "interp"}.nil? && @lexes.length > 0
  end

  def the_same other
    other.is_a?(Tok) && other.orth.downcase == @orth.downcase
  end

  def to_s
    "#<Tok #{@orth} [#{@lexes.map{|x| x.to_s}.join(" ")}]>"
  end

  def to_xml
    "<tok>\n<orth>#{@orth.xml_entity}</orth>\n"+@lexes.map{|x| x.to_xml}.join("\n")+"\n</tok>"
  end

  def select_all_interps
    @lexes.each {|l| l.select}
  end

  def golden_lex
    @golden_lex ||= @lexes.find{|x| x.golden}
    @golden_lex
  end

  def base
    golden_lex.base
  end

  def refresh_golden_tag
    @golden_lex = nil
    @golden_tag = golden_lex ? golden_lex.tag : ""
  end

  def positional_tag
    golden_lex.positional
  end

  def tag
    golden_lex ? golden_lex.tag : ""
  end

  def selected_tag= x
    @selected_tag = x
    @lexes << Lex.new(@orth, x, false, self) if !x.nil? && @lexes.find{|y| y.tag == x}.nil?
    @lexes.each do |l|
      if l.tag == x
        l.select
      else
        l.deselect
      end
    end
  end
end

class Lex
  attr :base
  attr :tag
  attr :golden
  attr :selected
  attr :token
  
  def select
    @selected = true
  end

  def deselect
    @selected = false
  end

  def initialize base, tag, golden, token
    @base = base
    @tag = tag
    @golden = golden
    @selected = true
    @token = token
  end

  def to_s
    "#{base}:#{tag}:#{golden ? '1' : '0'}:#{selected ? '1' : '0'}"
  end

  def to_xml
    "<lex#{golden ? " disamb_orig=\"1\"" : ""}#{selected ? " disamb=\"1\"" : ""}><base>#{base.xml_entity}</base><ctag>#{tag}</ctag></lex>"
  end

  def positional
    token.sentence.file.corpus.tagset.make_positional tag
  end
end

class CorpusFile
  attr :path, :dir, :filename
  attr_accessor :sents, :aligned_with
  attr :corpus

  def initialize corpus, filename, dir, sents = []
    @filename = filename
    @sents = sents
    @corpus = corpus
    @dir = dir
  end

  def self.from_xml dir, filename, corpus
    self.new(corpus, filename, dir).from_xml
  end

  def self.from_sents sents, dir, filename, corpus
    self.new corpus, filename, dir, sents
  end

  def from_xml 
    @path = File.join(@dir, @filename)
    Helpers::DEBUGSTREAM.puts "loading #{@path}", Helpers::INFO
    File.open(@path) do |f|
      data = f.read.gsub("\n", "")
      data =~ /<chunkList[^>]*>(.*)<\/chunkList>/
      chunks = $~[1]
      chunks.scan(/(<chunk[^>]*>(?:[^<]|<[^\/c]|<\/[^c]|<\/c[^h]|<c[^h])+<\/chunk>)/) do |chunk|
        @sents << Sentence.from_xml(chunk.first, self)
      end
#      doc = REXML::Document.new f
#      REXML::XPath.match(doc, "//chunk").select{|x| x.attributes["type"] == "s"}.each do |sent|
#        @sents << Sentence.from_dom(sent, self)
#      end
    end
    return self
  end

  def clone_untagged filename, corp
    f = CorpusFile.new(corp, filename, @dir)
    f.sents = @sents.map{|s| s.clone_untagged f}
    f.aligned_with = self
    @aligned_with = f
    f
  end

private

#partyzantka
  def apply_action tok, act
    tok.orth = act.first
    tok.lexes.clear
    act.last.each do |a|
      tok.lexes << Lex.new(a[0], a[1], a[2], tok)
    end
    tok.refresh_golden_tag
    tok
  end

  def align_toks toks1, id1, toks2, id2, ntoks2, ntoks1
    Configuration.special_alignment_corrections.each_pair do |match, action|
      if toks1[id1..id1+match.first.length-1].map{|x| [x.orth, x.tag]} == match.first &&
        toks2[id2..id2+match.last.length-1].map{|x| [x.orth, x.tag]} == match.last
        if action[0]
          if action[0] != :copy
            apply_action toks1[id1], action[0]
          end
          ntoks1 << toks1[id1]
        end
        if action[1]
          if action[1] != :copy
            apply_action toks2[id2], action[1]
          end
          ntoks2 << toks2[id2]
        end
        return action[2..3]
      end
    end
    return false
  end

public

  def align other
    alltoks = sents.map{|s| r = s.toks; r.last.eos = true; r}.flatten
    othertoks = other.sents.map{|s| r = s.toks; r.last.eos = true; r}.flatten

    id1, id2 = 0,0
    nt1, nt2 = [], []
    while id1 < alltoks.length
      if alltoks[id1].the_same othertoks[id2]
        nt1 << alltoks[id1]
        nt2 << othertoks[id2]
        id1 += 1
        id2 += 1
      else
        if (x = align_toks alltoks, id1, othertoks, id2, nt2, nt1) || 
          (y = align_toks othertoks, id2, alltoks, id1, nt1, nt2)
          if x
            id1 += x[0]
            id2 += x[1]
          else #y
            id1 += y[1]
            id2 += y[0]
          end
        else
          return false
        end
      end
    end

    newsents = []
    newothersents = []
    sent = []
    othersent = []
    nt1.each_with_index do |t, i|
      sent << t
      othersent << nt2[i]
      if t.eos || nt2[i].eos
        newsents << Sentence.from_array(sent, self)
        sent = []
        newothersents << Sentence.from_array(othersent, other)
        othersent = []
      end
    end
    newsents << Sentence.from_array(sent, self) unless sent.empty?
    newothersents << Sentence.from_array(othersent, other) unless sent.empty?
    @sents = newsents
    other.sents = newothersents

    @aligned_with = other
    other.aligned_with = self

    sents.each_with_index do |s, i|
      s.aligned_with = other.sents[i]
      other.sents[i].aligned_with = s
    end
    true
  end

  def to_xml include_headers = true
    ret = ""
    ret += Corpus.file_header if include_headers
    ret += sents.map{|x| x.to_xml}.join("\n")
    ret += "\n"+Corpus.file_footer if include_headers
    ret
  end

  def save
    File.open(File.join(@dir, @filename), "w") do |f|
      f.write to_xml
    end
  end

end

class Tagset
  attr :attrs, :attrnames, :poss, :val2attr, :allvals, :filename, :analyser_conversions

  def initialize file
    @filename = file
    lines = File.open(file) {|f| f.lines.to_a}

    attrs = {}
    poss = {}
    poss_attrs_optionality = {}
    val2attr = {}
    analyser_conversions = {}

    state = :no
    lines.map{|l| l.strip}.each do |l|
      next if l[0] == "#" || l.empty?
      if l =~ /^\[([A-Z_]+)\]$/
        state = $~[1].to_sym
        next
      end
      if state == :ATTR
        a = l.split("=")[0].strip
        vs = l.split("=")[1].split
        attrs[a] = vs
        vs.each do |v|
          val2attr[v] = a
        end
      end
      if state == :POS
        p = l.split("=")[0].strip
        vs = l.split("=")[1] || ""
        poss_attrs_optionality[p] = vs.split.map{|x| x[0] == "[" && x[-1] == "]"}
        vs = vs.split.map{|x| x[0] == "[" && x[-1] == "]" ? x[1..-2] : x}
        attrs["pos"] ||= []
        attrs["pos"] << p
        poss[p] = vs
        val2attr[p] = "pos"
      end
      if state == :ANALYSER_CORRECTOR
        p = l.split("=")[0].strip
        convs = l.split("=")[1].split(",").map{|x| x.strip.split}
        analyser_conversions[p] ||= {}
        h = analyser_conversions[p]
        convs.each do |cp|
          h[cp.first] = cp.last
        end
      end
    end
    attrs["pos"] << "ns"
    val2attr["ns"] = "pos"
    @attrs = attrs
    @attrnames = attrs.keys.sort
    @poss = poss
    @poss_attrs_optionality = poss_attrs_optionality
    @val2attr = val2attr
    @allvals = val2attr.keys.sort
    @analyser_conversions = analyser_conversions
  end

  def correct_analyser tag
    tp = tag.split(":")
    pos = tp.first
    pos = "ign" if @poss[pos].nil?
    ntag = pos
    @poss[pos].each_with_index do |attr, i|
      next if @poss_attrs_optionality[pos][i] && tp[i+1].nil?
      if !@attrs[attr].include?(tp[i+1])
        v = nil
        if @analyser_conversions[attr]
          v = @analyser_conversions[attr][tp[i+1]]
        end
        unless v
          $stderr.puts "unknown tag from analyser: #{tag}"
          $stderr.puts "attribute: #{attr}, tagset has #{@attrs[attr].inspect}, corrector has #{@analyser_conversions[attr].inspect}"
          $stderr.puts "tagset: #{@filename}"
          raise
        end
        ntag += ":" + v
      else
        ntag += ":" + tp[i+1]
      end
    end
    ntag
  end

  def make_positional tag
    atvs = tag.split(":")
    ats = {}
    atvs.each {|v| ats[@val2attr[v]] = v}
    @attrnames.map{|x| ats[x]}
  end
  
  def all_possible_tags
    @all_possible if @all_possible

    @all_possible = []
    @poss.each_pair do |pos, ats|
      tags = [pos]
      ats.each do |a|
        nt = []
        @attrs[a].each do |av|
          nt += tags.map{|x| x+":"+av}
        end
        tags = nt
      end
      @all_possible += tags
    end
    @all_possible
  end
end

class Corpus
  attr :tagset
  attr_accessor :basedir, :aligned_with, :files

  def self.file_header
    "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n"+
      "<!DOCTYPE cesAna SYSTEM \"xcesAnaIPI.dtd\">\n"+
      "<cesAna xmlns:xlink=\"http://www.w3.org/1999/xlink\" type=\"pre_morph\" version=\"IPI-1.2\">\n"+
      "<chunkList>\n"
  end

  def self.file_footer
    "</chunkList>\n"+
      "</cesAna>\n"
  end

  def initialize tagset = nil
    @tagset = tagset
    @files = []
  end

  def clone_untagged filename, tagsetfile
    c = Corpus.new
    c.load_tagset tagsetfile
    c.files = @files.map {|f| f.clone_untagged filename, c}
    c.aligned_with = self
    @aligned_with = c
    c.basedir = @basedir
    c
  end

  def self.from_xml dir, filename, tagsetfile
    f = self.new.from_xml dir, filename
    f.load_tagset tagsetfile
    f
  end

  def self.from_sents sents, tagset
     x = self.new tagset
     x.files = [CorpusFile.from_sents(sents, "x", ".", x)]
     x
  end

  def from_xml dir, filename
    @basedir = dir
    wd = Dir.pwd
    Dir.chdir dir
    paths = (0..5).to_a.map{|x| Dir[File.join(["*"]*x+[filename])]}.flatten.sort
    paths.each do |p|
      @files << CorpusFile.from_xml(p[0..-filename.length-2], filename, self)
    end
    Dir.chdir wd
    return self
  end

  def to_xml
    Corpus.file_header + @files.map{|x| x.to_xml false}.join("\n")+"\n" + Corpus.file_footer
  end

  def align other_corpus
    otherf_lookup = {}
    other_corpus.files.each {|f| otherf_lookup[f.path] = f}
    thisf_lookup = {}
    files.each {|f| thisf_lookup[f.path] = f}
    @files = @files.select{|f| otherf_lookup[f.path] && f.align(otherf_lookup[f.path])}
    other_corpus.files.replace other_corpus.files.select{|f| thisf_lookup[f.path]}
    @aligned_with = other_corpus
    other_corpus.aligned_with = self
  end

  def select_all_interps
    all_toks.each {|t| t.select_all_interps}
  end

  def serialize filename
    File.open(filename, "w") do |f|
      f.write Marshal::dump(self)
    end
  end

  def self.deserialize filename
    corp = nil
    File.open(filename, "r") do |f|
      corp = Marshal::load f.read
    end
    corp
  end

  def split_interps
    all_sents.each {|s| s.split_interps}
  end

  def all_sents
    files.map{|x| x.sents}.flatten
  end

  def all_toks
    all_sents.map{|x| x.toks}.flatten
  end

  def load_tagset filename
    @tagset = Tagset.new filename
  end

  def clear_selected_tags
    all_toks.each {|t| t.selected_tag = nil}
  end

  def save
    wd = Dir.pwd
    Dir.chdir @basedir
    files.each {|f| f.save}
    Dir.chdir wd
  end
end





