# coding: cp1251
import random
import timeit
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

def do_lists(a):
	s = 0
	for i in range(n):
		s += a[i][i] # суммируем диагональ
	return s

def do_numpy(a):
	s = sum(a.diagonal())
	# for i in range(n):
	# 	s += a[i][i]
	return s

def do_numpy_method(a):
	s = a.trace() # получаем след матрицы
	return s

# list'ы
# small
df = pd.DataFrame(columns=["lists", "numpy", "numpy_method", "size"]) # задаём датафрейм
k = 1
n = 10
a = []
for j in range(n):
	a.append(random.sample(range(n), n)) # заполняем матрицу случайными данными
b = np.random.randint(0, n, (n, n)) # заполняем матрицу случайными данными
c = np.random.randint(0, n, (n, n)) # заполняем матрицу случайными данными
for i in range(50):
	row = []
	row.append(timeit.timeit(lambda: do_lists(a), number=1)) # добавляем результат листов
	row.append(timeit.timeit(lambda: do_numpy(b), number=1)) # добавляем результат numpy
	row.append(timeit.timeit(lambda: do_numpy_method(c), number=1)) # добавляем результат метода numpy
	row.append(i)
	df.loc[i+1] = row

# print(df)
fig1 = df[["lists", "numpy", "numpy_method"]].plot()
fig1.set_xlabel("Number", fontsize=12)
fig1.set_ylabel("Time", fontsize=12)
plt.savefig("small.png")
# plt.show()

# medium
df = pd.DataFrame(columns=["lists", "numpy", "numpy_method", "size"])
n = 50
a = []
for j in range(n):
	a.append(random.sample(range(n), n)) # заполняем матрицу случайными данными
b = np.random.randint(0, n, (n, n)) # заполняем матрицу случайными данными
c = np.random.randint(0, n, (n, n)) # заполняем матрицу случайными данными
for i in range(50):
	row = []
	row.append(timeit.timeit(lambda: do_lists(a), number=1)) # добавляем результат листов
	row.append(timeit.timeit(lambda: do_numpy(b), number=1)) # добавляем результат numpy
	row.append(timeit.timeit(lambda: do_numpy_method(c), number=1)) # добавляем результат метода numpy
	row.append(i)
	df.loc[i+1] = row
# print(df)
fig2 = df[["lists", "numpy", "numpy_method"]].plot()
fig2.set_xlabel("Number", fontsize=12)
fig2.set_ylabel("Time", fontsize=12)
plt.savefig("medium.png")
# plt.show()

# big
df = pd.DataFrame(columns=["lists", "numpy", "numpy_method", "size"])
n = 100
a = []
for j in range(n):
	a.append(random.sample(range(n), n)) # заполняем матрицу случайными данными
b = np.random.randint(0, n, (n, n)) # заполняем матрицу случайными данными
c = np.random.randint(0, n, (n, n)) # заполняем матрицу случайными данными
for i in range(50):
	row = []
	row.append(timeit.timeit(lambda: do_lists(a), number=1)) # добавляем результат листов
	row.append(timeit.timeit(lambda: do_numpy(b), number=1)) # добавляем результат numpy
	row.append(timeit.timeit(lambda: do_numpy_method(c), number=1)) # добавляем результат метода numpy
	row.append(i)
	df.loc[i+1] = row
# print(df)
fig3 = df[["lists", "numpy", "numpy_method"]].plot()
fig3.set_xlabel("Number", fontsize=12)
fig3.set_ylabel("Time", fontsize=12)
plt.savefig("big.png")
plt.show() # выводим графики