"""
    qfitmod8.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 NLQFIT8 algorithm. The input parameter
        trmode specifies the resonance type (normally "transmission").

        This version fits real and imaginary parameters that describe
        a frequency dependence of the background (leakage vector). This
        is typically used when a transmission resonance sits on the
        'tail' of another resonance.

        Uncalibrated line should be de-embedded (if it has significant
        length) from the S-parameter data before calling the functions
        in this module.

        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: Removed unused variable.
                           Cleaner version of initialfit().
        :09-August-2024:   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
import math, cmath


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 optimisefit8(F, S, N, Fseed, sv, loop_plan, Tol, quiet):
    """
       Iterative non-linear fit, NLQFIT8 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. 43)
               (list of fitted parameters): [m1, m2, m4, m4, m8, m9, QL, FL]
    """

    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
    m8 = 0.0
    m9 = 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):
                den = complex(1.0,2*(m6*F[i]/Flwst-m5))
                E = S[i] - complex(m1,m2) - complex(m8,m9)*t - complex(m3,m4)/den
                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,8])
            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))
                u = c1*y*y*2
                u2 = -u*F[i]/Flwst
                FL = Flwst*m5/m6;
                t = 2*(F[i]-FL)/FL;
                M[i,:] = 1.0, 0.0, y.real, -y.imag, u.real, u2.real, t, 0.0
                M[i2,:] = 0.0, 1.0, y.imag, y.real, u.imag, u2.imag, 0.0, t
                v = c2 + c3*y + (m8+1j*m9)*t;
                r = S[i] - v  # residual
                G[i] = r.real; G[i2] = r.imag

            X = M.transpose()
            T = np.multiply(X, PV2)
            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]
            m8 += dm[6]; m9 += dm[7]
            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
                    den = complex(1.0,2*(m6*F[i]/Flwst-m5));
                    FL = Flwst*m5/m6;
                    t = 2*(F[i]-FL)/FL;
                    ip = PV[i]
                    E = S[i] - complex(m1,m2) - complex(m8,m9)*t - complex(m3,m4)/den
                    SumNum = SumNum + ip*(E.real*E.real + E.imag*E.imag)
                    SumDen = SumDen + ip

            RMS_Error = math.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,m8,m9,m5,m5*Flwst/m6], 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,m8,m9,m5,FL = 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': # 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)
