""" hebern_configs.py

    This file contains the configuration details for
    some of the historical Hebern machines. 
"""

import operator
from string import ascii_uppercase as UPPER
from copy import deepcopy
from permutations import Permutation,cycleForm,permForm
from ciphertools import phi,phiAvg
from datatools import transpose

identity = Permutation(UPPER)

def char2int(key):
    try:
        c = key.upper()
        return UPPER.index(c)
    except:
        return key

class Rotor(Permutation):

    def __init__(self,perm,key='A'):

        Permutation.__init__(self,perm,length=26,onLeft=False)
        self.setKey(key)

    def setKey(self,A):
        k = char2int(A)
        cycles = [[(x-k) % 26 for x in cycle] for cycle in self.C]
        perm = permForm(cycles)
        self.P = perm
        self.C = cycleForm(perm)
        self.key = UPPER[k]

    def step(self,j=1):
        cycles = [[(x-j) % 26 for x in cycle] for cycle in self.C]
        perm = permForm(cycles)
        self.P = perm
        self.C = cycleForm(perm)

        k = char2int(self.key)
        k = (k+j) % 26
        self.key = UPPER[k]

Hebern = {
    'keyboard': Permutation("XAKHSZJLYWGPMIOURDBFTNVCQE",onLeft=False),
    'wheels': {1 : Rotor("GADBOCTKNUZXIWHFQYJVPMELSR"),   # I
               2 : Rotor("IZNCTKUDPJEVOWLFHXSMGQAYBR"),   # II
               3 : Rotor("PJXFWLTAUGYBMHROVNCKSEQIZD"),   # III
               4 : Rotor("FLVARGWCMQBXNYIOTJUPSKEDHZ"),   # IV
               5 : Rotor("FQTGXANWCJOIVZPHYBDRKUSLEM")},  # V
    'lampboard': Permutation("TYOEUMXDFJQVKWBNSHCILRZAGP",onLeft=False),
    }

def checkPhi(R,text):
    N = len(text)/26
    tL = [list(text[26*k:26*(k+1)]) for k in range(N)]
    tL = transpose(tL)
    dL = []
    for t in tL:
        dL.append(map(R.decipher,t))
        R.step()
    sL = ["".join(l) for l in transpose(dL)]
    ps = map(phi,sL)
    return sum(ps)/len(ps)

def checkallPhi(R,text):
    l = []
    for k in range(26):
        l.append(checkPhi(R,text[k:]))
    return l

def fastRotor(text):
    text = "".join(map(Hebern['lampboard'].decipher,text))
    ll = []
    for k in range(1,6):
        rotor = deepcopy(Hebern['wheels'][k])
        ll.append(checkallPhi(rotor,text))
    return ll

def findOffset(text):
    l = []
    for j in range(26):
        currenttext = text[j:]
        N = len(currenttext)/26
        sL = [currenttext[26*k:26*(k+1)] for k in range(N)]
        pL = map(phi,sL)
        avg = sum(pL)/N
        l.append(avg)
    return l

def decode(rotor,text,key):
    rotor = deepcopy(rotor)
    rotor.setKey(key)
    output = []
    for c in text:
        rotor.step()
        output.append(rotor.decipher(c))
    return "".join(output)

def encode(rotor,text,key):
    rotor = deepcopy(rotor)
    rotor.setKey(key)
    output = []
    for c in text:
        rotor.step()
        output.append(rotor.encipher(c))
    return "".join(output)

def phiL(text):
    N = len(text)/26
    tL = [text[26*k:26*(k+1)] for k in range(N)]
    print display(map(phi,tL)),
    return phiAvg(tL)

def display(L):
    sL = [ "%0.4f" % (x,) for x in L ]
    return "  ".join(sL)

def byRotor(ptext,ctext):
    ptext = ptext.upper()
    ctext = ctext.upper()

    irotate = lambda p,c,k: ( (UPPER.index(p) + k)%26, (UPPER.index(c) + k)%26 )

    trans = set()
    for (k,p) in enumerate(ptext):
        c = ctext[k]
        trans.add(irotate(p,c,k))

    P = set([p for (p,c) in trans])
    C = set([c for (p,c) in trans])

    return len(trans)==len(P)==len(C)

def findRotor(ptext,ctext):
    ptext = ptext.upper()
    ctext = ctext.upper()

    irotate = lambda p,c,k: ((UPPER.index(p) + k)%26, (UPPER.index(c) + k)%26)

    trans = set()
    for (k,p) in enumerate(ptext):
        c = ctext[k]
        trans.add(irotate(p,c,k))

    P = set([p for (p,c) in trans])

    C = ''
    for k in range(26):
        if k in P:
            c = [c for (p,c) in trans if p==k][0]
            c = UPPER[c]
        else:
            c = ' '

        C = C+c

    return C

def spread(text,d=1):
    s = ' '*d
    return s.join(list(text))

#sample = 'GEEDNTKVLMTBTFYCYEDVWYZKQIEQFOgVVUSOJOXPBUHUFRLDLKSDLHYBYUVBTTMZXTIQQOfVTVGPKITTEDXZFHNLMGIBFOVDBHVOPCIRMVRTYIAXOTWOVVVBJcNNTEEBCGJGJUZDJOWYXXWCQCKTWLCJSAYNLGINWHKCGQSEAFCJAGJNgWFQWRITYYmVQZKFICGEXWKQVQAHUHKZBMOAXIYTMKAPCJHWHWGIPMQKVGRGLNPI'.upper()
#sampleLi5i = 'IOFSKCGFQKLASQRBQATWWEQDLNCRTEaRALAJEQARKZAMVAZDADVLRLADDVAWBRVFAOCWXXuWXPVCXNNTFSRPNBPTVPVHLRERSCNCHIISHDVZDDWUNRNQLVXNVnYQXNPVYYAKQNJZZTVMMVZUUBIPPIKUIUbURKBJCUGECFDIOQWTLOCQeOTvUSOTAMrRQVJXSRqJDJIDKJMJRWXETJNKRDKJTUGKDMTJAYELUSODCYDKXJYD'.upper()

sampleA = """

GEEDN TKVLM TBTFY CYEDV WYZKQ IEQFO gVVUS OJOXP BUHUF RLDLK
SDLHY BYUVB TTMZX TIQQO fVTVG PKITT EDXZF HNLMG IBFOV DBHVO
PCIRM VRTYI AXOTW OVVVB JcNNT EEBCG JGJUZ DJOWY XXWCQ CKTWL
CJSAY NLGIN WHKCG QSEAF CJAGJ NgWFQ WRITY YmVQZ KFICG EXWKQ
VQAHU HKZBM OAXIY TMKAP CJHWH WGIPM QKVGR GLNPI

""".strip().upper().replace(' ','').replace('\n','')

sampleA_Li = "".join(map(Hebern['lampboard'].decipher,sampleA))
sampleA_Li5i = decode(Hebern['wheels'][5],sampleA_Li,'L')

sampleB_Li5i = """
                          I O F S K C G F Q K L A S
Q R B Q A T W W E Q D L N C R T E a R A L A J E Q A
R K Z A M V A Z D A D V L R L A D D V A W B R V g A
O C W X X u W X P V C X N N T F S R P N B P T V P V
H L R E R S C N C H I I S H D V Z D D W U N R N Q L
V X N V n Y Q X N P V Y Y A K Q N J Z Z T V M M V Z
U U B I P P I K U I U b U R K B J C U G E C F D I O
Q W T L O C Q e O T v U S O T A M r R Q V J X S R q
J D J I D K J M J R W X E T J N K R D K J T U G K D
M T J A Y E L U S O D C Y D K X J Y D

""".strip().upper().replace(' ','').replace('\n','')

sampleB_Li = encode(Hebern['wheels'][5],sampleB_Li5i,'L')
sampleB = "".join(map(Hebern['lampboard'].encipher,sampleB_Li))

sampleC_Li5i = """
                          I O F S K C G F Q K L A S
Q R B Q A T W W E Q D L N C R T E A R A L A J E Q A
R K Z A M V A Z D A D V L R u A D D V A W B R V G A
O C W X X U x X P V C X N N T F S R P N B P T V P V
H L R E R S C N C H I I S H D V Z D D W U N R N Q L
V X N V j Y Q X N P V Y Y A K Q N J Z Z T V M M V Z
U U B I P P I K U I U B U R K B J C U G E C F D I O
Q W T L O C Q E O T V U S O T A M w R Q V J X S R Q
J D J I D K J M J R W X E T J N K R D K J T U G K D
M T J A Y E L U S O D C Y D K X J Y D

""".strip().upper().replace(' ','').replace('\n','')

sampleC_Li = encode(Hebern['wheels'][5],sampleC_Li5i,'L')
sampleC = "".join(map(Hebern['lampboard'].encipher,sampleC_Li))

sampleD = """

GEEDN TKVLM TBTFY CYEDV WYZKQ IEQFO fVVUS OJOXP BUHUF RLDLK
SDLhY BYUVB TTMjX TIQQO tvTVG PKITT EDXZF HNLMG IBFOV DBHVO
PCIRM VRTYI AXOTW OVVVB JcNNT EEBCG JGJUZ DJOWY XXWCQ CKTWL
CJSAg NLGIN WHKCG QSEAF CJAGJ NfWFb WRITY YmVQZ KFICh EXWKQ
VQAHU HKZBM OAXIY TMKAP CJHWH WGIPM QKVGR GLNPI

""".strip().upper().replace(' ','').replace('\n','')

sampleD_Li = "".join(map(Hebern['lampboard'].decipher,sampleD))
sampleD_Li5i = decode(Hebern['wheels'][5],sampleD_Li,'L')

message = """

Confidential for Jones STOP
I am returning tonight on the ten fifteen train STOP
Meet me at Pennsylvania Station and we will discuss the negotiations at dinner STOP
Phillips still insists on the same price for the stock STOP
If we can sweeten the deal somewhat he might budge STOP
Have an idea

""".strip().upper().replace(' ','').replace('\n','')


"""
compare(sampleA,sampleB)    compare(sampleB,sampleC)    compare(sampleA,sampleC)
 30 G F                      30                          30 G F
 53                          53 H J                      53 H J
 63 Z J                      63                          63 Z J
 70 F T                      70                          70 F T
 71                          71 V B                      71 V B
121 C R                     121 R C                     121
154 Y G                     154                         154 Y G
176 G F                     176                         176 G F
179 Q B                     179                         179 Q B
186 M X                     186 X M                     186
194 G H                     194                         194 G H

compare  A B C D
----------------
     30  G F F F
     53  H H J H
     63  Z J J J
     70  F T T T
     71  V V B V
    121  C R C C
    154  Y G G G
    176  G F F F
    179  Q B B B
    186  M X M M
    194  G H H H

sampleA
GEEDNTKVLMTBTFYCYEDVWYZKQIEQFOgVVUSOJOXPBUHUFRLDLKSDLhYBYUVBTTMzXTIQQOfvTVGPKITTEDXZFHNLMGIBFOVDBHVOPCIRMVRTYIAXOTWOVVVBJcNNTEEBCGJGJUZDJOWYXXWCQCKTWLCJSAyNLGINWHKCGQSEAFCJAGJNgWFqWRITYYmVQZKFICgEXWKQVQAHUHKZBMOAXIYTMKAPCJHWHWGIPMQKVGRGLNPI
sampleB                       |                                |      |                                                  |                                |                     |  |      |       |
GEEDNTKVLMTBTFYCYEDVWYZKQIEQFOfVVUSOJOXPBUHUFRLDLKSDLhYBYUVBTTMjXTIQQOtvTVGPKITTEDXZFHNLMGIBFOVDBHVOPCIRMVRTYIAXOTWOVVVBJrNNTEEBCGJGJUZDJOWYXXWCQCKTWLCJSAgNLGINWHKCGQSEAFCJAGJNfWFbWRITYYxVQZKFIChEXWKQVQAHUHKZBMOAXIYTMKAPCJHWHWGIPMQKVGRGLNPI
sampleC                                              |                 |                                                 |                                                                |
GEEDNTKVLMTBTFYCYEDVWYZKQIEQFOfVVUSOJOXPBUHUFRLDLKSDLjYBYUVBTTMjXTIQQOtbTVGPKITTEDXZFHNLMGIBFOVDBHVOPCIRMVRTYIAXOTWOVVVBJcNNTEEBCGJGJUZDJOWYXXWCQCKTWLCJSAgNLGINWHKCGQSEAFCJAGJNfWFbWRITYYmVQZKFIChEXWKQVQAHUHKZBMOAXIYTMKAPCJHWHWGIPMQKVGRGLNPI

sampleA_Li5i
IOFSKCGFQKLASQRBQATWWEQDLNCRTEuRALAJEQARKZAMVAZDADVLRlADDVAWBRVfAOCWXXvwXPVCXNNTFSRPNBPTVPVHLRERSCNCHIISHDVZDDWUNRNQLVXNVjYQXNPVYYAKQNJZZTVMMVZUUBIPPIKUIUnURKBJCUGECFDIOQWTLOCQsOTtUSOTAMwRQVJXSRjJDJIDKJMJRWXETJNKRDKJTUGKDMTJAYELUSODCYDKXJYD
sampleB_Li5i                  |                                |      |                                                  |                                |                     |  |      |       |
IOFSKCGFQKLASQRBQATWWEQDLNCRTEaRALAJEQARKZAMVAZDADVLRlADDVAWBRVgAOCWXXuwXPVCXNNTFSRPNBPTVPVHLRERSCNCHIISHDVZDDWUNRNQLVXNVnYQXNPVYYAKQNJZZTVMMVZUUBIPPIKUIUbURKBJCUGECFDIOQWTLOCQeOTvUSOTAMrRQVJXSRqJDJIDKJMJRWXETJNKRDKJTUGKDMTJAYELUSODCYDKXJYD
sampleC_Li5i                                         |                 |                                                 |                                                                |
IOFSKCGFQKLASQRBQATWWEQDLNCRTEaRALAJEQARKZAMVAZDADVLRuADDVAWBRVgAOCWXXuwXPVCXNNTFSRPNBPTVPVHLRERSCNCHIISHDVZDDWUNRNQLVXNVjYQXNPVYYAKQNJZZTVMMVZUUBIPPIKUIUbURKBJCUGECFDIOQWTLOCQeOTvUSOTAMwRQVJXSRqJDJIDKJMJRWXETJNKRDKJTUGKDMTJAYELUSODCYDKXJYD
"""