# -*- coding: utf-8 -*-

# Copyright (C) 2012 Adam Radziszewski. Part of IOBBER.
# This program is free software; you can redistribute and/or modify it
# under the terms of the GNU Lesser General Public License as published by the Free
# Software Foundation; either version 3 of the License, or (at your option)
# any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE.
#
# See the LICENCE, COPYING.LESSER and COPYING files for more details

import CRFPP # CRF++ Python wrapper
import subprocess, os # running crf_learn
import codecs

import config, corpio

DATA_SEP = '\t'

def open_tr_files(model_name, data_dir, channels):
	tr_files = {}
	for chan in channels:
		tr_files[chan] = codecs.open(corpio.f_name(model_name, data_dir,
					config.EXT_DATA, chan), 'wb', 'utf-8')
	return tr_files

def close_tr_files(tr_files):
	for chan in tr_files:
		tr_files[chan].close()

def write_example(tr_file, feat_vals, class_label):
	"""Writes a training example in simple tab-separated format."""
	tr_file.write(DATA_SEP.join(feat_vals))
	tr_file.write(DATA_SEP)
	tr_file.write(class_label)
	tr_file.write('\n')

def write_end_of_sent(tr_file):
	"""Writes end-of-sentence marker to the training file."""
	tr_file.write('\n')

def train_and_save(conf, model_name, config_dir, data_dir, chan_name):
	"""Trains a CRF classifier for the given chan_name. The trained model
	is saved to filenames (generated using model_name and conf)."""
	tr_fname = corpio.f_name(model_name, data_dir, config.EXT_DATA, chan_name)
	cr_fname = corpio.f_name(model_name, data_dir, config.EXT_CR, chan_name)
	cr_template = corpio.f_name(model_name, config_dir, config.EXT_TEXT, chan_name)
	crf_opts = conf.get(config.S_CLASSIFIER, config.O_PARAMS)
	# run crf_learn
	args = ['crf_learn', cr_template, tr_fname, cr_fname]
	args.extend(crf_opts.split()) # if any
	with open(os.devnull, 'w') as fnull:
		retval = subprocess.call(args,
			stdout = fnull, stderr = fnull)
		if retval != 0:
			raise IOError('Training CRF++ FAILED. Check .tab file for data validity. Call: %s' % ' '.join(args))

def load(conf, model_name, data_dir, chan_name):
	"""Tries to load a stored classifier.
	If doesn't exist, will return None."""
	cr_fname = corpio.f_name(model_name, data_dir, config.EXT_CR, chan_name)
	if os.path.isfile(cr_fname):
		return CRFPP.Tagger('-m %s' % cr_fname)
	raise IOError('can\'t open classifier from %s' % cr_fname)

def open_sent(crf_obj):
	"""
	Notify the trained classifier than a new sentence will be classified.
	"""
	crf_obj.clear()

def eat_token(crf_obj, feat_vals):
	"""Feed the trained classifier with a new token (instance). The output
	tag sequence for the sentence being processed will be reade after calling
	close_sent."""
	instance = DATA_SEP.join(feat_vals).encode('utf-8')
	crf_obj.add(instance)

def close_sent(crf_obj):
	"""Notify the trained classifier that a whole sentence has been fed and
	have the classifier classify each token."""
	crf_obj.parse()

def classify_token(crf_obj, tok_idx):
	"""Retrieve the class label (tag) for the token at given tok_idx. Assumes
	that a whole sentence has been fed to the trained crf_obj with open_sent,
	eat_token and close_sent calls."""
	return crf_obj.y2(tok_idx)
