"""
    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 absorption (notch)
         resonator by using the NLQFIT6 algorithm.

         Requires script qfitmod6.py
         Test data is read from file Figure27.txt (as used in Figure 27 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.
        :09-Aug-2024: Added code for calculating Qo with fix for pivoting of Q-circle
                      (uses new function GetQoNotchPivotedQCircle in qfitmod6.py).
                      Note that to get better agreement with Probst resonator_tools
"""
import numpy as np
from qfitmod6 import *

# Load frequency and S-parameter data from a file.
fn = 'Figure27.txt';
F, R, I = np.loadtxt(fn,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 if they are long (not needed here)

print('Demonstration of Q-factor fitting (pure Python and numpy implementation)')
print('S21 data is taken from file Figure27.txt (plotted in Fig. 27 of MAT58)')
print('Superconducting absorption resonator')
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 - Identify 10% largest outliers from previous fit. Remove the effect of these.
loop_plan ='fwfwc'
#loop_plan ='fwfrwc'

# Find notch in |S21| - this is used to give initial value of freq.
Mg = np.absolute(D)
index_min = np.argmin(Mg)
Fseed = F[index_min]
Tol = 1.0E-5*np.argmax(Mg) # argmax(Mg) could be much less than 1.0 as cables
                              # through cryostat wall attenuate the signal

# 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,'GHz, QL=',Qseed)

# Step 1: Initial unweighted fit --> solution vector
N = len(F)
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
##mv, weighting_ratio, number_iterations, RMS_Error = optimisefit7(F, D, N, Fseed, sv, loop_plan, Tol, quiet)
##m1,m2,m3,m4,QL,FL,m7a = mv
##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

ifail, p = GetUnloadedData(mv, 'AUTO', 'absorption', quiet)
Qo, cal_diam, cal_gamma_V, cal_gamma_T = p

print('Optimised solution (eqn. 22):')
print('  Fr = %10.8f' % FL)
print('  QL = %6.2f ' % QL)
print('  Qo = %6.2f ' % Qo)
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()

# For asymmetric notch resonator
# Added August 2024
ifail, PivotCorrData = GetQoNotchPivotedQCircle(mv, quiet)
corr_Qo, rad_touching_circle, corr_diam, corr_gamma_V, corr_gamma_T = PivotCorrData
print('  Qo(with experimental pivot correction) = %6.2f'%corr_Qo)
print()
print ('Freq. span for 70% Q-circle = ',float(2.*FL/QL), 'GHz')

#### 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),'-r',label='Fit')
ax.set_rmax(0.4)
ax.set_rticks([0.1,0.2,0.3,0.4])
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')
print('Note: Figure 27 in MAT 58 has some aesthetic adjustments (rotational and scaling factors)')

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 Figure27.txt (plotted in Fig. 27 of MAT58)
Superconducting absorption resonator

Initial values for iteration:  Freq= 6.0722455 GHz, QL= 145792.20888353878
Initial solution (eqn. 17):
  a = -17072.3098 + j*9047.0761
  b = 0.0063 + j*0.0168
  QL = 53871.014

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

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

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

Notch/absorption Qo calculation: Supplied scaling_factor_A is "Auto", so using fitted data to estimate it
  Q-circle diam = 0.970
Optimised solution (eqn. 22):
  Fr = 6.07225567
  QL = 56019.85
  Qo = 1846782.51
  Number of iterations = 4
  Weighting ratio = 4.714
  RMS_Error = 0.01394722

Notch/absorption Qo calculation with pivoted Q-circle
  Radius of touching circle = 1.012
  Corrected Q-circle diam = 0.958
  Qo(with experimental pivot correction) = 1333670.84

Freq. span for 70% Q-circle =  0.00021678942580485164 GHz

"""
