#!/usr/bin/ruby

$: << 'fig/lib'
require 'execrunner'

# Note: run this on the NLP machines because SPARQL server is running on jack.
if ARGV.size == 0
  puts "Usage:"
  puts "  ./sempre @mode=train          [additional options] # Train semantic parser"
  puts "  ./sempre @mode=interact       [additional options] # Interactive mode for testing"
  puts
  puts "  ./sempre @mode=cacheserver    [additional options] # Start cache server"
  puts "  ./sempre @mode=sparqlserver   [additional options] # Start sparql server"
  puts
  puts "  ./sempre @mode=filterfreebase [additional options] # Filter raw Freebase RDF dump"
  puts "  ./sempre @mode=indexfreebase  [additional options] # Build Freebase index"
  puts
  puts "  ./sempre @mode=test           [additional options] # Run TestNG test suite"
  puts
  puts "Additional options can be any of the following:"
  puts "  - Additional program options (e.g., -BeamParser.beamSize 3)"
  puts "  - Execrunner options, which select the options to include (@data=0)"
  exit 1
end

system "mkdir -p state/execs"
system "touch state/lastExec"

def header
  l(
    letDefault(:q, 0), sel(:q, l(), l('fig/bin/q', '-shareWorkingPath', o('mem', '5g'), o('memGrace', 10), '-add', '---')),
    'fig/bin/qcreate',
    (File.exists?('/u/nlp/bin/java7') ? '/u/nlp/bin/java7' : 'java'),
    '-ea',
    '-Xmx15g',
    '-cp', 'classes:'+Dir['lib/*.jar'].join(':'),
    letDefault(:prof, 0), sel(:prof, l(), '-Xrunhprof:cpu=samples,depth=100,file=_OUTPATH_/java.hprof.txt'),
  nil)
end

def testHeader
  l(
    'java',
    '-ea',
    '-Xmx12g',
    '-cp', 'classes:'+Dir['lib/*.jar'].join(':'),
    'org.testng.TestNG',
  nil)  
end

$lexdir=6

def sparqlOpts
  l(
    required(:sparqlserver, 'host:port of the Sparql server'),
    o('SparqlExecutor.endpointUrl', lambda{|e| 'http://'+e[:sparqlserver]+'/sparql'}),
  nil)
end

$defaultFeatureDomains = [
  'basicStats',
  'alignmentScores',
  'tokensDistance',
  'context',
  'skipPos',
  'joinPos',
  #'bridgeBinaryMatch',
  'wordSim',
  #'lexAlign',
  'tokenMatch',
  #'rule',
  'opCount',
  'constant',
  'denotation',
  'whType',
  #'lemmaAndBinaries',
nil].compact

def defaultOpts
  l(
    o('execDir', '_OUTPATH_'), o('overwriteExecDir'),
    o('addToView', 0),

    selo(0, 'LexiconFn.lexiconClassName',
       'edu.stanford.nlp.sempre.fbalignment.lexicons.Lexicon',
       'edu.stanford.nlp.sempre.fbalignment.lexicons.WordNetExpansionLexicon',
       'edu.stanford.nlp.sempre.fbalignment.lexicons.GraphPropLexicon'),

    letDefault(:executor, 1),
    sel(:executor,
      l(),
      sparqlOpts,
    nil),
    selo(:executor, 'Builder.executor',
      'FormulaMatchExecutor', # Measure accuracy using logical forms (faster); good proxy
      'SparqlExecutor',       # Measure accuracy using answers (slower); executes SPARQL query on database
    nil),

    letDefault(:binarylex,0),
    sel(:binarylex,
      l(
        o('BinaryLexicon.binaryLexiconFilesPath','lib/fb_data/'+$lexdir.to_s+'/binaryInfoStringAndAlignment.txt'),
        o('BinaryLexicon.keyToSortBy','Intersection_size_typed'),
      nil),
      l(
        sel(nil,
          o('BinaryLexicon.binaryLexiconFilesPath','lib/fb_data/graphprop/1/jaccardLexicon','lib/fb_data/graphprop/1/graphpropLexicon'),
          o('BinaryLexicon.binaryLexiconFilesPath','lib/fb_data/graphprop/2/jaccardLexicon','lib/fb_data/graphprop/2/graphpropLexicon'),
        nil),
        o('BinaryLexicon.keyToSortBy','graphprop_estimate'),
      nil),
      l(
        o('BinaryLexicon.binaryLexiconFilesPath','lib/fb_data/graphprop/1/jaccardLexicon'),
        o('BinaryLexicon.keyToSortBy','jaccard'),
      nil),
    nil),

    # binary
    o('BinaryLexicon.useOnlyJaccard',false),

    # unary
    o('UnaryLexicon.unaryLexiconFilePath','lib/fb_data/'+$lexdir.to_s+'/unaryInfoStringAndAlignment.txt'),

  nil)
end

def unbalancedTrainDevSplit
  l(o('Dataset.trainFrac', 0.8), o('Dataset.devFrac', 0.2))
end
def balancedTrainDevSplit
  l(o('Dataset.trainFrac', 0.5), o('Dataset.devFrac', 0.5))
end

def cachePaths(lexiconFnCachePath, sparqlExecutorCachePath)
  l(
    required(:cacheserver, 'none (don\'t cache to disk), local (write to local file), or <hostname>:<port> (hit the cacheserver)'),
    lambda { |e|
      cacheserver = e[:cacheserver]
      case cacheserver
      when 'none' then l()
      when 'local' then l( # Use files directly - don't run more than one job that does this!
        o('Lexicon.cachePath', 'LexiconFn.cache'),
        o('SparqlExecutor.cachePath', 'SparqlExecutor.cache'),
      nil)
      else l(
        o('Lexicon.cachePath', cacheserver + ':' + lexiconFnCachePath),
        o('SparqlExecutor.cachePath', cacheserver + ':' + sparqlExecutorCachePath),
      nil)
      end
    },
  nil)
end

# tag is either "free917" or "webquestions"
def emnlp2013AblationExperiments(tag)
  l(
    letDefault(:ablation, 0),
    # Ablation experiments (EMNLP)
    sel(:ablation,
      l(), # (0) Just run things normally
      selo(nil, 'BeamParser.beamSize', 10, 50, 200), # (1) Vary beam size
      selo(nil, 'Dataset.trainFrac', 0.1, 0.2, 0.4, 0.6), # (2) Vary training set size
      sel(nil, # (3) Structural: only do join or only do bridge
        o('Grammar.tags', l(tag, 'join')),
        o('Grammar.tags', l(tag, 'bridge')),
      nil),
      sel(nil, # (4) Features
        o('FeatureExtractor.featureDomains', *($defaultFeatureDomains+['lexAlign'])), # +lexAlign
        o('FeatureExtractor.featureDomains', *($defaultFeatureDomains+['lexAlign']-['alignmentScores'])), # +lexAlign -alignmentScores
        o('FeatureExtractor.featureDomains', *($defaultFeatureDomains-['denotation'])), # -denotation
        o('FeatureExtractor.featureDomains', *($defaultFeatureDomains-['skipPos', 'joinPos'])), # -syntax features (skipPos, joinPos)
      nil),
    nil),

    letDefault(:split, 0), selo(:split, 'Dataset.splitRandom', 1, 2, 3),
  nil)
end

def free917
  l( # Data
    letDefault(:data, 0),
    sel(:data,
      l(o('Dataset.inPaths', 'train,data/free917.train.examples.canonicalized.json'), unbalancedTrainDevSplit), # (0) train 0.8, dev 0.2
      l(o('Dataset.inPaths', 'train,data/free917.train.examples.canonicalized.json', 'test,data/free917.test.examples.canonicalized.json')), # (1) Careful, this is test set!
    nil),

    # Grammar
    o('Grammar.inPaths', 'data/emnlp2013.inexactent.grammar'),
    o('Grammar.tags', 'free917', 'bridge', 'join'),

    # Features
    o('FeatureExtractor.featureDomains', *$defaultFeatureDomains),

    emnlp2013AblationExperiments('free917'),

    l(
      o('EntityLexicon.exactMatchIndex','lib/lucene/4.4/free917/'),
      o('Lexicon.entitySearchStrategy','exact'),
      cachePaths('LexiconFn.cache', 'SparqlExecutor.cache'),
    nil),
  nil)
end

def webquestions
  l(
    # Data
    letDefault(:data, 0),
    sel(:data,
      l( # Webquestions (dev)
        o('Dataset.inPaths',
          'train,lib/data/webquestions/dataset_11/webquestions.examples.train.json'),
        unbalancedTrainDevSplit,
      nil),
      l( # Webquestions (test)
        o('Dataset.inPaths',
          'train,lib/data/webquestions/dataset_11/webquestions.examples.train.json',
          'test,lib/data/webquestions/dataset_11/webquestions.examples.test.json'),
      nil),
    nil),

    # Grammar
    o('Grammar.inPaths', 'data/emnlp2013.inexactent.grammar'),
    o('Grammar.tags', 'webquestions', 'bridge', 'join'),

    o('BeamParser.beamSize', 200),

    # Features
    o('FeatureExtractor.featureDomains', *($defaultFeatureDomains+['lexAlign'])),

    emnlp2013AblationExperiments('webquestions'),

    # Caching
    cachePaths('LexiconFn.cache', 'SparqlExecutor.cache'),

    o('EntityLexicon.inexactMatchIndex','lib/lucene/4.4/inexact/'),
    o('EntityLexicon.maxEntries',50),

    # Learner reward
    o('Learner.partialReward', true),
  nil)
end

def jayant
  l(
    # Data
    letDefault(:data, 1),
    sel(:data,
      l(o('Dataset.inPaths', 'train,data/jayant-emnlp2012-validation.examples.json'), balancedTrainDevSplit), 
      l(o('Dataset.inPaths', 'train,data/jayant-emnlp2012-validation.examples.json','test,data/jayant-emnlp2012-test.examples.json')), 
    nil),

    # Grammar
    o('Grammar.inPaths', 'data/emnlp2013.inexactent.grammar'),
    o('Grammar.tags', 'webquestions', 'bridge', 'join'),

    o('BeamParser.beamSize', 500), 

    # Caching
    cachePaths('LexiconFn.cache', 'SparqlExecutor.cache'),

    o('EntityLexicon.inexactMatchIndex','lib/lucene/4.4/inexact/'),
  nil)
end

def selectDomain
  l(
    letDefault(:domain, 0),
    sel(:domain, {
      'none' => l(),
      'webquestions' => webquestions,
      'free917' => free917,
      'jayant' => jayant,
    }),
  nil)
end

def train
  l(
    header,
    'edu.stanford.nlp.sempre.Main',
    defaultOpts, selectDomain,
    o('Learner.maxTrainIters', 4),
  nil)
end

def interact
  l(
    # After training, run interact, which loads up a set of parameters and
    # puts you in a prompt.
    system('which rlwrap') ? 'rlwrap' : nil,
    header,
    'edu.stanford.nlp.sempre.Main',
    defaultOpts, selectDomain,
    o('Dataset.inPaths'),
    o('Learner.maxTrainIters', 0),
    required(:load, 'none or exec number (e.g., 812) to load'),
    lambda { |e|
      if e[:load] == 'none' then
        l()
      else
        execPath = "lib/models/#{e[:load]}.exec"
        l(
          o('Builder.inParamsPath', execPath+'/params'),
          o('Grammar.inPaths', execPath+'/grammar'),
          o('Master.logPath', lambda{|e| 'state/' + e[:domain] + '.log'}),
          o('Master.newExamplesPath', lambda{|e| 'state/' + e[:domain] + '.examples'}),
          o('Master.onlineLearnExamples', true),
          # Make sure features are set properly!
        nil)
      end
    },
    o('Main.interactive'),
    sel(:executeTopOnly,
      l(),
      l(
        o('Parser.executeTopFormulaOnly'),
        a('FeatureExtractor.disableDenotationFeatures', true),
      nil),
    nil),
  nil)
end

# Start the cache server that serves files with key-value maps.
def cacheserver
  l(
    'java', '-Xmx15g', '-ea', '-cp', 'classes:lib/fig.jar',
    'edu.stanford.nlp.sempre.StringCacheServer',
    o('port', 4000),
  nil)
end

############################################################
# Freebase RDF database

# Scratch directory
def scrOptions
  letDefault(:scr, '/u/nlp/data/semparse/rdf/scr/' + `hostname | cut -f 1 -d .`.chomp)
end

# (1) Filter RDF Freebase dump (do this once) [takes about 1 hour]
def filterfreebase
  l(
    scrOptions,
    l(
      'fig/bin/qcreate', o('statePath', lambda{|e| e[:scr] + '/state'}),
      'java', '-ea', '-Xmx20g', '-cp', 'classes:lib/*',
      'edu.stanford.nlp.sempre.freebase.FilterFreebase',
      o('inPath', '/u/nlp/data/semparse/scr/freebase/freebase-rdf-2013-06-09-00-00.canonicalized'),
      sel(:keep, {
        'all' => o('keepAllProperties'),
        'geo' => l(
          o('keepTypesPaths', 'data/geo.types'),
          o('keepPropertiesPath', 'data/geo.properties'),
          o('keepGeneralPropertiesOnlyForSeenEntities', true),
        nil),
      }),
      o('execDir', '_OUTPATH_'), o('overwriteExecDir'),
    nil),
  nil)
end

# (2) Start the SPARQL server (do this every time).
def sparqlserver
  l(
    scrOptions,
    letDefault(:exec, 93),
    'scripts/virtuoso', 'start',
    lambda{|e| 'lib/freebase/'+e[:exec].to_s+'.exec/vdb'}, # DB directory
    lambda{|e| 3000+e[:exec]}, # port
  nil)
end

# (3) Index the filtered RDF dump [takes 48 hours]
# Afterwards, stop the server and copy the exec directory to
# /u/nlp/data/semparse/scr/freebase/state/execs (this will serve as the master
# copy).
def indexfreebase
  l(
    letDefault(:stage, nil),
    scrOptions,
    required(:exec),
    sel(:stage,
      l(
        'scripts/virtuoso', 'add',
        lambda{|e| e[:scr]+'/state/execs/'+e[:exec].to_s+'.exec/0.ttl'}, # ttl file
        lambda{|e| 3000+e[:exec]}, # port
        lambda{|e| e[:offset] || 0}, # offset
      nil),
      l(
        'scripts/extract-freebase-schema.rb',
        lambda{|e| 'http://localhost:'+(3000+e[:exec]).to_s+'/sparql'}, # port
        lambda{|e| e[:scr]+'/state/execs/'+e[:exec].to_s+'.exec/schema.ttl'},
      nil),
    nil),
  nil)
end

# Query a single logical form or SPARQL
def query
  l(
    'java', '-ea',
    '-cp', 'classes:'+Dir['lib/*.jar'].join(':'),
    'edu.stanford.nlp.sempre.SparqlExecutor',
    sparqlOpts,
  nil)
end

def test
  l(
    testHeader,
    lambda { |e|
      if e[:class]
        l('-testclass', 'edu.stanford.nlp.sempre.test.' + e[:class])
      else
        'testng.xml'
      end
    },
  nil)
end

run!(
  sel(:mode, {
    'train' => train,
    'interact' => interact,
    'cacheserver' => cacheserver,
    'filterfreebase' => filterfreebase,
    'indexfreebase' => indexfreebase,
    'sparqlserver' => sparqlserver,
    'query' => query,
    'test' => test,
  }),
nil)
