"""Simulateur fidèle d'Enigma I (Wehrmacht/Luftwaffe) : rotors I-V, réflecteur B,
Ringstellung, Steckerbrett, double pas du rotor du milieu."""
A = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
ROTORS = {
    "I":   ("EKMFLGDQVZNTOWYHXUSPAIBRCJ", "Q"),
    "II":  ("AJDKSIRUXBLHWTMCQGZNPYFVOE", "E"),
    "III": ("BDFHJLCPRTXVZNYEIWGAKMUSQO", "V"),
    "IV":  ("ESOVPZJAYQUIRHXLNFTGKDCMWB", "J"),
    "V":   ("VZBRGITYUPSDNHLXAWMJQOFECK", "Z"),
}
REFLECTOR_B = "YRUHQSLDPXNGOKMIEBFZCWVJAT"

class Enigma:
    def __init__(self, rotors, rings, positions, plugboard=""):
        self.fwd = [[A.index(c) for c in ROTORS[r][0]] for r in rotors]
        self.bwd = [[f.index(i) for i in range(26)] for f in self.fwd]
        self.notch = [A.index(ROTORS[r][1]) for r in rotors]
        self.ring = [A.index(c) for c in rings]
        self.pos = [A.index(c) for c in positions]
        self.plug = list(range(26))
        for pair in plugboard.split():
            a, b = A.index(pair[0]), A.index(pair[1])
            self.plug[a], self.plug[b] = b, a
        self.refl = [A.index(c) for c in REFLECTOR_B]

    def step(self):
        p, n = self.pos, self.notch
        if p[1] == n[1]:            # double pas : le milieu entraîne la gauche et avance lui-même
            p[0] = (p[0] + 1) % 26
            p[1] = (p[1] + 1) % 26
        elif p[2] == n[2]:
            p[1] = (p[1] + 1) % 26
        p[2] = (p[2] + 1) % 26

    def press(self, c):
        self.step()
        x = self.plug[A.index(c)]
        for i in (2, 1, 0):
            o = self.pos[i] - self.ring[i]
            x = (self.fwd[i][(x + o) % 26] - o) % 26
        x = self.refl[x]
        for i in (0, 1, 2):
            o = self.pos[i] - self.ring[i]
            x = (self.bwd[i][(x + o) % 26] - o) % 26
        return A[self.plug[x]]

    def crypt(self, text):
        return "".join(self.press(c) for c in text if c in A)

def crypt(text, rotors, rings, positions, plugboard=""):
    return Enigma(rotors, rings, positions, plugboard).crypt(text)
