#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Simple module that implements the playfair cipher"""

import string, itertools, re
from operator import attrgetter
# create an alphabet of ASCII letters without 'J'
alphabet = ''.join(letter for letter in string.ascii_uppercase if letter != 'J')

def chunkize(iterable, size):
    """Splits an iterable into size-defined chunks"""
    result = list()

    current_chunk = list()
    for index, item in enumerate(iterable):
        current_chunk.append(item)
        if index % size == size - 1:
            result.append(current_chunk)
            current_chunk = list()

    if current_chunk:
        result.append(current_chunk)

    return result

class MatrixLetter(str):
    def __new__(cls, text):
        return str.__new__(cls, text)

    def __init__(self, letter, *args, **kwargs):
        str.__init__(self, letter, *args, **kwargs)
        self.link_previous = None
        self.link_next = None

    def __lshift__(self, other):
        """Appends at the end"""
        self.link_next = other
        other.link_previous = self

    @property
    def previous(self):
        if self.link_previous is not None:
            letters_before = self.link_previous.previous
            letters_before.append(self.link_previous)
            return letters_before
        else:
            return []

    @property
    def next(self):
        if self.link_next is not None:
            letters_after = [self.link_next]
            letters_after.extend(self.link_next.next)
            return letters_after
        else:
            return []

    @property
    def position(self):
        return len(self.previous)

    @property
    def coords(self):
        return divmod(self.position, 5)

    def in_row(self, other):
        own_coords = self.coords
        other_coords = other.coords
        return True if own_coords[0] == other_coords[0] else False

    def in_column(self, other):
        own_coords = self.coords
        other_coords = other.coords
        return True if own_coords[1] == other_coords[1] else False

    def at_position(self, position):
        all = self.previous + [self] + self.next
        return all[position]

    def find(self, letter):
        all = self.previous + [self] + self.next
        return all[all.index(letter)]

    @property
    def right_letter(self):
        row, column = self.coords
        position = row * 5 + ((column + 1) % 5)
        return self.at_position(position)

    @property
    def left_letter(self):
        row, column = self.coords
        position = row * 5 + ((column - 1) % 5)
        return self.at_position(position)

    @property
    def lower_letter(self):
        row, column = self.coords
        position = ((row + 1) % 5) * 5 + column
        return self.at_position(position)

    @property
    def upper_letter(self):
        row, column = self.coords
        position = ((row - 1) % 5) * 5 + column
        return self.at_position(position)

    def flip(self, other):
        own_row, own_column = self.coords
        other_row, other_column = other.coords
        position = own_row * 5 + other_column
        return self.at_position(position)

def create_pairs(text):
    clean_text = text.upper()
    clean_text = clean_text.replace('J', 'I')

    # filter out all wrong characters
    clean_chars = [letter for letter in clean_text if letter in alphabet]

    # replace OO by OXO etc.
    previous_letter = None
    clean_letters = list()
    for letter in clean_chars:
        if previous_letter != letter:
            clean_letters.append(letter)
        else:
            clean_letters.extend(('X', letter))
        previous_letter = letter

    if len(clean_letters) % 2 == 1:
        # if odd, add X
        clean_letters.append('X')

    return chunkize(clean_letters, 2)

def create_matrix(key, width=5):
    """Creates a matrix in letter form by adding letters of the key to it and
    filling it with the rest of the alphabet"""
    matrix = list()
    # first populate by the key
    for letter in key:
        # uppercase it
        letter = letter.upper()
        # add it to the matrix it it is not already there and if it exists
        # in the alphabet
        if not letter in matrix and letter in alphabet:
            matrix.append(letter)

    # not fill it with the rest of the alphabet
    for letter in alphabet:
        if letter not in matrix:
            matrix.append(letter)

    # the matrix in list form
    list_matrix = chunkize(''.join(matrix), width)
    # the matrix as string
    joined_matrix = '\n'.join(''.join(letters) for letters in list_matrix)
    return joined_matrix

def create_letters(matrix):
    clean_matrix = matrix.replace('\n', '')
    first = MatrixLetter(clean_matrix[0])
    latest = first

    for char in clean_matrix[1:]:
        new = MatrixLetter(char)
        latest << new
        latest = new
    return first

def uni_crypt(pairs, matrix, crypt=True):
    # the attribute accessors, chosen depending on the setting of crypt
    row_letter = attrgetter('right_letter' if crypt else 'left_letter')
    column_letter = attrgetter('lower_letter' if crypt else 'upper_letter')

    for pair in pairs:
        # convert the letters to matrix letters
        matrix_letters = [matrix.find(letter) for letter in pair]
        # some checks
        if matrix_letters[0].in_row(matrix_letters[1]):
            yield [row_letter(letter) for letter in matrix_letters]
        elif matrix_letters[0].in_column(matrix_letters[1]):
            yield [column_letter(letter) for letter in matrix_letters]
        else:
            yield [matrix_letters[0].flip(matrix_letters[1]),
                    matrix_letters[1].flip(matrix_letters[0])]

def unpair(pairs):
    return ''.join(''.join(pair) for pair in pairs)

def prettify(decoded):
    # replace charXchar by charchar
    for group in re.findall(r'(\w)X(\w)', decoded):
        if group[0] == group[1]:
            decoded = decoded.replace('%sX%s' % tuple(group), ''.join(group))

    # get rid of trailing X if the sting has odd lenght and a trailing X
    if len(decoded) % 2 and decoded.endswith('X'):
        decoded = decoded[:-1]
    return decoded

def main():
    """Demonstration by using a simple matrix (key) and a simple plain text"""
    letter_matrix = create_matrix(u'Der Knabe im Moor')
    # TODO: fix umlaut-escaping
    pairs = create_pairs(u"Oh schaurig ist's, über's Moor zu gehn")
    matrix = create_letters(letter_matrix)
    enc = uni_crypt(pairs, matrix)
    #print unpair(enc)
    dec = uni_crypt(enc, matrix, crypt=False)
    decoded = unpair(dec)
    print decoded
    print prettify(decoded)


if __name__ == '__main__':
    main()