from core.orthonormal_basis import generate_laguerre_basis
from core.sk_iter import generate_starting_poles
import numpy as np
import skrf as rf
import matplotlib.pyplot as plt
import sympy as sp
from scipy.linalg import null_space


network = rf.Network("/tmp/paramer/simulation/4000/4000.s2p")
freqs = network.f
s = freqs * 2j * np.pi
vf = rf.VectorFitting(network)
vf.vector_fit(n_poles_real=2, n_poles_cmplx=1,parameter_type="y")

poles = vf.poles
residues = vf.residues
y = vf.network.y

fig, ax = plt.subplots(2, 2)
fig.set_size_inches(12, 8)
vf.plot("mag",0,0,freqs,ax=ax[0][0],parameter="y")
rms_error11 = vf.get_rms_error(0,0,"y")
print("rms_error11",rms_error11)
vf.plot("mag",1,0,freqs,ax=ax[1][0],parameter="y")
rms_error21 = vf.get_rms_error(1,0,"y")
print("rms_error21",rms_error21)
vf.plot("mag",0,1,freqs,ax=ax[0][1],parameter="y")
rms_error12 = vf.get_rms_error(0,1,"y")
print("rms_error12",rms_error12)
vf.plot("mag",1,1,freqs,ax=ax[1][1],parameter="y")
rms_error22 = vf.get_rms_error(1,1,"y")
print("rms_error22",rms_error22)
fig.tight_layout()
plt.show()
plt.savefig(f"img.png")

diag_values = [y[i][0][0] for i in range(len(y))]
H = np.diag(diag_values)
P = 4
start_poles = generate_starting_poles(P,beta_min=1e8,beta_max=freqs[-1])
basis = generate_laguerre_basis(start_poles,s).T
print("start_poles",start_poles)
print("basis",basis)

# first step iteration
# A*x = b

A11 = np.real(basis)
print("A11 shape:",A11.shape)
A12 = np.real(- H @ basis[:,1:])
print("A12 shape:",A12.shape)
# print("A11",A11)
A21 = np.imag(basis)
print("A21 shape:",A21.shape)
A22 = np.imag(- H @ basis[:,1:])
print("A22 shape:",A22.shape)
A1 = np.hstack([A11,A12])
A2 = np.hstack([A21,A22])
A = np.vstack([A1,A2])
print ("A shape:",A.shape)

b1 = np.real(H @ basis[:,0])
b2 = np.imag(H @ basis[:,0])
b = np.hstack([b1,b2])
Q, R = np.linalg.qr(A, mode='reduced')
print("Q :",Q)
print("R :",R)
print("b shape:",b.shape)
# x = np.linalg.solve(R, Q.T @ b)
x_star, residuals, rank, s = np.linalg.lstsq(A, b, rcond=None)

# print("x_qr",x_star)
print("residuals",residuals)
# print("rank",rank)
# print("s",s)

x_ne = np.linalg.inv(A.T @ A) @ A.T @ b
print("x_ne",x_ne)

# sk iteration
target_vec = np.array([1+0j] + list(x_ne[len(x_ne)//2+1:]))
D_t = basis @ target_vec
print("D_t",D_t)
K = 25

for i in range(K):
    print(f"Iteration {i+1}/{K}")
    A11 = np.real(basis / D_t[:, np.newaxis])
    A12 = np.real(- H @ basis[:,1:] / D_t[:, np.newaxis])
    A21 = np.imag(basis / D_t[:, np.newaxis])
    A22 = np.imag(- H @ basis[:,1:] / D_t[:, np.newaxis])
    A1 = np.hstack([A11,A12])
    A2 = np.hstack([A21,A22])
    A = np.vstack([A1,A2])
    b1 = np.real(H @ basis[:,0])
    b2 = np.imag(H @ basis[:,0])
    b = np.hstack([b1,b2])
    x_star, residuals, rank, s = np.linalg.lstsq(A, b, rcond=None)
    # print("x_lstsq",x_star)
    # print("residuals",residuals)
    # print("rank",rank)
    # print("s",s)
    x_ne = np.linalg.inv(A.T @ A) @ A.T @ b
    # print("x_ne",x_ne)
    target_vec = np.array([1+0j] + list(x_ne[len(x_ne)//2+1:]))
    D_t_pre = D_t
    D_t = basis @ target_vec
    print("D_t/D_t_pre",np.linalg.norm(D_t)/np.linalg.norm(D_t_pre))
    # print("D_t",D_t)
    n_target_vec = np.array(list(x_ne[:len(x_ne)//2+1]))
    N_t = basis @ n_target_vec
    # print("H = N_t / D_t",np.linalg.norm(H - N_t / D_t)/np.linalg.norm(H))