f = open('frost_cipher.txt')
cipher = f.readlines()
cipher = [c.strip('\n').split(' ') for c in cipher]

clean = ''
for c in cipher:
    clean += "".join(c)
#    print("".join(c))
ciphertext = clean.upper()


colors = ['red', 'blue', 'darkgreen', 'orange', 'black', 'yellow']

width = 110
blocknum = 4
blocks = ["" for i in range(0,blocknum)]
output0 = "\\texttt{"
output1 = "\\texttt{"

for i in range(len(ciphertext)):
    output0 += ciphertext[i]

    for j in range(blocknum):
        if i % blocknum == j:
            output1 += "\\textcolor{" + colors[j] +  "}{" + ciphertext[i] + "}"

    blocks[i%blocknum] += ciphertext[i]

    if (i+1)%width == 0:
        output0 += '\n'
        output1 += '\n'

output0 += "}"
output1 += "}"

f0 = open('friedman_0.tex', 'w')
f0.write(output0)
f0.close()

f1 = open('friedman_{}.tex'.format(blocknum), 'w')
f1.write(output1)
f1.close()

def compute_coincidence(text):
    letter_count = {}
    letters = [chr(i) for i in range(ord('A'), ord('Z')+1)]

    for c in letters:
        letter_count[c] = 0

    for c in text:
        if c in letter_count.keys():
            letter_count[c] += 1

    N = len(text)
    s = 0
    for n_i in letter_count.values():
        s += n_i * (n_i - 1) 
    kappa =  float(s) / (N * (N - 1))
    print("kappa_o = ", kappa)
    return kappa


def output_block(block, color, width=20):
    output = "\\texttt{\\textcolor{" + color + "}{"
    for i in range(len(block)):
        output += block[i]
        if (i+1)%width == 0:
            output += '\n'

    output += "}}"
    
    kappa = compute_coincidence(blocks[j])
    output += "\\textcolor{"+color+"}{$$" + "\\kappa_o = {:6.4f}".format(kappa) + "$$}"
    print(output)

    f = open('friedman_{}_'.format(blocknum) + color + '.tex', 'w')
    f.write(output)
    f.close()

for j in range(blocknum):
    output_block(blocks[j], colors[j])