"""
    Test script for Q-factor fitting (pure Python version using numpy)
    ==================================================================

      Description
      -----------

         This software is from:
         "Q-factor Measurement by using a Vector Network Analyser",
         A. P. Gregory, National Physical Laboratory Report MAT 58 (2021)

         Fits to transmission (S21) data for two-port resonator
         (with frequency-dependent leakage) by using the NLQFIT8 algorithm.

         Requires script qfitmod8.py
         Test data is read from file Figure23.txt (shown in Figure 23 of MAT 58)

         Tested with Python 2.7 and Python 3.8


      Creative Commons CC0 Public Domain Dedication
      ---------------------------------------------

         This software is released under the Creative Commons CC0
         Public Domain Dedication. You should have received a copy of
         the Legal Code along with this software.
         If not, see <http://creativecommons.org/publicdomain/zero/1.0/>.


     Change log
     ----------

        :30-Oct-2021: Released under CC0.
        :23-Aug-2022: Added code for calculating fitted Trace, and plotting on a
                      polar chart with matplotlib. Uncomment these lines if required.
                      S-parameter data file is now read with loadtxt.

"""
import numpy as np
from qfitmod8 import *
import cmath, math

# Load frequency and S-parameter data from a file.
fn = 'Figure23.txt';
F, R, I, _Mg, _Ph = np.loadtxt(fn,comments='%',unpack=True)
N = len(F)
D = [complex(R[i],I[i]) for i in range(N)]
#D = R+1j*I # or use numpy array

# De-embed cables
ncablelen = 1.2 # root_eps * cable length in metres
D = [D[i]*cmath.exp(complex(0.0,-2E9*math.pi*F[i]*ncablelen/2.99792458E8)) for i in range(N)]

print('Demonstration of Q-factor fitting (pure Python and numpy implementation)')
print('S21 data is taken from file Figure23.txt (plotted in Fig. 23 of MAT58)')
print ()

# The convergence algorithm uses a number of steps set by loop_plan, a string of characters as follows:
#   f - fit once without testing for convergence
#   c - repeated fit, iterating until convergence is obtained
#   w - re-calculate weighting factors on basis of previous fit
#       Initially the weighting factors are all unity.
#   r - Exclude 10% of points with largesr residual error from the fit.
loop_plan ='fwfwfwc'

# Find peak in |S21| - this is used to give initial value of freq.
# Tol is 1.0E-5 * |S21| at peak.
Mg = np.absolute(D)
index_max = np.argmax(Mg)
Tol = Mg[index_max]*1.0E-5
Fseed = F[index_max]

# Set Qseed: An order-of-magnitude estimate for Q-factor
mult = 5.0 # Not critical. A value of around 5 will work well for a wide sweep-range as 
#            shown in Fig. 6(a), and for a narrow (optimised) sweep-range as shown in Fig 6(b).
Qseed = mult*Fseed/(F[-1]-F[0])


quiet = False
if not quiet: print('Initial values for iteration:  Freq=',Fseed,' QL=',Qseed)

# Step 1: Initial unweighted fit --> solution vector
sv = initialfit(F, D, N, Fseed, Qseed)
print('Initial solution (eqn. 17):');
print('  a = %6.4f + j*%6.4f' % (sv[0],sv[1]));
print('  b = %6.4f + j*%6.4f' % (sv[2],sv[3]));
print('  QL = %3.3f' % sv[4]);

print('\nOptimising ....')
# Step 2: Optimised weighted fit --> result vector
mv, weighting_ratio, number_iterations, RMS_Error = optimisefit8(F, D, N, Fseed, sv, loop_plan, Tol, quiet)
m1,m2,m3,m4,m8,m9,QL,FL = mv

print('Optimised solution (eqn. 43):')
print('  Fr = %10.8f' % FL)
print('  QL = %6.2f' % QL)
print('  Number of iterations = %i' % number_iterations)
if weighting_ratio is not None: print('  Weighting ratio = %5.3f' % weighting_ratio)
print('  RMS_Error = %10.8f' % RMS_Error)
print()

# Now calculate unloaded Q-factor and some other useful quantities.
# Require scaling_factor_A (the reciprocal of |S21| of a thru in place of resonator)
scaling_factor_A = 1/0.949    # 1/|S21_thru|
trmode='transmission'
ifail, p = GetUnloadedData(mv, scaling_factor_A, trmode, quiet)
Qo, cal_diam, cal_gamma_V, cal_gamma_T = p
print('Q-factor of uncoupled two-port resonator (unloaded Q-factor)')
print ('  scaling factor A = 1 / |S21|_thru')
print('  Qo = %6.2f' % Qo)
print('  Qcircle diameter = %6.4f' % cal_diam)

#### Calculate fitted trace and then plot (uncomment if required) ###
"""
import matplotlib.pyplot as plt
from numpy import angle, absolute, empty, vectorize
m6a = QL/FL
t = 2*(F-FL)/FL
vcmplex = vectorize(complex)
y = 1.0/vcmplex(1.0,2*(m6a*F-QL))
Fit = complex(m1,m2) + complex(m8,m9)*t + complex(m3,m4)*y

fig, ax = plt.subplots(subplot_kw={'projection': 'polar'})
ax.plot(angle(D), Mg,'-k',label='Measured')
ax.plot(angle(Fit), absolute(Fit),'ro',markevery=5,markersize=2,label='Fit')
ax.set_rmax(0.008)
ax.set_rticks([0.002,0.004,0.006,0.008])
ax.set_rlabel_position(-20)  # Move radial labels away from plotted line
ax.grid(True)
ax.legend();
ax.set_title("Measured and fitted Q-circles", va='bottom')
plt.show()

fig2, ax2 = plt.subplots()
ax2.plot(F, Mg,'-k',label='Measured')
ax2.plot(F, absolute(Fit),'ro',markevery=5,markersize=2,label='Fit')
ax2.grid(True)
plt.xlabel('Frequency (GHz)')
plt.ylabel('|S21|')
ax2.legend();
ax2.set_title("Measured and fitted magnitude", va='bottom')
plt.show()
"""

"""
OUTPUT:

Demonstration of Q-factor fitting (pure Python and numpy implementation)
S21 data is taken from file Figure23.txt (plotted in Fig. 23 of MAT58)

Initial values for iteration:  Freq= 9.76052255  QL= 11817.41443729872
Initial solution (eqn. 17):
  a = 8.9408 + j*2.1298
  b = 0.0054 + j*-0.0045
  QL = 4664.241

Optimising ....
Op f, Iteration 1, RMS_Error 0.00039672

Op w, Calculate weights
Op f, Iteration 2, RMS_Error 0.00006278

Op w, Calculate weights
Op f, Iteration 3, RMS_Error 0.00000923

Op w, Calculate weights
Op c, Iteration 4, RMS_Error 0.00000910
      Iteration 5, RMS_Error 0.00000910

Optimised solution (eqn. 43):
  Fr = 9.76015571
  QL = 4760.04
  Number of iterations = 5
  Weighting ratio = 5.078
  RMS_Error = 0.00000910

Q-factor of uncoupled two-port resonator (unloaded Q-factor)
  scaling factor A = 1 / |S21|_thru
  Qo = 4789.49
  Qcircle diameter = 0.0061

"""
