import numpy as np
from matplotlib import pyplot as mt


def Rz(y):
    userInput = y.copy()
    userInput.pop(-1)
    emptyList = []
    for i in range(len(userInput)):
        if userInput[i] == -volt:
            emptyList.append(-volt)
            emptyList.append(0)
        else:
            emptyList.append(volt)
            emptyList.append(0)
    T = np.arange(0, len(y), 0.5)
    T = list(T)
    T.pop((-1))
    rem = len(T) - len(emptyList)
    emptyList.extend([0] * rem)
    return T, emptyList


def polar(y):
    userInput = y.copy()
    for i in range(len(userInput)):
        if 0 == userInput[i]:
            userInput[i] = -volt
    return userInput


def plot(n, d, x, b):
    mt.subplot(1, 1, n)
    mt.step(d, x, where='post')
    mt.grid(True, which='both')
    mt.xlim(0, len(y)+1)
    mt.ylim(-b-0.1, volt+1)


y = input('Enter Binary Code: ')
volt = int(input('Enter the Voltage Level: '))
print('Binary of ', 'is', y)
y = y + '0'
y = int(y)
y1 = volt * y
y = list(map(int, str(y1)))
y1 = y.copy()
t = list(range(0, len(y1)))
t1 = t.copy()

fig, ax = mt.subplots(1, 1)
fig.tight_layout()
fig = mt.gcf()
ax.remove()

y22 = polar(y1)
T, y2 = Rz(y22)
plot(1, T, y2, volt)
mt.title('RZ Polar')

mt.show()