Laba6

import math
import random
import numpy as np
import matplotlib.pyplot as plt

class complex:
    re = 0
    im = 0

def dif(a, b):
    r = complex()
    r.re = a.re - b.re
    r.im = a.im - b.im
    return r

def sum(a, b):
    r = complex()
    r.re = a.re + b.re
    r.im = a.im + b.im
    return r

def prod(a, b):
    r = complex()
    r.re = a.re * b.re - a.im * b.im
    r.im = a.re * b.im + a.im * b.re
    return r


def prn(a, n):
    print()
    x = []
    y = []
    for i in range(1, n+1):
       x.append(i)
       y.append(a[i].re)
    plt.plot(x, y)

def fft(x, m):
    u = s = t = g = complex()
    r = n = le = le2 = ip = 0
    n = int(math.pow(2, m))

    for l in range(1, m+1):
        le = int(math.pow(2, m + 1 - l))
        le2 = int(le / 2)
        u.re = 1
        u.im = 0
        r = math.pi / le2
        s.re = math.cos(r)
        s.im = math.sin(r)
        for j in range(1, le2 + 1):
            for i in range(j, n + 1, le):
                ip = i + le2
                t = sum(x[i], x[ip])
                x[ip] = prod(dif(x[i], x[ip]), u)
                x[i] = t
            u = prod(u, s)

    nd2 = int(n / 2)
    nm1 = n - 1
    j = 1
    k = 0

    for q in range(1, int(nm1) + 1):
        if q < j:
            x[j] = x[q]
            x[q] = x[j]
        k = nd2
        while k < j:
            j -= k
            k /= 2

        j += k

def sig(k1, k2, t1):
    return 3 * (1 + 0.3 * np.cos(2*np.pi*k1*t1/n)) * np.cos(2*np.pi*k2*t1/n)


k1 = 5
k2 = 17
n = 1024
m = 10
x = []
for i in range(n+1):
    x.append(complex())

for i in range(1, n+1):
    x[i].re = sig(k1, k2, i)
    x[i].im = 0

prn(x, n)
plt.show()
fft(x, m)
prn(x, n)
plt.show()