import copy
from random import randint
import math



def generowanie_listy(N,d1,d2):
    tab = []
    for i in range(N):
        tab.append(0)
    i = 0
    while i < (d1+d2):
        a = randint(0,N-1)
        if tab[a] == 1 or tab[a] == -1:
            continue
        else:
            if i >= d1:
                tab[a] = -1
                i += 1
            else:
                tab[a] = 1
                i += 1
    return tab
def centralLift(q,poly):
    a = list(poly)
    for i in range(len(a)):
        if a[i] >= (q/2):
            a[i] = a[i] - q
        elif a[i] <= (-q/2):
            a[i] = a[i] + q

    return QRZ(a)

def generowanie_f_fp_fq(N,d1,d2):
    corrFlag = 0
    while (corrFlag == 0):
        corrFlag = 1
        f = QRZ(generowanie_listy(N,d1,d2))
        try:
            fp = QRP(list(f))^(-1)
        except ZeroDivisionError:
            corrFlag = 0

        try:
            fq = QRQ(list(f))^(-1)
        except ZeroDivisionError:
            corrFlag = 0
    return (f,fp,fq)

class NTRU:
    def __init__(self,N,p,q,d1,d2,d3):
        self.N  = N
        self.p  = p
        self.q  = q
        self.d1 = d1
        self.d2 = d2
        self.d3 = d3
        self.f,self.fp,self.fq =  generowanie_f_fp_fq(self.N,self.d1+1,self.d1)
        self.g =  QRZ(generowanie_listy(self.N,self.d2,self.d2))
        self.h = self.p * self.fq*QRQ(list(self.g))

    def szyfrowanie(self,m):
        m = QRQ(list(m))
        r = generowanie_listy(self.N,self.d3,self.d3)
        r = QRQ(r)
        e = r * self.h + m
        return e

    def deszyfrowanie(self,e):
        a = QRQ(list(self.f))*QRQ(list(e))
        a = QRZ(list(a))
        a = centralLift(self.q,a)
        d = self.fp * QRP(list(a))
        d = QRZ(list(d))
        d = centralLift(self.p,d)
        return d




N = 503
p = 3
q = 256
df = 215
dg = 72
dr = 18

Fp=GF(p)
Fq=GF(q)

FP.<u>=PolynomialRing(Fp)
FQ.<v>=PolynomialRing(Fq)
FZ.<l>=PolynomialRing(ZZ)

QRZ.<z>=FZ.quotient(l^(N)-1)
QRP.<x>=FP.quotient(u^(N)-1)
QRQ.<y>=FQ.quotient(v^(N)-1)
wiadomosc = QRZ(genList(N,33,33))

print(wiadomosc)
NTRU=NTRU(N,p,q,df,dg,dr)

e = NTRU.szyfrowanie(wiadomosc)
print(e)

d = NTRU.deszyfrowanie(e)
print(d)
if d == wiadomosc:
    print("DZIALA")