"""
    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 FL and QL to transmission (S21) data for two-port resonator
         by using the NLQFIT6 algorithm.

         Requires script qfitmod6.py
         Test data is read from file Figure6b.txt (as used in Figure 6(b) 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 qfitmod6 import *


# Load frequency and S-parameter data from a file.
# File 'Figure6b.txt' contains S21 data for Fig. 6(b) in MAT 58
fn = 'Figure6b.txt';
F, R, I = 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

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

# De-embed cables if they are long (not needed here)

# 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 ='fwfwc'

# 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 = optimisefit6(F, D, N, Fseed, sv, loop_plan, Tol, quiet)
m1,m2,m3,m4,QL,FL = mv

print('Optimised solution (eqn. 22):')
print('  Fr = %10.8f   (expect 3.987848 GHz)' % FL)
print('  QL = %6.2f  (expect 7454)' % 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.874    # 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  (expect 7546)' % Qo)
print('  Qcircle diameter = %6.4f' % cal_diam)
print('  GammaV = %10.8f +j* %10.8f (S21 detuned = leakage vector)' % (cal_gamma_V.real, cal_gamma_V.imag))
print('  GammaT = %10.8f +j* %10.8f (S21 tuned)' % (cal_gamma_T.real, cal_gamma_T.imag))

# Tests
assert abs(FL-3.987848)<0.000001
assert abs(Qo-7546)<1.0
assert abs(QL-7454)<1.0


#### Calculate fitted trace and then plot (uncomment if required) ###
"""
import matplotlib.pyplot as plt
from numpy import angle, absolute, empty, vectorize
m6a = QL/FL
vcmplex = vectorize(complex)
y = 1.0/vcmplex(1.0,2*(m6a*F-QL))
Fit = complex(m1,m2) + 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.015)
ax.set_rticks([0.005,0.01, 0.015])
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 Figure6b.txt (plotted in Fig. 6(b) of MAT58)

Initial values for iteration:  Freq= 3.98783686  QL= 18636.493410597937
Initial solution (eqn. 17):
  a = 0.8104 + j*-1.6928
  b = 0.0077 + j*-0.0071
  QL = 7440.848

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

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

Op w, Calculate weights
Op c, Iteration 3, RMS_Error 0.00001216
      Iteration 4, RMS_Error 0.00001216

Optimised solution (eqn. 22):
  Fr = 3.98784835   (expect 3.987848 GHz)
  QL = 7454.48  (expect 7454)
  Number of iterations = 4
  Weighting ratio = 5.149
  RMS_Error = 0.00001216

Q-factor of uncoupled two-port resonator (unloaded Q-factor)
  scaling factor A = 1 / |S21|_thru
  Qo = 7545.58  (expect 7546)
  Qcircle diameter = 0.0121
  GammaV = -0.00008895 +j* 0.00003852 (S21 detuned = leakage vector)
  GammaT = 0.00849357 +j* -0.00845349 (S21 tuned)

"""
