import matplotlib.pyplot as plt
from skrf import Network
from core.freqency import auto_select_multple_ports
from core.freqency import auto_select
from core.vectorFitting import VectorFitting
import numpy as np
from typing import Literal
import skrf as rf

def vector_fitting_info(vf:VectorFitting,sampled_freqs,sampled_responses,title:str='vf',parameter_type:Literal['y','s','z']='y'):
    for i in range(vf.network.nports):
        for j in range(vf.network.nports):
            rms_error = vf.get_rms_error(i,j,parameter_type=parameter_type)
            print(f'RMS Error Port {i+1} to Port {j+1}: {rms_error}')
            fitted_points = vf.get_model_response(i,j,sampled_freqs)
            if parameter_type=='y':
                input_points = vf.network.y[:,i,j]
            elif parameter_type=='s':
                input_points = vf.network.s[:,i,j]
            elif parameter_type=='z':
                input_points = vf.network.z[:,i,j]
            plt.figure(figsize=(20,12))
            plt.suptitle(f'{title} Port {i+1} to Port {j+1}')
            plt.subplot(2,2,1)
            plt.title('Magnitude')
            plt.plot(sampled_freqs,np.abs(sampled_responses[:,i,j]),'o', ms=4, color='red', label='Samples')
            plt.plot(vf.network.f,np.abs(input_points),'x', ms=4, color='blue', label='Input Samples')
            plt.plot(sampled_freqs,np.abs(fitted_points),'-', lw=2, color='k', label='Fit')
            plt.ylabel('Magnitude (dB)')
            plt.grid()
            plt.legend()
            plt.subplot(2,2,2)
            plt.title('Phase')
            plt.plot(sampled_freqs,np.angle(sampled_responses[:,i,j]),'o', ms=4, color='red', label='Samples')
            plt.plot(vf.network.f,np.angle(input_points),'x', ms=4, color='blue', label='Input Samples')
            plt.plot(sampled_freqs,np.angle(fitted_points),'-', lw=2, color='k', label='Fit')
            plt.ylabel('Phase (deg)')
            plt.xlabel('Frequency (GHz)')
            plt.grid()
            plt.legend()
            plt.subplot(2,2,3)
            plt.title('RMS_Error')
            plt.plot(rms_error, 'b', label='RMS Error')
            plt.ylabel('RMS Error (dB)')
            plt.grid()
            plt.legend()
            plt.savefig(f"outputs/{title.replace(' ','_')}_{i+1}_{j+1}.png")


original_network = Network('/tmp/paramer/simulation/3000/3000.s2p')
# original_network = rf.data.ring_slot

H,freqs = auto_select_multple_ports(original_network.y,original_network.f,max_points=20)
fitted_network = Network(y=H,f=freqs)
vf = VectorFitting(fitted_network)
vf.vector_fit(n_poles_cmplx=2,n_poles_real=0,parameter_type='y')
# vf.auto_fit(parameter_type='y')

vector_fitting_info(vf,original_network.f,original_network.y,parameter_type='y')