#!/usr/bin/env python

from __future__ import division
import optparse, sys, subprocess, os, shutil
from math import *
import random
import re, glob

# begin configuration

#pcode = "/home/stud1/epschern/misc/code/pcode/pcode"
pcode = "/extra/debian/work/thesis/pcode/pcode"

# end configuration

class Class:
    def __init__(self):
        self.pics = []
        self.codes = []
        self.path = None

def exe(command):
    print " ".join(['"' + x + '"' for x in command])
    p = subprocess.Popen(command,
        shell = False,
        stdout = subprocess.PIPE,
        stderr = subprocess.STDOUT)

    p.wait()
    return p.stdout.read()

stderr = ""
def exe_visible(command):
    global stderr
    print " ".join(['"' + x + '"' for x in command])
    p = subprocess.Popen(command,
        shell = False,
        stdout = subprocess.PIPE,
        stderr = subprocess.PIPE)
    p.wait()
    stderr = p.stderr.read()
    return p.stdout.read()

def extract_options():
    max_id = 0
    sub_dirs = ["L", "R"]
    min_pics = 0
    if options.extract_only:
        eo = options.extract_only.split(",")
        max_id = int(eo[0])
        sub_dirs = [x for x in eo[1]]
        min_pics = int(eo[2])
    return max_id, sub_dirs, min_pics

def extract_textures(casia_dir):
    max_id, sub_dirs, min_pics = extract_options()
    dirs = glob.glob(os.path.join(casia_dir, "*"))
    pics = []
    for d in dirs:
        dn = d[d.rfind("/") + 1:]
        try:
            id = int(dn)
        except ValueError:
            continue
        if max_id > 0 and id > max_id: continue
        for sub in sub_dirs:
            spics = glob.glob(os.path.join(casia_dir, d, sub, "*.jpg"))
            if min_pics == 0 or len(spics) >= min_pics:
                pics += spics
    for pic in pics:
        command = [pcode]
        command += ["-i", pic]
        if not options.no_lids:
            command += ["-l"]
        command += ["-w", "finaltexture mask circles"]
        command += ["-p", options.output]
        exe(command)

def mask_from_texture(t):
    x = t.find("_finaltexture.tiff")
    name = t[:x]
    m = name + "_mask.tiff"
    return m

def textures_to_classes(d):
    max_id, sub_dirs, min_pics = extract_options()
    files = glob.glob(d + "/*_finaltexture.tiff")
    classes = {}
    for t in files:
        
        m = mask_from_texture(t)
        cname = m[:-len("_mask.tiff") - 2]
        
        slash = cname.rfind("/")
        if slash >= 0: cname = cname[slash + 1:]
        
        # S1000L
        if max_id > 0 and int(cname[2:5]) > max_id: continue
        if not cname[5] in sub_dirs: continue
    
        c = classes.get(cname, Class())
        c.path = cname
        c.pics += [(t, m)]
        classes[cname] = c
    if min_pics == 0: return classes
    classes2 = {}
    for k, v in classes.items():
        if len(v.pics) >= min_pics:
            classes2[k] = v
    return classes2

def codes_to_classes(d):
    max_id, sub_dirs, min_pics = extract_options()
    files = glob.glob(d + "/*.code")
    classes = {}
    for code in files:
        cname = code
        x = cname.rfind("/")
        if x >= 0: cname = cname[x + 1:]
        x = cname.rfind("_")
        if x >= 0: cname = cname[:x]
        
        cname = cname[:-2]

        # S1001L01
        if max_id > 0 and int(cname[2:5]) > max_id: continue
        if not cname[5] in sub_dirs: continue
    
        c = classes.get(cname, Class())
        c.path = cname
        c.codes += [code]
        classes[cname] = c
    if min_pics == 0: return classes
    classes2 = {}
    for k, v in classes.items():
        if len(v.codes) >= min_pics:
            classes2[k] = v
    return classes2

def get_info():
    info = ""
    if options.name: info += options.name
    return info

def generate_texture_transformed_with_key_per_class(classes):    
    for c in classes:
        x = c.rfind("/") 
        cname = c[x + 1:]
        if options.no_transform:
            cpath = "%s/%s_copied" % (options.output, cname)
        else:
            s = random.randint(0, 2**32 - 1)
            if options.permutate:
                cpath = "%s/%s_%0x08x_%s" % (
                    options.output, cname, s, options.grid)
            else:
                cpath = "%s/%s_%0x08x_%s_%s_%s" % (
                    options.output, cname, s, options.size, options.xamount,
                    options.yamount)
        classes[c].path = cpath

        if not os.path.isdir(cpath):
            os.mkdir(cpath)
        for t, m in classes[c].pics:
            for pic in [t, m]:
                name = warped_name(pic)
                if options.no_transform:
                    shutil.copyfile(pic, cpath + "/" + name)
                elif options.permutate:
                    permutate(pic, cpath, name, s)
                else:
                    warp(pic, cpath, name, s)
    return classes

def base_name(pic, extra = ""):
    x = pic.rfind("/")
    y = pic.rfind(".")
    return pic[x + 1:y] + extra

def generate_exported_txt(classes, output):
    name = options.name
    if name == None: name = ""
    ex = file("%s/exported%s.txt" % (output, name), "w")
    tn = 0
    if options.training: tn = int(options.training)
    for c in classes.values():
        i = 0
        ex.write("%s:\n" % c.path)
        for code in c.codes:
            if i == tn: ex.write("-\n")
            i += 1
            ex.write("%s\n" % code)
        if i <= tn:
            ex.write("-\n")
    ex.close()
            
def generate_codes(classes, name = ""):
    if name == None: name = ""
    for c in classes.values():
        for t, m in c.pics:
            x = t.rfind("_finaltexture")
            y = t.rfind("/")
            base = t[y + 1:x]
            command = [pcode]
            command += ["-t", t]
            command += ["-m", m]

            if options.algorithm == "lg":
                command += ["-a", "lg"]
            elif options.algorithm == "d42":
                command += ["-a", "d42"]
            else:
                command += ["-a", "qsw"]
            if not options.no_lids:
                command += ["-l"]
            what = "code"
            if options.keep_files: what += " algocode"
            if options.half: what += " half"
            command += ["-w", what]
            command += ["-p", options.output]
            command += ["-b", base]
            if options.half: command += ["-H"]
            exe(command)

def rates(intra, inter):
    """
    intra: 512 bins with matching score
    inter: 512 bins with matching score
    
    The matching score is the not normalized number of samples falling
    into the bin.
    """
    FMR = []
    prev = 0
    intern = 0
    for i in range(512):
        # Assume, we use i as treshold for saying two codes match. Then
        # the FMR is how many times inter results are below that threshold.
        FMR.append(inter[i] + prev)
        prev = FMR[i]
        intern += inter[i]
    
    FNMR = []
    prev = 0
    intran = 0
    for i in range(512):
        # Assume, we use i as treshold for saying two codes match. Then
        # the FNMR is how many times intra results are above that threshold.
        FNMR.append(intra[511 - i] + prev)
        prev = FNMR[i]
        intran += intra[511 - i]
    FNMR.reverse()

    # Normalize the sample counts to % of the total samples.
    for i in range(512):
        FMR[i] *= 100.0 / intern
        FNMR[i] *= 100.0 / intran
    
    # Find EER
    tmin, dmin = None, None
    for t in range(512):
        d = abs(FMR[t] - FNMR[t])
        if tmin == None or d < dmin:
            tmin, dmin = t, d
    EER = (FMR[tmin] + FNMR[tmin]) / 2.0

    return FNMR, FMR, EER

def verify(name = ""):
    if not name: name = ""
    command = [pcode]
    command += ["-f", "%s/exported%s.txt" % (options.verify, name)]
    command += ["-w", "verify"]
    if not options.no_lids:
        command += ["-l"]
    rof = "16"
    if options.rotation_offsets: rof = options.rotation_offsets
    command += ["-r", rof]
    output = exe_visible(command)
    print output
    intra = []
    inter = []
    for line in output.splitlines():
        if line.startswith("intra: "):
            intra = [int(x) for x in line.split()[1:]]
        elif line.startswith("inter: "):
            inter = [int(x) for x in line.split()[1:]]
    if len(intra) != 512 or len(inter) != 512:
        sys.stderr.write(output)
        sys.stderr.write(stderr)
        sys.stderr.write("Could not read distribution. Aborting.\n")
        sys.exit(1)
    return intra, inter

def identify(name = ""):
    if not name: name = ""
    command = [pcode]
    command += ["-f", "%s/exported%s.txt" % (options.identify, name)]
    command += ["-w", "identify"]
    if not options.no_lids:
        command += ["-l"]
    rof = "16"
    if options.rotation_offsets: rof = options.rotation_offsets
    command += ["-r", rof]
    output = exe_visible(command)
    print output

def results(intra, inter):
    name = options.name
    if not name: name = ""
    # normalize
    intran = 0.0
    intern = 0.0
    for i in range(512):
        intran += float(intra[i])
        intern += float(inter[i])

    for i in range(512):
        intra[i] = float(intra[i]) * 100.0 / intran
        inter[i] = float(inter[i]) * 100.0 / intern
    
    result = file("%s/distribution%s.txt" % (options.verify,
        name), "w")
    for i in range(512):
        result.write("%f %f %f\n" % (i * 100.0 / 512.0, intra[i], inter[i]))
    result.close()

    FNMR, FMR, EER = rates(intra, inter)
    result = file("%s/rates%s.txt" % (options.verify,
        name), "w")
    for i in range(512):
        result.write("%f %f %f\n" % (i * 100.0 / 512.0, FNMR[i], FMR[i]))
    result.close()
     
    return EER

def match(codes):
    command = [pcode]
    if options.algorithm == "lg":
        command += ["-a", "lg"]
    else:
        command += ["-a", "qsw"]
    if not options.no_lids:
        command += ["-l"]
    command += ["-r", "16"]

    for c in codes:
        command += ["-c", c]
    output = exe(command)

    result = file("%s/scores.txt" % options.verify, "w")
    for line in output.splitlines():
        if ":" in line:
            colon = line.find(":")
            slash = line.find("/")
            x = int(line[:colon])
            a = int(line[colon + 1:slash])
            b = int(line[slash + 1:])
            result.write("%d %f\n" % (x, 100.0 * a / b))
    result.close()

def gnuplot_distribution():
    name = ""
    if options.name: name = options.name
    info = get_info()
    txtout = options.verify + "/distribution%s.txt" % name
    epsout = options.verify + "/distribution%s.eps" % name
    gpout = options.verify + "/distribution%s.gnuplot" % name
    gp = file(gpout, "w")
    gp.write("""# autogenerated file
set terminal postscript eps color solid
set output '%(epsout)s'
set xrange [0:100]
set xlabel "differing bits (%%)"
set ylabel "relative match count (%%)"
set xtics 10
plot "%(txtout)s" using 1:3 title "inter class (%(info)s)" with steps, \\
    "%(txtout)s" using 1:2 title "intra class (%(info)s)" with steps
""" % locals())
    gp.close()

    exe(["gnuplot", gpout])
    print exe(["epstopdf", epsout])
    os.remove(epsout)

def gnuplot_rates(eer):
    name = ""
    if options.name: name = options.name
    info = get_info()
    txtout = options.verify + "/rates%s.txt" % name
    epsout = options.verify + "/rates%s.eps" % name
    gpout = options.verify + "/rates%s.gnuplot" % name
    gp = file(gpout, "w")
    gp.write("""# autogenerated file
set terminal postscript eps color solid
set output '%(epsout)s'
set xrange [0:100]
set xlabel "FMR (%%)"
set ylabel "FNMR (%%)"
set xtics 10
plot "%(txtout)s" using 3:2 title "ROC (%(info)s)" with lines, \\
    "-" title "EER (%(eer)f)" with points
%(eer)f %(eer)f
e
""" % locals())
    gp.close()

    exe(["gnuplot", gpout])
    print exe(["epstopdf", epsout])
    os.remove(epsout)

def gnuplot_scores():
    name = ""
    if options.name: name = options.name
    txtout = options.verify + "/scores%s.txt" % name
    epsout = options.verify + "/scores%s.eps" % name
    gpout = options.verify + "/scores%s.gnuplot" % name
    gp = file(gpout, "w")
    gp.write("""# autogenerated file
set terminal postscript eps color solid
set output '%(epsout)s'
set xlabel "sample"
set ylabel "match score"
set yrange [0:100]
plot "%(txtout)s" title "score" with steps
""" % locals())
    gp.close()

    exe(["gnuplot", gpout])
    print exe(["epstopdf", epsout])
    os.remove(epsout)

p = optparse.OptionParser(description = """
Python script which calls pcode.
""".strip())
p.add_option("-a", "--algorithm")
p.add_option("-e", "--extract",
    help = "Extract textures from CASIA and write to the output dir. " +
    "For example: -e /data/CASIA-v3i/ -o /data/textures/")
p.add_option("-E", "--extract-only",
    metavar = "id,sub,min",
    help = "For example: -E 50,L,5 takes all ids up to 50, left eyes, at " +
    "least 5 pictures")
p.add_option("-c", "--codes",
    help = "Generate feature codes from textures. Example: " +
    "-c /data/textures/ -o /data/codes/")
p.add_option("-H", "--half", action = "store_true")
p.add_option("-i", "--identify")
p.add_option("-K", "--keep-files")
# Note: As we do not generate the texture nor the mask, this only affects the
# final matching (which will completely ignore the mask). There will still
# be transformed mask pictures floating around though instead of an empty mask,
# but that sholdn't matter.
p.add_option("-L", "--no-lids", action = "store_true",
    help = "do not use masks")
p.add_option("-T", "--training", metavar = "T",
    help = "The first T samples in each class are used for training. " +
    "If there is not enough samples left for testing, the class will not " +
    "take part in testing, so make sure there's enough!"
    )
p.add_option("-n", "--name", help = "extra name to add to output files")
p.add_option("-o", "--output")
p.add_option("-R", "--rotation-offsets", help = "default is 16")
p.add_option("-v", "--verify")

g = optparse.OptionGroup(p, "Custom options",
    "Nothing here yet.")

p.add_option_group(g)
    
options, args = p.parse_args()

if options.verify:
    classes = codes_to_classes(options.verify)
    generate_exported_txt(classes, options.verify)
    intra, inter = verify()
    EER = results(intra, inter)
    gnuplot_distribution()
    gnuplot_rates(EER)
    sys.exit(0)

if options.identify:
    classes = codes_to_classes(options.identify)
    generate_exported_txt(classes, options.identify)
    identify()
    sys.exit(0)

if not options.output:
    print "-o or -v is mandatory."
    p.print_help()
    sys.exit(-1)

try:
    os.mkdir(options.output)
except OSError:
    sys.stderr.write("Please remove the existing output folder or use a " +
        "new one - this program will *never* overwrite data.\n")
    sys.exit(-1)   

if options.extract:
    extract_textures(options.extract)

elif options.codes:
    if not options.algorithm:
        print "-a option is mandatory with -c"
        print "-h for help"
        sys.exit(-1)
    classes = textures_to_classes(options.codes)
    generate_codes(classes, options.name)
	        
