"""
    qfitmod7.py
    ===========

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

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

        Python implementation of the NLQFIT7 algorithm. The input parameter
        trmode specifies the resonance type (normally "reflection_method1"
        or "reflection_method2" - see MAT 58 Section 2.5.1).

        Uncalibrated line should be de-embedded (if it has significant
        length) from the S-parameter data before calling the functions
        in this module. The remaining length (phase) is fitted by optimisefit7.

        The user must supply string of characters 'loop_plan',
        which defines order of steps used by the fitting process.


      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/>.


      System
      ------

        Compatible with Python 2.7 and 3.X
        Requires numpy 1.4 or later.


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

        :30-October-2021:  Released under CC0.
        :18-February-2024: Added function estimate_length_uncal_line(). This can be used
                           to prefit length of uncalibrated line when a rough estimate is
                           insufficient (often applies when Q < 100 and cable-length is multiple
                           wavelengths).
                           Removed unused variable.
                           Cleaner version of initialfit().
        :09-August-2024:   Added function GetQoNotchPivotedQCircle(). Calculates Qo for planar-circuit
                           notch resonators when the Q-circle is pivoted as a result of asymmetry.
                           Fix for numpy deprecation warning.
                           New option 'r' in loop_plan to drop the 10% of points with largest residual error.
                           
"""

import numpy as np
from math import sqrt, pi, log, ceil
from cmath import exp as cexp


def angularwts(F, Fres, QL):
    """
        Calculate diagonal elements of weights matrix.
        ----------------------------------------------

          The weights are needed when F consists of equally-spaced
          frequency points (rather than points equally spaced around
          the Q-circle). MAT 58 eqn. (28).
    """

    F2 = np.array(F)
    ptmp = 2.0*QL*(F2-Fres)/Fres
    PV = 1.0/(ptmp**2 + 1.0);
    return PV


def initialfit(F, S, N, Fres, Qseed):
    """
        Initial linear fit, step (1)
        ----------------------------

          NLQFIT6, NLQFIT7 and NLQFIT8 Step (1).

          Code amended to make more efficient, Feb. 2024.
          Gives the same results as an earlier version given
          in MAT 58 issue 1 and 2.

          Input parameters:
            F     - List of frequencies
            S     - List of complex data S-parameter to be fitted.
            N     - Number of points
            Fres  - Estimated resonant frequency (not fitted)
            Qseed - Estimated QL (will be improved by fitting).
                    It is often sufficient to supply QL = 1000.0

          Output parameters (MAT58 eqn. 17):
             sv = [a', a'', b', b'', QL]
    """

    N2 = 2*N;
    M = np.zeros([N2, 5])
    G = np.zeros(N2)
    P = np.zeros(N2)

    for i in range(N):
        i2 = i+N
        t = 2.0*(F[i]/Fres - 1.0)
        y = abs(1.0/complex(1.0, Qseed*t))
        P[i] = P[i2] = y*y
        v1 = S[i]
        G[i] = v1.real;    G[i2] = v1.imag
        v2 = v1*t
        M[i,:] =  t, 0.0,  1.0, 0.0, v2.imag
        M[i2,:] = 0.0, t,  0.0, 1.0, -v2.real

    X = M.transpose()
    T = np.multiply(X,P)
    C = np.matmul(T,M)
    q = np.matmul(T,G)
    sv = np.linalg.solve(C,q)
    return sv


def estimate_length_uncal_line(F, SA, N, Fres, Qseed, lowerbound, upperbound, clight, checkbounds=False):
    """
        Estimate cable length from initial fits by using golden-section search
        ----------------------------------------------------------------------

        Robust fitting of the length from reflection-coefficient data. The value obtained
        can be used to improve robustness of optimised fits with optimisefit7().

          Input parameters:
            F     - List of frequencies.
            SA    - List of complex data S-parameter (reflection only) to be fitted.
            N     - Number of points
            Fres  - Estimated resonant frequency (not fitted)
            Qseed - Estimated QL (will be improved by fitting).
            lowerbound, upperbound - search range for cable length.
            clight - velocity of light in the line (units consist with lengths and Fres)
            checkbounds - optional check.

          Output parameters (MAT58 eqn. 17):
             len_ul - estimated length

          Source: Wikipedia, December 2023 (edited).
             https://en.wikipedia.org/wiki/Golden-section_search

    """

    invphi = (sqrt(5) - 1) / 2  # 1 / phi
    invphi2 = (3 - sqrt(5)) / 2  # 1 / phi^2

    def GetErrSum(F, S, N, Fres, Qseed, len_ul):
        # For S11 or S22 only
        D = [S[n]*cexp(complex(0.0, 4*pi*F[n]*len_ul/clight)) for n in range(N)]
        sv = initialfit(F, D, N, Fres, Qseed)
        errsum = 0.0
        for i in range(N):
            t = 2.0*(F[i]/Fres - 1.0)
            r = D[i]-(complex(sv[0],sv[1])*t + complex(sv[2],sv[3]))/complex(1.0, sv[4]*t)
            errsum+=r.real**2 + r.imag**2
        return errsum

    # Golden-section search
    tol = 0.1*clight/Fres  # one tenth of a wavelength
    a = min(lowerbound, upperbound)
    b = max(lowerbound, upperbound)
    h = b - a
    assert h>2.0*tol, 'Improper bounds for length fit'

    c = a + invphi2 * h
    d = a + invphi * h
    yc = GetErrSum(F, SA, N, Fres, Qseed, c)
    yd = GetErrSum(F, SA, N, Fres, Qseed, d)

    if checkbounds:
        assert GetErrSum(F, SA, N, Fres, Qseed, c-tol)>yc, 'Length-fit unsuccessful for specified bounds'
        assert GetErrSum(F, SA, N, Fres, Qseed, d-tol)>yd, 'Length-fit unsuccessful for specified bounds*'

    n = int(ceil(log(tol / h) / log(invphi))) # n must be 2 or above because h>2.0*tol
    nn = max(n-1,2)
    for k in range(nn):
        if yc < yd:
            b = d
            d = c
            yd = yc
            h = invphi * h
            c = a + invphi2 * h
            yc = GetErrSum(F, SA, N, Fres, Qseed, c)
        else:
            a = c
            c = d
            yc = yd
            h = invphi * h
            d = a + invphi * h
            yd = GetErrSum(F, SA, N, Fres, Qseed, d)

    if yc < yd:
      len_ul = 0.5*(a+d)
    else:
      len_ul = 0.5*(b+c)

    return len_ul


def optimisefit7(F, S, N, Fseed, sv, loop_plan, Tol, quiet):
    """
       Iterative non-linear fit, NLQFIT7 Step (2)
       ------------------------------------------

         Optimised fit of Q-factor (QL) and resonant frequency (FL)
         by the gradient-descent method.

         Uses the results of the initial fit (sv) as the starting
         values for the iteration.

         Input parameters:
            F         - List of frequencies.
            S         - List of complex data S-parameter to be fitted.
            N         - Number of points.
            Fseed     - Estimated resonant frequency.
            sv        - Initial solution (numpy vector or a list) found with InitialFit
            loop_plan - String of characters which defines order of
                        steps used by the fitting process, e.g. 'fwfwc':
                         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.
                         r - Identify 10% largest outliers from previous fit. Remove the effect of these
                             by giving them zero weighting.
            Tol       - Criterion for the convergence test.
                        Recommend using 1.0E-5 for reflection
                        or max(abs(Gamma))*1.0E-5 for transmission.
            quiet     - Boolean flag controlling output of information to
                        the console.

         Output (MAT58 eqn. 26):
               list of fitted parameters: [m1, m2, m4, m4, QL, FL, m7/Flwst]

    """

    if loop_plan[-1] == 'w' or loop_plan[-1] == 'r':
        assert 0, 'Last item in loop_plan must not be w (weight calculation) or r (remove outliers)'
    if loop_plan[0] == 'w' or loop_plan[0] == 'r':
        assert 0, 'First item in loop_plan must not be w (weight calculation) or r (remove outliers)'
    if loop_plan[-1] != 'c': print('Warning: Last item in loop_plan is not c so convergence not tested!')

    N2 = N*2;
    PV = np.ones(N)     # default weights vector
    PV2 = np.ones(N2)

    m1 = sv[1]/sv[4]  # a''/QL
    m2 = -sv[0]/sv[4]
    m3 = sv[2]-m1
    m4 = sv[3]-m2
    m5 = sv[4]
    Flwst = F[0]   # lowest freq. is a convenient normalisation factor.
    m6 = Flwst*m5/Fseed
    m7 = 0.0
    last_op = 'n'
    del sv
    weighting_ratio = None
    number_iterations = 0
    OL2 = np.ones(2*N, dtype=int) # Outliers indicated by zeros.

    ## Loop through all of the operations specified in loop_plan
    for op in loop_plan:

        if op == 'w':          # Fr                       QL
            PV = angularwts(F, Flwst*float(m5)/float(m6), float(m5))
            weighting_ratio = max(PV)/min(PV)
            PV2 = np.concatenate((PV, PV))
            if not quiet: print('Op w, Calculate weights')
            last_op = 'n'
            continue

    ## Find outliers. The 10% of frequencies with largest residual errors are excluded.
    ## Currently, the residuals are not weighted.
    ## Added August 2024
        if op == 'r': 
            ResidsArray = np.zeros(N)
            if not quiet: print('Op r, Remove effect of largest outliers')
            for i in range(N):
                fdn = F[i]/Flwst - m5/m6
                den = complex(1.0, 2*(m6*F[i]/Flwst-m5))
                pj = complex(0.0,m7*fdn)
                E = S[i] - (complex(m1,m2) + complex(m3,m4)/den)*cexp(pj)
                ResidsArray[i] = E.real**2 + E.imag**2
            ZeroWeightList = ResidsArray.argsort()[::-1][:round(N/10)]
            OL = np.ones(N, dtype=int)
            for i in ZeroWeightList: OL[i] = 0
            OL2 = np.concatenate((OL, OL))
            del OL, ResidsArray, ZeroWeightList
            continue

        if op == 'c': seek_convergence = True
        elif op == 'f': seek_convergence = False
        else: assert 0,'Unexpected character in loop_plan'

        TerminationConditionMet = False
        RMS_Error = None
        while not (TerminationConditionMet):
            number_iterations += 1
            M = np.zeros([N2,7])
            G = np.zeros(N2)
            c1 = complex(-m4,m3); c2 = complex(m1,m2); c3 = complex(m3,m4)
            for i in range(N):
                i2 = i+N
                y = 1.0/complex(1.0,2*(m6*F[i]/Flwst-m5))
                fdn = F[i]/Flwst - m5/m6
                pj = complex(0.0, m7*fdn)
                expm7 = cexp(pj)
                ym = y*expm7
                u = c1*y*ym*2
                u2 = -u*F[i]/Flwst
                v = (c2 + y*c3)* expm7;
                u3 = v*fdn;
                M[i,:] = expm7.real, -expm7.imag, ym.real, -ym.imag, u.real, u2.real, -u3.imag
                M[i2,:] = expm7.imag, expm7.real, ym.imag, ym.real, u.imag, u2.imag, u3.real
                r = S[i] - v  # residual
                G[i] = r.real; G[i2] = r.imag

            X = M.transpose()
            PV3 = PV2*OL2 # Exclude outliers if these have been defined.
            T = np.multiply(X, PV3)
            C = np.matmul(T,M)
            q = np.matmul(T,G)
            dm = np.linalg.solve(C,q)
            m1 += dm[0]; m2 += dm[1]; m3 += dm[2]
            m4 += dm[3]; m5 += dm[4]; m6 += dm[5]
            m7 += dm[6]
            del G, X, T, C, dm
            if RMS_Error is not None: Last_RMS_Error = RMS_Error
            else: Last_RMS_Error = None

            SumNum = 0.0
            SumDen = 0.0
            for i in range(N):
                if OL2[i]: # Outliers not included in RMS_Error
                    fdn = F[i]/Flwst - m5/m6
                    den = complex(1.0, 2*(m6*F[i]/Flwst-m5))
                    pj = complex(0.0,m7*fdn)
                    E = S[i] - (c2 + c3/den)*cexp(pj)
                    ip = PV[i]
                    SumNum = SumNum + ip*(E.real*E.real + E.imag*E.imag)
                    SumDen = SumDen + ip

            RMS_Error = sqrt(SumNum/SumDen);
            if not quiet:
                if last_op =='c':
                    print('      Iteration %i, RMS_Error %10.8f' % (number_iterations, RMS_Error))
                else:
                    print('Op %c, Iteration %i, RMS_Error %10.8f' % (op, number_iterations, RMS_Error))
            last_op = op

            if seek_convergence:
                if Last_RMS_Error is not None:
                    delta_S = abs(RMS_Error-Last_RMS_Error)
                    TerminationConditionMet = delta_S < Tol
            else:
                TerminationConditionMet = True

        # After last operation, we end up here ...
        if not quiet: print()
    return [m1,m2,m3,m4,m5,m5*Flwst/m6, m7/Flwst], weighting_ratio, number_iterations, RMS_Error


def QCircleData(A, m1, m2, m3, m4):
    """
      Use MAT 58 eqn. (31) to calculate calibrated diameter
      -----------------------------------------------------

      Also return calibrated gamma_V and gamma_T
    """
    aqratio = complex(m1,m2)
    b = complex(m1+m3, m2+m4)
    caldiam = abs(b - aqratio)*A
    calgamma_V =  complex(m1,m2)*A
    calgamma_T = b*A
    return caldiam, calgamma_V, calgamma_T


def GetUnloadedData(mv, scaling_factor_A, trmode, quiet):
    """
       Calculate unloaded Q-factor and various 'calibrated' quantities
       ---------------------------------------------------------------

       Input parameters:
         mv               - solution produced by OptimiseFit
         trmode           - 'transmission', 'reflection_method1',
                            'reflection_method2' or 'absorption'
         scaling_factor_A - scaling factor as defined in MAT 58.
                            For reflection_method1, can specify as 'AUTO'
                            to use the magnitude of the fitted detuned reflection coefficient (gammaV)
    """

    if type(scaling_factor_A) is str:
        if (scaling_factor_A.upper()=='AUTO'): auto_flag = True
        else: return 1, 'Illegal Scaling factor'
    else:
        # Also permit negative scaling factor to indicate 'Auto'
        if scaling_factor_A<0.0: auto_flag = True
        else: auto_flag = False

    m1,m2,m3,m4,m5,FL,m7_flwst = mv

    if trmode=='transmission':
        if auto_flag: return 1, 'Scaling factor must not be "Auto" for transmission case'
        cal_diam, cal_gamma_V, cal_gamma_T = QCircleData(scaling_factor_A, m1, m2, m3, m4)
        if cal_diam == 1.0:
            return 1, 'Divide by zero forestalled in calculation of Qo'
        Q0 = m5/(1.0-cal_diam)

    elif trmode=='reflection_method1':
        if auto_flag:
            if not quiet: print('Supplied scaling_factor_A is "Auto", so using fitted data to estimate it')
            scaling_factor_A = 1.0/abs(complex(m1,m2)) # scale to gammaV if 'AUTO'
        cal_diam, cal_gamma_V, cal_gamma_T = QCircleData(scaling_factor_A, m1, m2, m3, m4)
        cal_touching_circle_diam = 2.0
        if not quiet: print('  Q-circle diam = %5.3f, touching_circle_diam = %5.3f' % (cal_diam, cal_touching_circle_diam))
        den = cal_touching_circle_diam/cal_diam - 1.0
        Q0 = m5*(1.0 + 1.0/den)

    elif trmode=='reflection_method2':
        if auto_flag: return 1, 'Scaling factor must not be "Auto" for Method 2'
        cal_diam, cal_gamma_V, cal_gamma_T = QCircleData(scaling_factor_A, m1, m2, m3, m4)
        gv = abs(cal_gamma_V)
        gv2 = gv*gv
        mb = abs(cal_gamma_T)
        cosphi = (gv2 + cal_diam*cal_diam - mb*mb)/(2.0*gv*cal_diam) # Cosine rule
        cal_touching_circle_diam = ((1.0-gv2)/(1.0-gv*cosphi))
        if not quiet: print('  Q-circle diam = %5.3f, touching_circle_diam = %5.3f' % (cal_diam, cal_touching_circle_diam))
        den = cal_touching_circle_diam/cal_diam - 1.0
        Q0 = m5*(1.0 + 1.0/den)

    elif trmode=='notch' or trmode=='absorption': # Absorption resonator measured by transmission
        if auto_flag:
            if not quiet: print('Notch/absorption Qo calculation: Supplied scaling_factor_A is "Auto", so using fitted data to estimate it')
            scaling_factor_A = 1.0/abs(complex(m1,m2)) # scale to gammaV if 'AUTO'
        cal_diam, cal_gamma_V, cal_gamma_T = QCircleData(scaling_factor_A, m1, m2, m3, m4)
        if not quiet: print('  Q-circle diam = %5.3f' % (cal_diam))
        if cal_diam == 1.0:
            return 1, 'Divide by zero forestalled in calculation of Qo'
        den = 1.0/cal_diam - 1.0 # Gao thesis (2008) 4.35 and 4.40
        Q0 = m5*(1.0 + 1.0/den)  # https://resolver.caltech.edu/CaltechETD:etd-06092008-235549
        # For this type of resonator, critical coupling occurs for cal_diam = 0.5.

    else:
        return 1, 'Unknown trmode'

    return 0, (Q0, cal_diam, cal_gamma_V, cal_gamma_T)


def GetQoNotchPivotedQCircle(mv, quiet):
    """
       Calculate unloaded Q-factor and Touching-Circle diameter for notch resonator with pivoted Q-circle
       ---------------------------------------------------------------------------------------------------

       Use the touching-circle method to evaluate Qo and correct the Q-circle parameters.
       This is an alternative to the Diameter-Corrected-Method (DCM) developed by Khalil (https://doi.org/10.1063/1.3692073).
       It is used only for planar-circuit notch resonators when asymmetry results from pivoting of the Q-circle around Sv.
       The scaling factor is determined automatically, then normalised to the touching-circle radius.

       On the basis of limited tests so far, good agreement with DCM is obtained.

       Input parameters:
         mv               - solution produced by OptimiseFit

    """
    m1,m2,m3,m4,m5,FL,m7_flwst = mv
    if not quiet: print('Notch/absorption Qo calculation with pivoted Q-circle')
    scaling_factor_A = 1.0/abs(complex(m1,m2))

    cal_diam, cal_gamma_V, cal_gamma_T = QCircleData(scaling_factor_A, m1, m2, m3, m4)

    # Find centre of Q-circle, add radius to get diameter of touching circle.
    rad_touching_circle = abs(cal_gamma_V+cal_gamma_T)/2 + cal_diam/2
    if not quiet: print('  Radius of touching circle = %5.3f' % (rad_touching_circle))

    corr_diam = cal_diam*abs(cal_gamma_V)/rad_touching_circle
    corr_gamma_V = cal_gamma_V/rad_touching_circle
    corr_gamma_T = cal_gamma_T/rad_touching_circle

    if not quiet: print('  Corrected Q-circle diam = %5.3f' % (corr_diam))
    if corr_diam == 1.0:
        return 1, 'Divide by zero forestalled in calculation of Qo'

    corr_Qo = float(m5)*(1.0 + 1.0/(1.0/corr_diam - 1.0))

    return 0, (corr_Qo, rad_touching_circle, corr_diam, corr_gamma_V, corr_gamma_T)
