from contextlib import redirect_stderr
import numpy as np
from matplotlib import pyplot as plt
import csv
import scipy.constants as sp


def main():
    print("Code is running.")
    plot_UI(open_CSV("U(I).csv"))
    plot_UIs(open_CSV("1,5mA.csv"), '[1,5mA]')
    plot_UIs(open_CSV("3mA.csv"), '[3mA]')
    plot_UIs(open_CSV("5mA.csv"), '[5mA]')


def open_CSV(filename):
    with open('/home/dmitrii/Physics Labs/3.5.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_UI(data):
    plt.scatter(data[:, 1][0:23], data[:, 0][0:23], color="g")
    plt.scatter(data[:, 1][23:len(data)], data[:, 0]
                [23:len(data)], color="orange")
    plt.xticks(np.arange(min(data[:, 1])-0.25, max(data[:, 1])+0.25, 0.25))
    plt.yticks(np.arange(min(data[:, 0])-1, max(data[:, 0])+1, 1))
    A = np.vstack([data[:, 1][5:12], np.ones(len(data[:, 1][5:12]))]).T
    alpha = np.dot((np.dot(np.linalg.inv(np.dot(A.T, A)), A.T)),
                   data[:, 0][5:12])
    var = str(round(alpha[1]*10**3))
    delta = str(abs(round(alpha[0]*10**3)))
    print("Максимальное дифференциальное сопротивление разряда: {}".format(
        var) + u" \u00B1 " + "{}, Ом".format(delta))
    plt.plot(data[:, 1][5:12], alpha[0]*data[:, 1][5:12]+alpha[1], 'r')
    plt.ylabel('U, В', fontsize=20)
    plt.xlabel('I, мкА', fontsize=20)
    plt.grid()
    plt.savefig('U(I).png')
    plt.show()


def centre(data):
    xdiff = [data[n]-data[n-1] for n in range(1, len(data))]
    return sum(xdiff)/2


def bfl(data):
    line = np.polyfit(data[:, 0], data[:, 1], 1)
    return line


def plot_UIs(data, num):
    ax = plt.gca()
    ax.spines['top'].set_color('none')
    ax.spines['left'].set_position('zero')
    ax.spines['right'].set_color('none')
    ax.spines['bottom'].set_position('zero')
    plt.scatter(data[:, 0], data[:, 1], color="g",
                marker="+", s=100)  # Make a scatter plot
    data = np.array(sorted(data, key=lambda x: x[0], reverse=True))
    line1 = bfl(data[len(data)-5:len(data)])
    line2 = bfl(data[0:5])
    aprox = bfl(data[4:len(data)-5])
    x1 = (line1[1]-aprox[1])/(aprox[0]-line1[0])
    y1 = line1[0]*x1+line1[1]
    plt.plot(x1, y1, marker='+', color='red')
    x2 = (line2[1]-aprox[1])/(aprox[0]-line2[0])
    y2 = line2[0]*x2+line2[1]
    plt.plot(x2, y2, marker='+', color='red')
    plt.plot((data[:, 0][len(data)-5:len(data)]), (data[:, 0]
             [len(data)-5:len(data)]*line1[0]+line1[1]), color="orange")
    plt.plot([x1, 0], [y1, line1[1]], linestyle="dashed", color="orange",
             label='Ток насыщения: ' + str(round(line1[1], 2)) + ' ,А')
    plt.plot(data[:, 0][0:5], data[:, 0]
             [0:5]*line2[0]+line2[1], color="orange")
    plt.plot([x2, 0], [y2, line2[1]], linestyle="dashed", color="orange",
             label='Ток насыщения: ' + str(round(line2[1], 3)) + ' ,А')
    S = sp.pi*0.2*10**-4*5.2*10**-4
    me = 22*1.66*10**-27
    Te = x2/2*11400
    ni = line2[1]/0.4/sp.e/S/np.sqrt(2*sp.k*Te/me)/10**14
    wp = np.sqrt(4*np.pi*ni*sp.e**2/me)*10**10
    rde = np.sqrt(sp.k*Te/4/sp.pi/ni/sp.e**2)*10**-5
    rd = np.sqrt(sp.k*300/4/sp.pi/ni/sp.e**2)
    Nd = 4/3*sp.pi*rd**3*ni
    print("-----Exp{}-----".format(num))
    print("Ток насыщения: {},{}".format(
        round(line1[1], 3), round(line2[1], 3)))
    print("Te = {}".format(Te   ))
    print("Ni = {}".format(ni))
    print("wp = {}".format(wp))
    print("rde = {}".format(rde))
    print("rd = {}".format(rd))
    print("Nd = {}".format(Nd))
    print("---------------")
    plt.xlabel('U, В', fontsize=20)
    plt.ylabel('I, мкА', fontsize=20)
    plt.grid()
    ax.legend()
    plt.savefig(num + '.png')
    plt.show()
    return plt


def plotAll(plot1, plot2, plot3):
    return True


if __name__ == "__main__":
    main()