#!/usr/bin/python

# make the clustering class-specific

import sys,os,re,glob,math,glob,signal,traceback
if "DISPLAY" not in os.environ:
    import matplotlib
    matplotlib.use("AGG")
from scipy.ndimage import interpolation
from pylab import *
from optparse import OptionParser
from multiprocessing import Pool
import ocrolib
from ocrolib import number_of_processors

signal.signal(signal.SIGINT,lambda *args:sys.exit(1))

class CostTooHigh:
    def __init__(self,total=None,total_thr=None,avg=None,avg_thr=None):
        self.total = total
        self.total_thr = total_thr
        self.avg = avg
        self.avg_thr = avg_thr
    def __str__(self):
        result = ""
        for k in "total total_thr avg avg_thr".split():
            v = getattr(self,k,None)
            if v is None: continue
            if result!="": result += " "
            result += "%s=%s"%(k,v)
        return result

class AlignmentFailed:
    pass

def safe_str(x):
    try:
        return str(x)
    except e:
        traceback.print_exc()
        return "<bad ASCII>"

class Log:
    def __init__(self):
        self.lines = []
    def __del__(self):
        self.output()
    def log(self,*args):
        args = [safe_str(x) for x in args]
        self.lines.append(" ".join(args))
    def error(self,*args):
        self.log("ERROR",*args)
    def output(self):
        if len(self.lines)>0:
            print "\n".join(self.lines)
            self.lines = []
    def __call__(self,*args):
        self.log(*args)

parser = OptionParser("""
usage: %prog [options] [text.txt langmod.fst image.png ...]

Performs recognition and optional alignment using the given classifier
and language models. The classifier should be an isolated character classifier.

Arguments can be a mix of text files, language models, and images.

If a language model is given, that's used for aligning/recognizing
subsequent images.

If a text file is given, it is compiled into a language model and
then used for recognizing subsequent images.

When alignment is performed, rseg.gt.png, cseg.gt.png,
and gt.txt files are written.
""")

# these options control alignment
parser.add_option("-m","--model",help="model file",default="unlv.model")
parser.add_option("-s","--segmenter",help="segmenter",default="DpSegmenter")
parser.add_option("-l","--langmod",help="language model",default=None)
parser.add_option("-C","--maxccost",help="maximum per-character cost added to lattice",type="float",default=10.0)
parser.add_option("-N","--nbest",help="use nbest classifier outputs only",type="float",default=10)
parser.add_option("-D","--maxdist",help="maxdist for grouper",type="int",default=5)
# this can be used to allow for ambiguous characters in alignment
parser.add_option("-p","--precomp",help="precompose extra transducer",default=None)
# parser.add_option("-A","--noambigs",help="don't use ambiguous classes",action="store_true")
parser.add_option("-M","--linemodel",help="character size model",default=None)

# file handling
parser.add_option("-S","--suffix",help="suffix for writing rseg/cseg files",default="gt")
parser.add_option("-O","--overwrite",help="overwrite rseg/cseg files",action="store_true")
parser.add_option("-x","--gtextension",help="extension used for ground truth (ex: .txt, .gt.txt, .fst,...)",default=None)

# rejection of alignments
parser.add_option("-t","--cthreshold",help="reject if average character cost is worse than this",type="float",default=7.0)
parser.add_option("-T","--gthreshold",help="reject if total line cost is worse than this",type="float",default=200.0)
# NB: the following option is not effective if -C is less than the given value
parser.add_option("-B","--badthreshold",help="reject if any character cost is worse than this",type="float",default=30.0)

# debugging
parser.add_option("-d","--display",help="display progress",action="store_true")
parser.add_option("-w","--wait",help="wait after each line",action="store_true")
parser.add_option("-v","--verbose",help="verbose",action="store_true")
parser.add_option("-Q","--parallel",type=int,default=number_of_processors(),help="number of parallel processes to use")
(options,args) = parser.parse_args()

if len(args)==0:
    parser.print_help()
    sys.exit(0)

suffix = options.suffix
if suffix!="" and suffix[0]!=".": suffix = "."+suffix

assert options.precomp is None,"precomp not implemented yet"

def error(*args):
    message = " ".join(args)
    print "ERROR:",message

def die(*args):
    error(*args)
    sys.exit(1)

segmenter = ocrolib.SegmentLine().make(options.segmenter)
grouper = ocrolib.StandardGrouper()
grouper.pset("maxdist",options.maxdist) # use 5 to handle "/''

ion()

print "loading",options.model

cmodel = ocrolib.load_component(options.model)
assert cmodel is not None
linerec = ocrolib.CmodelLineRecognizer(cmodel=cmodel,
                                      best=options.nbest,
                                      maxcost=options.maxccost)

default_lmodel_name = None
default_lmodel = None

if options.display:
    ion()
    show()
    
nlines = 0
goodlines = 0

def align1(t):
    log = Log()
    log("[[[")
    global default_lmodel_name,default_lmodel,nlines,goodlines
    print "***",t
    imagefile,lname = t
    if lname!=default_lmodel_name:
        default_lmodel_name = lname
        default_lmodel = ocrolib.read_lmodel_or_textlines(lname)
    nlines += 1
    prefix = re.sub(r'\.[^/.]*$','',imagefile)

    ## load the line image
    log("load\t",imagefile)
    image = ocrolib.read_image_gray(imagefile)
    if options.display:
        clf(); gray(); imshow(image); draw()

    ## perform line recognition, yielding a recognition lattice
    try:
        lattice,rseg = linerec.recognizeLineSeg(image)
    except:
        log.error(imagefile,"recognition failed")
        traceback.print_exc()
        return 0
    assert rseg.dtype=='i'
    # lattice.save("_lattice.fst")
    lattice.save(prefix+suffix+".fst")

    ## compute and output the raw best path
    s = lattice.bestpath()
    cost = 0.0
    log("lraw %6.2f\t%3d\t%s"%(cost,len(s),s))

    ## if we can find a language model for it then perform alignment
    lmodel = default_lmodel
    if options.gtextension is not None:
        base = re.sub(r'\.[^/.]*$','',imagefile)
        gtname = base+options.gtextension
        log("gt",gtname)
        if not os.path.exists(gtname):
            log.error(gtname,"not found")
            return 0
        try:
            lmodel = read_lmodel_or_textlines(gtname)
        except:
            log.error(gtname,"failed to load")
            return 0

    if lmodel is not None:
        try:
            r = ocrolib.compute_alignment(lattice,rseg,lmodel)
            result = r.output
            cseg = r.cseg
            costs = r.costs
            # print costs
            log("costs",sum(costs),mean(costs),amax(costs),median(costs))
            log("result",result)

            # compute various cost statistics and reject if they are too high
            tcost = sum(costs)
            if tcost>10000.0:
                raise AlignmentFailed()
            mcost = mean(costs)
            if mcost>options.cthreshold:
                raise CostTooHigh(avg=mcost,avg_thr=options.cthreshold)
            if tcost>options.gthreshold:
                raise CostTooHigh(total=tcost,total_thr=options.gthreshold)

            # check whether the output file exists; that's usually an error (either
            # the user left some previous outputs lying around, in which case we might
            # get a mix of old and new aligned files, or it's precious ground truth
            # that we shouldn't overwrite)
            rseg_file = prefix+".rseg"+suffix+".png"
            cseg_file = prefix+".cseg"+suffix+".png"
            if not options.overwrite:
                if os.path.exists(rseg_file): die(rseg_file,"already exists")
                if os.path.exists(cseg_file): die(cseg_file,"already exists")
                # don't bother checking for costs file

            # output the aligned segmentations
            ocrolib.write_line_segmentation(rseg_file,rseg)
            ocrolib.write_line_segmentation(cseg_file,cseg)
            with open(prefix+suffix+".costs","w") as stream:
                for i in range(len(costs)):
                    stream.write("%d %g\n"%(i,costs[i]))
            if options.gtextension is None:
                ocrolib.write_text(prefix+suffix+".txt",result)
            goodlines += 1
        except AlignmentFailed,e:
            log("FAIL alignment failed")
        except CostTooHigh,e:
            log("FAIL cost threshold exceeded: %s"%e)
        except Exception,e:
            log(traceback.format_exc())
    log("[",goodlines,nlines,goodlines*100.0/nlines,"]")
    if options.wait:
        raw_input()
    log("]]]")
    log.output()
    return 1

jobs = []
lmodel = None
for arg in args:
    if arg[-4:]==".fst" or arg[-4:]==".txt":
        assert options.gtextension is None,"either specify gtextension or language model"
        lmodel = arg
        continue
    jobs += [(arg,lmodel)]

if options.parallel<2:
    for arg in jobs:
        align1(arg)
else:
    pool = Pool(processes=options.parallel)
    result = pool.map(align1,jobs)
