#! /usr/bin/python
# -*- coding: utf-8 -*-

import sys
import os

BASE_DIR = os.path.abspath( os.path.join( os.path.dirname( __file__ ), ".." ) )
path     = os.path.abspath( os.path.join( BASE_DIR, "python" ) )
sys.path.append( path )

import glob
import momo
import random
import numpy as np
from optparse import OptionParser

def list_callback( option, opt_str, value, parser ):
  setattr( parser.values, option.dest, [int( v ) for v in value.split( "," )] )

# Process command line
parser = OptionParser( usage = "usage: %prog [options] dataset features algorithm [params...]" )
parser.add_option( 
  "--ids", help = "Trajectory ids to learn from", type = "string", 
  action = "callback", dest = "ids", callback = list_callback, default = []
)
parser.add_option( 
  "--delta", help = "Grid resolution", type = "float", 
  dest = "delta", default = 0.25
)
parser.add_option( 
  "--path", help = "Data directory", type = "string", 
  dest = "root", default = os.path.join( BASE_DIR, "data" )
)

options, args = parser.parse_args()

if len( args ) < 2:
  parser.error( "Incorrect number of arguments" )
  exit( 1 )

dataset        = args[0]
feature_name   = args[1]
algorithm_name = args[2]
algorithm_args = args[3:]

delta = options.delta
ids  = options.ids
root = options.root

# Prepare file paths
start_path = os.path.join( root, "models", dataset )
if not os.path.exists( start_path ):
  os.makedirs( start_path )

output_name  = os.path.join( start_path, "%s-%s" % ( feature_name, algorithm_name ) )
if len( algorithm_args ) > 0:
  output_name += "-%s" % "_".join( algorithm_args )
output_name += ".txt"

# Perform learning
file_path = os.path.join( root, "filtered", "%s*.txt" % dataset )
files = glob.glob( file_path )

data = []

for f in files:
  data.append( momo.read_data( os.path.join( f ) ) )
features  = momo.features.__dict__[feature_name]
algorithm = momo.learning.__dict__[algorithm_name]

a = momo.irl_assembler( features, algorithm, *algorithm_args, data = data, ids = ids, delta = delta )
f = open( output_name, "w" )
a.save( f )
f.close()

