import matplotlib.pyplot as plt
import numpy as np
from numpy import linalg
from math import pi, sin

def Create_MNA(mat_cons):
    Y = np.zeros(mat_cons[0])
    B = np.zeros(mat_cons[1])
    C = np.zeros(mat_cons[2])
    D = np.zeros(mat_cons[3])

    return [Y,B,C,D]

class Voltage_Supply:
    def __init__(self, func, a, b):
        self.node_a = a - 1
        self.node_b = b - 1
        self.v = 0
        self.func = func
        self.branch_current_node = 0
        self.vp = 0
        self.ip = 0

    def set_voltages_currents(self, v, c):
        if self.node_b > 0 and self.node_a > 0:
            self.vp = v[self.node_a] - v[self.node_b]
        elif self.node_b <= 0 and self.node_a > 0:
            self.vp = v[self.node_a]
        elif self.node_b > 0 and self.node_b <= 0:
            self.vp = v[self.node_b]

        if self.needs_branches:
            self.ip = c[self.branch_current_node]

    def set_time(self, t):
        self.v = self.func(t)

    def needs_branches(self):
        return True

    def set_branch_current_node(self, a):
        self.branch_current_node = a

    def get_matrix(self, mat_cons):
        Y,B,C,D = Create_MNA(mat_cons)

        if self.node_b >= 0:
            B[self.node_b][self.branch_current_node] = 1
            C[self.branch_current_node][self.node_b] = 1
        if self.node_a >= 0:
            B[self.node_a][self.branch_current_node] = -1
            C[self.branch_current_node][self.node_a] = -1

        M = np.block([[Y,B],[C,D]])
        RHS_Top = np.zeros((mat_cons[0][0], 1))
        RHS_Bot = np.zeros((mat_cons[2][0], 1))
        RHS_Bot[self.branch_current_node] = self.v

        RHS = np.block([[RHS_Top], [RHS_Bot]])

        return [M, RHS]

class Resistance:
    def __init__(self, func, a, b):
        self.node_a = a - 1
        self.node_b = b - 1
        self.r = 0
        self.func = func
        self.vp = 0
        self.ip = 0

    def set_voltages_currents(self, v, c):
        if self.node_b > 0 and self.node_a > 0:
            self.vp = v[self.node_a] - v[self.node_b]
        elif self.node_b <= 0 and self.node_a > 0:
            self.vp = v[self.node_a]
        elif self.node_b > 0 and self.node_b <= 0:
            self.vp = v[self.node_b]

    def set_time(self, t):
        self.r = self.func(t)

    def needs_branches(self):
        return False

    def get_matrix(self, mat_cons):
        Y,B,C,D = Create_MNA(mat_cons)

        a = self.node_a
        b = self.node_b

        if b >= 0:
            Y[b,b] = self.r
        if a >= 0:
            Y[a,a] = self.r

        if a >= 0 and b >= 0:
            Y[a,b] = -self.r
            Y[b,a] = -self.r

        M = np.block([[Y,B], [C,D]])
        RHS = np.zeros((mat_cons[0][0] + mat_cons[2][0], 1))

        return [M, RHS]

class Inductor:
    def __init__(self, func, a, b):
        self.node_a = a - 1
        self.node_b = b - 1
        self.l = 0
        self.func = func
        self.previous_t = 0.0
        self.current_t = 0.0
        self.vp = 0
        self.ip = 0

    def set_voltages_currents(self, v, c):
        if self.node_b > 0 and self.node_a > 0:
            self.vp = v[self.node_a] - v[self.node_b]
        elif self.node_b <= 0 and self.node_a > 0:
            self.vp = v[self.node_a]
        elif self.node_b > 0 and self.node_b <= 0:
            self.vp = -v[self.node_b]

        if self.needs_branches:
            self.ip = c[self.branch_current_node]

    def set_time(self, t):
        self.previous_t = self.current_t
        self.current_t = t
        self.l = self.func(t)

    def needs_branches(self):
        return True

    def set_branch_current_node(self, a):
        self.branch_current_node = a

    def get_matrix(self, mat_cons):
        Y,B,C,D = Create_MNA(mat_cons)

        a = self.node_a
        b = self.node_b
        h = self.current_t - self.previous_t
        lh = self.l / h

        if self.node_b >= 0:
            B[self.node_b][self.branch_current_node] = -1
            C[self.branch_current_node][self.node_b] = -1
        if self.node_a >= 0:
            B[self.node_a][self.branch_current_node] = 1
            C[self.branch_current_node][self.node_a] = 1

        D[self.branch_current_node][self.branch_current_node] = -lh

        M = np.block([[Y,B],[C,D]])
        RHS_Top = np.zeros((mat_cons[0][0], 1))
        RHS_Bot = np.zeros((mat_cons[2][0], 1))
        RHS_Bot[self.branch_current_node] = -lh * self.ip

        RHS = np.block([[RHS_Top], [RHS_Bot]])

        return [M, RHS]

class Capacitor:
    def __init__(self, func, a, b):
        self.node_a = a - 1
        self.node_b = b - 1
        self.c = 0
        self.func = func
        self.previous_t = 0.0
        self.current_t = 0.0
        self.vp = 0
        self.ip = 0

    def set_voltages_currents(self, v, c):
        if self.node_b > 0 and self.node_a > 0:
            self.vp = v[self.node_a] - v[self.node_b]
        elif self.node_b <= 0 and self.node_a > 0:
            self.vp = v[self.node_a]
        elif self.node_b > 0 and self.node_b <= 0:
            self.vp = -v[self.node_b]

    def set_time(self, t):
        self.previous_t = self.current_t
        self.current_t = t
        self.c = self.func(t)

    def needs_branches(self):
        return False

    def set_branch_current_node(self, a):
        self.branch_current_node = a

    def get_matrix(self, mat_cons):
        Y,B,C,D = Create_MNA(mat_cons)

        a = self.node_a
        b = self.node_b

        h = self.current_t - self.previous_t
        ch = self.c / h

        if b >= 0:
            Y[b,b] = ch
        if a >= 0:
            Y[a,a] = ch

        if a >= 0 and b >= 0:
            Y[a,b] = -ch
            Y[b,a] = -ch

        M = np.block([[Y,B],[C,D]])
        RHS_Top = np.zeros((mat_cons[0][0], 1))
        RHS_Bot = np.zeros((mat_cons[2][0], 1))
        
        if a > 0:
            RHS_Top[a] = ch * self.vp
        if b > 0:
            RHS_Top[b] = - ch * self.vp

        RHS = np.block([[RHS_Top], [RHS_Bot]])

        return [M, RHS]

def run( netlist, t0, t1, points):
    counter = 0
    nodes = list()
    for j in netlist:
        if j.needs_branches():
            j.set_branch_current_node(counter)
            counter += 1
        nodes.append(j.node_a)
        nodes.append(j.node_b)

    max_node = max(nodes) + 1

    Y = (max_node,max_node)
    B = (max_node,counter)
    C = (counter,max_node)
    D = (counter,counter)

    M = (max_node+counter, max_node+counter)
    RHS = (max_node+counter, 1)

    mat_cons = [Y,B,C,D]

    t_space = np.linspace(t0,t1,points)[1:]

    answer_array = list()
    answer_array.append(list())

    for i in range(0,M[0]):
        answer_array.append(list())

    for t in range(0, len(t_space)):
        mat = np.zeros(M)
        rhs = np.zeros(RHS)
        for i in netlist:
            i.set_time(t_space[t])
            m, r = i.get_matrix(mat_cons)
            mat += m
            rhs += r

        ans = linalg.lstsq(mat,rhs)[0]

        voltages = ans[:max_node]
        currents = ans[max_node:]

        voltages = [v[0] for v in voltages]
        currents = [c[0] for c in currents]

        for j in netlist:
            j.set_voltages_currents(voltages, currents)

        answer_array[0].append(t_space[t])
        for i in range(0, len(ans)):
            answer_array[i+1].append(ans[i][0])

    return answer_array

def R_func(t):
    return 1/1000.0

def V1_func(t):
    #if t > 0.00001 * 0.5:
    #    return 5
    #else:
    #    return 0
    v = 0
    for i in range(1, 10, 2):
        v += (10.0/i)*sin(2*pi*i*1000000.0*t)
    return v

def L_func(t):
    return 1.0e-9

def C_func(t):
    return 1000.0e-12

def C1_func(t):
    return 1.0e-12

def V2_func(t):
    if t == 0.0:
        return -2.5
    else:
        return 0

netlist = [
    Voltage_Supply(V1_func, 0, 1),
    #Inductor(L_func, 1, 2),
    Capacitor(C1_func, 1, 2),
    Resistance(R_func, 1, 2),
    Capacitor(C_func, 2, 0),
    Resistance(R_func, 2, 0)
    #Voltage_Supply(V2_func, 0, 3)
]


ans = run(netlist, 0.0, 0.00001, 100000)

plt.plot(ans[0],ans[1], 'r')
plt.plot(ans[0],ans[2], 'g')
#plt.plot(t, outs[3])
plt.show()

plt.plot(ans[0], ans[3], 'r')
plt.show()