#Rail fence cipher (multiple rails)

def rail_fence_cipher(message, numbRails, action):
    
    #Initialises the list rails, creating a row for each rail and populating each row with an empty string
    rails = []
    for row in range(numbRails):
         rails.append('')

    #Inserts the characters in the input message, including spaces, into the list rails in a zig zag pattern
    row = 0
    downUp = 1
    for char in message:
            rails[row] = rails[row] + char
            row = row + downUp
            #Changes direction when top or bottom rail is reached
            if row == numbRails:
                downUp = -1
                row = row - 2
            if row == -1:
                downUp = 1
                row = row + 2
    #Handles encryption
    if action == '1':           
        #Joins the rails together to create the text to be output
        outputText = ''
        for row in range(numbRails):
            outputText = outputText + rails[row]

    #Handles decryption
    else:     
        #Initialises the list countCharsOnRail giving each element of the list an initial value of 0
        countCharsOnRail = []
        for row in range(numbRails):
            countCharsOnRail.append(0)
        for row in range(numbRails):
            countCharsOnRail[row] = len(rails[row])

        #Initialises nested list newRail which is used to re-assemble encrypted message
        newRail = []
        for row in range(numbRails):
           newRail.append([])
        #Each element of list newRail holds the characters from a rail
        begin = 0
        for row in range(0, numbRails):
            end = begin + countCharsOnRail[row]
            #The built-in type conversion function list used to convert slices of original message into lists
            newRail[row] = list(message[begin:end])
            begin = end

        #Assembles text to be output by popping characters sequentially from newRail
        row = 0
        outputText = ''
        upDown = 1
        for index in range(len(message)):
            outputText = outputText + newRail[row].pop(0)
            row = row + upDown
            if row == numbRails:
                upDown = -1
                row = row - 2
            if row == -1:
                upDown = 1
                row = row + 2

    return outputText
###end function###
  
###Main program###
print('RAIL FENCE')
print('\nOPTIONS')
flag = True
while flag:
    print('')
    print('1. Encrypt a message')
    print('2. Decrypt a message')
    print('3. Exit the program')
    action = input('\nPlease select an option. ')

    if action =='3':
        flag = False
        print('Good bye')
    else:
        message = input('\nEnter the message: ')
        numbRails = int(input('Enter the number of rails: '))
        if numbRails < 2:
            print('\nThe cipher only works with two or more rails.')
        else:
            print('\nThe messgage is: '),
            print(rail_fence_cipher(message, numbRails, action))
   

