from re import X
from matplotlib import pyplot as plt
import numpy as np
import scipy as sp
import csv

table1 = '1.csv'
table2 = '2.csv'
table3 = '3.csv'


def open_CSV(filename):
    with open('/home/dmitrii/Physics Labs/3.6.1/' + filename, newline="") as csvfile:
        datareader = csv.reader(csvfile, delimiter=",", quotechar="|")
        data = []
        for row in datareader:
            res = [float(i) for i in row]
            data.append(res)
        data = np.array(data)
        return data


def plot(arr, xlab, ylab, plotName):
    x = arr[:, 0]
    y = arr[:, 1]
    params = np.polyfit(x, y, 1)
    print("Next line:")
    print("Approx line params:")
    print(params)
    plt.plot(x, x*params[0]+params[1], color='orange', linewidth=1)
    plt.scatter(x, y, marker='+', s=200, color='green')
    plt.xlabel(xlab, fontsize=20)
    plt.ylabel(ylab, fontsize=20)
    plt.title(plotName, fontsize=20)
    plt.grid()
    plt.show()


def main(table1, table2, table3):
    plot(table1, r'$\frac{1}{\tau}, 10^2$',
         r'$\Delta \nu, Hz, 10^2$', r'$\Delta \nu(\frac{1}{\tau})$')
    plot(table2, r'$f_{period}, Hz, 10^3$',
         r'$\Delta \nu, Hz, 10^3$', r'$\Delta \nu(f_{period})$')
    plot(table3, r'$\frac{A_{side}}{A_{main}}, 1$',
         r'$m, 1$', r'$m(\frac{A_{side}}{A_{main}})$')


if __name__ == '__main__':
    main(open_CSV(table1), open_CSV(table2), open_CSV(table3))
