import numpy as np
from core.sk_iter import generate_starting_poles
from scipy.linalg import block_diag
import skrf as rf
from skrf import VectorFitting
from core.freqency import auto_select_multple_ports
import matplotlib.pyplot as plt
import random as rnd

class MultiplePortQR:
    def __init__(self,H,freqs,poles,weights=None,passivity=True,dc_enforce=True,fit_constant=True,fit_proportional=False):
        self.least_squares_rms_error  = None
        self.least_squares_condition = None
        self.eigenval_condition = None
        self.eigenval_rms_error = None

        self.dc_tol = 1e-18

        self.dc_enforce = dc_enforce
        self.fit_constant = fit_constant
        self.fit_proportional = fit_proportional
        
        # self.H = H
        # self.freqs = freqs
        self.freqs = freqs
        self.H = H
        self.ports = H.shape[1]
        self.s = self.freqs * 2j * np.pi
        self.P = len(poles)
        self.poles = poles
        self.Phi = self.generate_basis(self.s, self.poles)
        self.A = self.matrix_A(self.poles)
        self.B = self.vector_B(self.poles)
        self.Cw,self.w0,self.e = self.fit_denominator(self.H, weights=weights)
        self.D = self.w0

        self.Cr = None
        
        z = np.linalg.eigvals(self.A - self.B @ self.Cw)
        p_next = -z

        if passivity:
            self.next_poles = self.passivity_enforce(p_next)
        else:
            self.next_poles = p_next
        # z = np.where(np.real(z) < 0, z, -np.conj(z))  # enforce LHP
        # self.next_poles = np.sort_complex(z)
        self.eigenval_condition = np.linalg.cond(self.A - self.B @ self.Cw)
        self.eigenval_rms_error = np.sqrt(np.mean(np.abs(np.real(z) - np.real(poles))**2 + np.abs(np.imag(z) - np.imag(poles))**2))

        self.Dt = self.eval_Dt_state_space()
        self.delta = self.Dt / weights if weights is not None else self.Dt
        pass

    def passivity_enforce(self,poles):
        """enforce poles' real parts to be negative"""
        enforced_poles = []
        for pole in poles:
            if pole.real > 0:
                pole = -np.conj(pole)
            enforced_poles.append(pole)
        return enforced_poles

    def eval_Dt_state_space(self):
        """Return D(s_k)=C(s_k I - A)^(-1)B + D for all k (complex 1D array)."""
        s = 1j * 2*np.pi * np.asarray(self.freqs, float).ravel()
        A = np.asarray(self.A, np.complex128); n = A.shape[0]
        B = np.asarray(self.B, np.complex128).reshape(n, 1)
        C = np.asarray(self.Cw, float).reshape(1, n)
        D = self.D
        I = np.eye(n, dtype=np.complex128)
        out = np.empty_like(s, dtype=np.complex128)
        for k, sk in enumerate(s):
            DS = D + (C @ np.linalg.inv(sk*I - A) @ B)
            out[k] = DS[0, 0]
        return out
    

    def generate_basis(self,s, poles):
        """Real basis of (15)-(16); returns Φ(s) and a layout for packing C."""
        cols = []
        i = 0
        while i < len(poles):
            p = poles[i]
            if p.real > 0:
                raise ValueError("poles must be in the LHP")
            if i+1 < len(poles) and np.isclose(poles[i+1], np.conj(p)):
                pc = poles[i+1]
                phi1 = 1/(s - p) + 1/(s - pc)                 # eq (15)generate_basis
                phi2 = 1j*(1/(s - p) - 1/(s - pc))            # eq (16)  (fixed sign)
                cols += [phi1, phi2]
                i += 2
            else:
                cols.append(1/(s - p))
                i += 1
        Phi = np.column_stack(cols).astype(np.complex128)
        return Phi

    def matrix_A(self, poles):
        def A_block(p):
            if abs(p.imag) < 1e-14:
                return np.array([[p.real]], float)                  # A_p = [ p ]
            return np.array([[p.real,  p.imag],                     # A_p = [[Re p, Im p],
                              [-p.imag, p.real]], float)            #         [-Im p, Re p]]
        A = None; i = 0
        while i < len(poles):
            p = poles[i]
            Ab = A_block(p)
            if i+1 < len(poles) and np.isclose(poles[i+1], np.conj(p)): i += 2
            else: i += 1
            A = Ab if A is None else block_diag(A, Ab)
        return A
    
    def vector_B(self, poles):
        def B_block(p):
            return np.array([[1.0]], float) if abs(p.imag)<1e-14 else np.array([[2.0],[0.0]], float)
        B = None; i = 0
        while i < len(poles):
            p = poles[i]
            Bb = B_block(p)
            if i+1 < len(poles) and np.isclose(poles[i+1], np.conj(p)): i += 2
            else: i += 1
            B = Bb if B is None else np.vstack([B, Bb])
        return B
    
    def fit_denominator(self, H, weights=None, d0 = 1.0):
        """
        Solve formula (70) on the real basis Φ to obtain:
          - d (real)  → packs into C for this state's block structure
          - gamma (complex)
        Optional 'weights' (K,) apply row scaling: SK weighting if 1/|D_prev|.
        """
        K, N = self.Phi.shape
        one  = np.ones((K, 1), np.complex128)
        Phi   = self.Phi
        dc_tol = 1e-18
        has_dc = self.dc_enforce and self.freqs[0] < dc_tol
        keep = np.ones(K, dtype=bool)

        # SK weighting (applied only to the (73) rows we keep in LS)

        if has_dc:
            # Enforce DC response exactly:
            k0   = int(np.argmin(np.abs(self.freqs)))
            keep[k0] = False

        if self.fit_constant:  
            Phi_w = np.hstack([one, Phi]) 
            index = 0
            M_kp = None
            for i in range(self.ports):
                for j in range(self.ports):
                    M0 = np.zeros((K,N*self.ports**2),dtype=complex)
                    M0[:,index*N:(index+1)*N] = Phi
                    M0 = np.hstack([M0, -(H[:,i,j].reshape(-1,1) * Phi_w)]).reshape((K, -1))[keep,:]       # (K, 2N),   complex
                    index+=1
                    M_kp = M0 if M_kp is None else np.vstack([M_kp, M0])
            assert M_kp is not None
        else:
            index = 0
            M_kp = None
            for i in range(self.ports):
                for j in range(self.ports):
                    M0 = np.zeros((K,N*ports**2),dtype=complex)
                    M0[:,index*N:(index+1)*N] = Phi
                    M0 = np.hstack([M0, -(H[:,i,j].reshape(-1,1) * Phi)]).reshape((K, -1))[keep,:]       # (K, 2N),   complex
                    index+=1
                    M_kp = M0 if M_kp is None else np.vstack([M_kp, M0])
            assert M_kp is not None

        if weights is None:
            weights_kp = np.diag(np.ones(len(freqs[keep]) * self.ports**2, np.complex128))
        else:
            weights_kp0 = weights[keep]
            weights0 = []
            for i in range(self.ports **2 ):
                for res in weights_kp0:
                    weights0.append(1/res)
            weights_kp = np.diag(np.array(weights0))


        if has_dc:
            M_w_kp     = weights_kp @ M_kp
            A_re  = np.real(M_w_kp)
            A_im  = np.imag(M_w_kp)
            mask    = np.ones(K, dtype=bool); mask[k0] = False
            # exact (unweighted) DC rows:
            # A_dc_re = np.real(M_kp).reshape(1, -1)
            # A_dc_im = np.imag(M_kp).reshape(1, -1)
        else:
            M_w_kp    = weights_kp @ M_kp
            A_re = np.real(M_w_kp)
            A_im = np.imag(M_w_kp)
            # A_dc_re = A_dc_im = None

        A_blocks = [A_re, A_im]

        if self.fit_constant:
            Hk_kp = None
            for i in range(self.ports):
                for j in range(self.ports):
                    Hk_kp0 = H[:,i,j][keep]
                    Hk_kp = Hk_kp0 if Hk_kp is None else np.hstack([Hk_kp, Hk_kp0])
            assert Hk_kp is not None
            Hk_sum = np.sum(np.abs(Hk_kp)**2)
            beta     = float(np.sqrt(Hk_sum))
            mean_row = (beta / weights_kp.shape[0]) * np.sum(Phi_w, axis=0)
            A_w0      = np.concatenate([np.zeros(N*self.ports**2, float),
                                np.real(mean_row).astype(float)]
                                ).reshape(1, -1)
            b_w0      = np.array([beta], float)

            A_blocks += [A_w0]
            m = A_re.shape[0] + A_im.shape[0]
            b = np.zeros(m, float)
            b = np.concatenate([b, b_w0])
        else:
            H_kp = None
            for i in range(self.ports):
                for j in range(self.ports):
                    H_kp0 = weights_kp @ (H[:,i,j]).reshape(1,-1)[keep,:]
                    H_kp = H_kp0 if H_kp is None else np.hstack([H_kp, H_kp0])
            assert H_kp is not None
            H_kp = H_kp.reshape(-1,1)
            
            b_re = np.real(d0 * H_kp)
            b_im = np.imag(d0 * H_kp)
            b = np.concatenate([b_re.ravel(), b_im.ravel()]).astype(float)

        # ---- build final stacked-real system ----
        

        # if A_dc_re is not None:
        #     A_blocks += [A_dc_re, A_dc_im]
        #     b = np.concatenate([b, np.zeros(2, float)])   # DC rows → 0

        # ---- QR solve for x = [c_H (N); c_w (N+1)] ----
        A = np.vstack(A_blocks).astype(float)
        Q, R = np.linalg.qr(A, mode="reduced")
        
        if self.fit_constant:
            Q2 = Q[:,Phi.shape[1] * self.ports**2:]
            R22 = R[Phi.shape[1] * self.ports**2:,Phi.shape[1] * self.ports**2:]
        else:
            Q2 = Q[:,Phi.shape[1] * self.ports**2:]
            R22 = R[Phi.shape[1] * self.ports**2:,Phi.shape[1] * self.ports**2:]

        x = np.linalg.solve(R22, Q2.T @ b)

        # diagnostics
        resid = Q2 @ R22 @ x - b
        self.least_squares_rms_error = float(np.sqrt(np.mean(resid**2)))
        self.least_squares_condition = float(np.linalg.cond(R))

        # split cw and return
        # cw = x[N:]                     # last (N+1) entries = [w0, w_1..w_N]
        # w0 = float(cw[0])
        # Cw = cw[1:].reshape(1, N)      # row vector (1, N)
        return self.extract_Cw_d_e(x,N,d0)
    
    def extract_Cw_d_e(self,C,N,d0=1.0):
        if self.fit_proportional and self.fit_constant:
            d = C[1]
            e = C[0]
            return C[2:].reshape(1, -1), d, e
        elif self.fit_proportional and not self.fit_constant:
            d = 0.0
            e = C[0]
            return C[1:].reshape(1, -1), d, e
        elif not self.fit_proportional and self.fit_constant:
            d = C[0]
            e = 0.0
            return C[1:].reshape(1, -1), d, e
        else:
            return C.reshape(1, -1), d0, 0.0
    
    
    def non_bias_Cr(self,w0):
        A = np.asarray(self.Phi)
        den = np.diag((w0 + self.Phi @ self.Cw.T).ravel())
        Cr = []
        for i in range(self.ports):
            Cr.append([])
            for j in range(self.ports):
                b = np.asarray(den) @ self.H[:,i,j].reshape(-1,1)
                Cr_ij, residuals, rank, s = np.linalg.lstsq(A, b, rcond=None)
                Cr[i].append(Cr_ij)
        return Cr
    
    def evaluate(self,freqs, w0):
        H  = np.zeros((len(freqs),self.ports,self.ports),dtype=complex)
        s = 1j * 2*np.pi * np.asarray(freqs, float).ravel()
        phi = self.generate_basis(s, self.poles)
        den = w0 + phi @ self.Cw.T
        if self.Cr is None:
            self.Cr = self.non_bias_Cr(w0=w0)
        for i in range(self.ports):
            for j in range(self.ports):
                num = phi @ self.Cr[i][j]
                H[:,i,j] = (num / den).reshape(1,-1)
        return H
    
def noise(n:complex,coeff:float=0.05):
    noise_r = rnd.gauss(-coeff * n.real, coeff * n.real)
    noise_i = rnd.gauss(-coeff * n.imag, coeff * n.imag)
    return complex(n.real + noise_r, n.imag + noise_i)

if __name__ == "__main__":
    start_point = 0
    network = rf.Network("/tmp/paramer/simulation/3500/3500.s2p")
    ports = network.nports
    K = 10

    full_freqences = network.f[start_point:]
    noised_sampled_points = network.y[start_point:,:,:]
    sampled_points = network.y[start_point:,:,:]

    # noised_sampled_points = - network.y[start_point:,0,1].reshape(-1,1,1)
    # sampled_points = network.y[start_point:,1,1].reshape(-1,1,1)

    H,freqs = auto_select_multple_ports(noised_sampled_points,full_freqences,max_points=20)
    poles = generate_starting_poles(2,beta_min=1e4,beta_max=freqs[-1]*1.1)
    
    Dt_1 = np.ones((len(freqs),1),np.complex128)
    # Levi step (no weighting):
    basis = MultiplePortQR(H,freqs,poles=poles)
    Dt = basis.Dt
    poles = basis.next_poles

    print("Levi step (no weighting):")
    print("A:",basis.A)
    print("B:",basis.B)
    print("C:",basis.Cw)
    print("D:",basis.D)
    print("next_pozles:",basis.next_poles)
    print("Dt:",Dt, "norm:",np.linalg.norm(Dt))

    # SK weighting (optional, after first pass):
    least_squares_condition = []
    least_squares_rms_error = []
    eigenval_condition = []
    eigenval_rms_error = []
    for i in range(K):
        basis = MultiplePortQR(H,freqs,poles=poles,weights=Dt)
        Dt_1 = Dt
        Dt = basis.Dt
        poles = basis.next_poles
        print(f"SK Iteration {i+1}/{K}")
        print("A:",basis.A)
        print("B:",basis.B)
        print("C:",basis.Cw)
        print("D:",basis.D)
        print("z:",basis.next_poles)
        print("Dt:",Dt)
        print("Dt/Dt-1",np.linalg.norm(Dt) / np.linalg.norm(Dt_1))
        least_squares_condition.append(basis.least_squares_condition)
        least_squares_rms_error.append(basis.least_squares_rms_error)
        eigenval_condition.append(basis.eigenval_condition)
        eigenval_rms_error.append(basis.eigenval_rms_error)

    # H11_evaluated = basis.evaluate_pole_residue(network.f[1:],poles,basis.C[0])
    H_evaluated = basis.evaluate(full_freqences, w0=basis.w0)
    fitted_points = H_evaluated
    sliced_freqences = freqs

    input_points = H
    for i in range(ports):
        for j in range(ports):
            fig, axes = plt.subplots(3, 2, figsize=(15, 16), sharex=False)
            ax00 = axes[0][0]    
            ax00.plot(full_freqences, np.abs(sampled_points[:,i,j]), 'o', ms=4, color='red', label='Samples')
            ax00.plot(full_freqences, np.abs(fitted_points[:,i,j]), '-', lw=2, color='k', label='Fit')
            ax00.plot(sliced_freqences, np.abs(input_points[:,i,j]), 'x', ms=4, color='blue', label='Input Samples')
            ax00.set_title(f"Response i={i+1}, j={j+1}")
            ax00.set_ylabel("Magnitude")
            ax00.legend(loc="best")

            ax01 = axes[0][1]
            ax01.set_title(f"Response i={i+1}, j={j+1}")
            ax01.set_ylabel("Phase (deg)")
            ax01.plot(full_freqences, np.angle(sampled_points[:,i,j],deg=True), 'o', ms=4, color='red', label='Samples')
            ax01.plot(full_freqences, np.angle(fitted_points[:,i,j],deg=True), '-', lw=2, color='k', label='Fit')
            ax01.plot(sliced_freqences, np.angle(input_points[:,i,j],deg=True), 'x', ms=4, color='blue', label='Input Samples')
            ax01.legend(loc="best")

            # ax00 = axes[0][0]    
            # ax00.plot(full_freqences, np.real(sampled_points[:,i,j]), 'o', ms=4, color='red', label='Samples')
            # ax00.plot(full_freqences, np.real(fitted_points[:,i,j]), '-', lw=2, color='k', label='Fit')
            # ax00.plot(sliced_freqences, np.real(input_points[:,i,j]), 'x', ms=4, color='blue', label='Input Samples')
            # ax00.set_title(f"Response i={i+1}, j={j+1}")
            # ax00.set_ylabel("Real Part")
            # ax00.legend(loc="best")

            # ax01 = axes[0][1]
            # ax01.set_title(f"Response i={i+1}, j={j+1}")
            # ax01.set_ylabel("Imag Part")
            # ax01.plot(full_freqences, np.imag(sampled_points[:,i,j]), 'o', ms=4, color='red', label='Samples')
            # ax01.plot(full_freqences, np.imag(fitted_points[:,i,j]), '-', lw=2, color='k', label='Fit')
            # ax01.plot(sliced_freqences, np.imag(input_points[:,i,j]), 'x', ms=4, color='blue', label='Input Samples')
            # ax01.legend(loc="best")

            ax10 = axes[1][0]
            ax10.plot(least_squares_condition, label='Least Squares Condition')
            ax10.set_title("least_squares_condition")
            ax10.set_ylabel("Magnitude")
            ax10.legend(loc="best")

            ax11 = axes[1][1]
            ax11.plot(least_squares_rms_error, label='Least Squares RMS Error')
            ax11.set_title("least_squares_rms_error")
            ax11.set_ylabel("Magnitude")
            ax11.legend(loc="best")

            ax20 = axes[2][0]
            ax20.plot(eigenval_condition, label='Eigenvalue Condition')
            ax20.set_title("eigenval_condition")
            ax20.set_ylabel("Magnitude")
            ax20.legend(loc="best")

            ax21 = axes[2][1]
            ax21.plot(eigenval_rms_error, label='Eigenvalue RMS Error')
            ax21.set_title("eigenval_rms_error")
            ax21.set_ylabel("Magnitude")
            ax21.legend(loc="best")
            fig.tight_layout()
            plt.savefig(f"MultiplePortQR_port_{i+1}{j+1}.png")
            print(f"Saved MultiplePortQR_port_{i+1}{j+1}.png")