from operator import itemgetter 
from scipy.stats import rv_discrete
import numpy as np
import sys, math
import random
from joblib import Parallel, delayed

#verbose=True
verbose=False

delta = 0.01

ni = {}
table = [] # rows: 0 - empty, 1...s; columns 0,1...s
table_ids = ['F00']
ignore = []
tests = []
faults = []
non_triv_faults = []

def num_samples(mu1, mu2):
    if mu1[2] != 0 or mu2[2] != 0:
        raise ZeroDivisionError
    p = mu1[0]
    return int(math.log(delta)/(math.log(2)+0.5*math.log(p*(1-p))))

def parse_file(file):
    fault_cnt = 0
    ignore_index = []
    with open(file) as f:
        for line in f:
            line = line.strip()
            line_parts = line.split('=');
            prefix = line_parts[0].split(':')
            if float(prefix[1]) - 0.5 == 0:
                ignore.append(prefix[0])
                ignore_index.append(fault_cnt)
                continue
            tuples = line_parts[1].split('|')
            fault_triples = [eval(t) for t in tuples]
            table.append(fault_triples)
            table_ids.append(prefix[0])
            if verbose: print "=== Fault-%s mapped to Test-%d ===" % (prefix[0], fault_cnt)
            fault_cnt += 1

    if len(ignore_index) > 0:
        if verbose: print "************** Modified Table ****************"
        # also remove columns Fi for i in ignore
        for t in table:
            for i in ignore_index:
                t.pop(i+1)
            if verbose: print t
    
        if verbose: print "**********************************************"
    # add empty row in beginning for easy indexing
    table.insert(0,[0]*(len(table)+1))
    return (table, ignore)

def init(f):
    global num_gates, table, ignore
    global tests, faults, non_triv_faults
    (table, ignore) = parse_file(f)
    num_gates = len(table) - 1 # adjust extra row
    tests = range(1,num_gates+1)
    faults = range(0, num_gates+1)
    non_triv_faults = range(1, num_gates+1)

def compute_ni(table):
    global ni

    ni = {}
    for r in tests:
        mu1 = table[r][0]
        mu2 = table[r][r]
        if mu1 == mu2:
            print r
            raise ZeroDivisionError
        num = 0
        if mu1 == (1,0,0) or mu1 == (0,1,0):
            num = 1
        else:
            num = int(num_samples(mu1, mu2))
        print " ** %d samples needed for Test-%d: %s" % (num, r, table[r])
        ni[r] = num

    # sort tests by increasing value of ni
    tests.sort(cmp = lambda i,j : ni[i] - ni[j])
    print "New sequence of tests:", tests
    #print max(ni.values()),reduce(lambda x,y:x+y,ni.values())

def DKL(q, p):
    d = 0.0
    
    for c in [0,1,2]:
        if p[c] == 0 and q[c] == 0:
            pass
        elif p[c] == 0 or q[c] == 0:
            d = d + 10000 # add infinity
        else:
            d = d + q[c]*math.log((q[c]/p[c]),2)

    return d

elements = np.array([0,1,2])

def sample_mc(distr, num_samples):
    probabilities = np.array(distr)
    samples = np.random.choice(elements, num_samples, p=list(probabilities))
    sample_distr = (np.count_nonzero(samples == 0), np.count_nonzero(samples == 1), np.count_nonzero(samples == 2))
    sample_distr = [sample_distr[i]*1.0/num_samples for i in [0,1,2]]
    s = []
    if num_samples < 20: s = samples
    if verbose: print "%d samples from %s is %s [%s]" % (num_samples, distr, sample_distr, s)
    return sample_distr

def compute_tau(C,i):
    distr = table[i][C]
    tau = sample_mc(distr, ni[i])
    return tau

def exp(i,C): # runs i-th test on the circuit C
    if verbose: print "  ==> Beginning exp ",i
    tau = compute_tau(C,i)
    if tau[2] != 0: return -1
    d1 = DKL(tau, table[i][0])
    d2 = DKL(tau, table[i][i])
    if verbose: print "Comparing sample distr %s against FF distr %s and Fault-%d distr %s" % (tau, table[i][0], i, table[i][i])
    if d1 <= d2:
        return 0
    else:
        return i

def detect_rand(C): # return True if Faulty and False if Fault-free
    i = random.randint(1, num_gates)
    samples = ni[i]
    f = exp(i,C)
    if f == -1 or f != 0:
        return (True, samples)
    else:
        return (False, samples)

def detect_det(C): # return True if Faulty and False if Fault-free
    samples = 0
    for i in tests:
        f = exp(i,C)
        samples = samples + ni[i]
        if f == -1:
            return (True, samples)
        # if tau(2) == 0:
        if verbose: print "Test-%d for Fault-%d returned F-%d" % (i,C,f) 
        if f != 0:
            return (True, samples)
    return (False,samples)

def detect(f, num_iter):
    missed_f = 0
    samples_f = 0
    if verbose: print "****** Computing SUSP for Fault-",table_ids[f]
    
    for iter in range(num_iter):
        if verbose: print "====> Iteration: ", iter
        (is_faulty, s) = detect_det(f)
        if (f == 0 and is_faulty) or (f > 0 and not is_faulty):
           missed_f = missed_f+1
        else:
           samples_f = samples_f + s

    return (f, missed_f, samples_f)

def main(num_iter):
    if len(sys.argv) < 2:
        print "Error. python cluster.py <diagnostic-table.txt>"
        sys.exit(1)

    init(sys.argv[1])
    #duplicates = remove_duplicates(table)
    compute_ni(table)
    print "   Testing %d faults. %d faults undetectable from fault-free: %s" % (num_gates+1, len(ignore), ignore)

    missed = {}
    samples = {}

    ret = Parallel(n_jobs=3, verbose=2)(delayed(detect)(f, num_iter) for f in faults)

    for (f, m, s) in ret:
        missed[f] = m
        samples[f] = s

    print "----------- RESULTS ------------"
    for f in faults:
        print "Fault: %d, Total = %d, missed = %d (%.2f %%), number of samples (when detected) = %f" % \
        (f, num_iter, missed[f], missed[f]*100.0/num_iter, samples[f]*1.0/(num_iter - missed[f]))

main(10000)