#!/usr/bin/env python

# Copyright (c) 2024
#      Daniel O'Connor <darius@dons.net.au>.  All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY AUTHOR AND CONTRIBUTORS ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL AUTHOR OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.
#

# This script uses the "Y Factor Method" as discussed here:
# https://scdn.rohde-schwarz.com/ur/pws/dl_downloads/dl_application/application_notes/1ma178/1MA178_5e_NoiseFigure.pdf
# https://www.analog.com/en/resources/technical-articles/noise-figure-measurement-methods-and-formulas--maxim-integrated.html
#
# It requires an ENR and has hard coded values for the source we have.

import argparse
import math
import misc
import numpy
import rsib
import scipy
import scpi
import sys

def findenr(frq):
    # ENR values from the noise source
    enrdb = numpy.array([15.55, 15.96, 15.68, 15.11, 15.07, 14.84, 14.77, 14.82, 14.86, 14.79, 14.83, 14.93, 14.93, 15.07, 15.19, 15.08, 15.14, 14.87, 14.97, 14.59])
    enrfrq = numpy.array([0.01e9, 0.1e9, 1.0e9, 2.0e9, 3.0e9, 4.0e9, 5.0e9, 6.0e9, 7.0e9, 8.0e9, 9.0e9, 10.0e9, 11.0e9, 12.0e9, 13.0e9, 14.0e9, 15.0e9, 16.0e9, 17.0e9, 18.0e9])

    # Convert back to linear values
    enr = 10 ** (enrdb / 10)

    # Interpolate
    rtn = numpy.interp([frq], enrfrq, enr)

    # Convert to dB
    rtndb = 10 * math.log10(rtn[0])

    return rtndb

def setup(r, freq, span, sweeps, bw, atten, time):
    # Reset to defaults
    r.write('*RST')

    # Set to single sweep mode
    r.write('INIT:CONT OFF')

    # Enable display updates
    r.write('SYST:DISP:UPD ON')

    # Set frequency range
    r.write('SENSE1:FREQ:CENT %f Hz' % (freq))
    r.write('SENSE1:FREQ:SPAN %f Hz' % (span))

    # Switch marker 1 on in screen A
    r.write('CALC:MARK1 ON')

    # Set marker to centre frequency
    r.write('CALC:MARK1:X %f Hz ' % (freq))

    # Enable noise measurement
    r.write('CALC:MARK1:FUNC:NOIS ON')

    # Turn averaging on
    r.write('AVER:STAT ON')

    # Set number of sweeps
    r.write('SWE:COUN %d' % (sweeps))

    # Use RMS detector
    r.write('SENS:DET RMS')

    # Set sweep time
    r.write('SWE:TIME %f' % (time))

    # Set resolution bandwidth
    r.write('SENS1:BAND:RES %f Hz' % (bw))

    # Set video bandwidth (10x res BW)
    r.write('SENS1:BAND:VID %f Hz' % (bw * 10))

    r.write('INP:ATT %s' % atten)

def getnoise(r):
    # Trigger the sweep
    r.write('INIT;*WAI')

    # Wait for it to be done
    r.ask('*OPC?')

    # Set data format
    r.write('FORM:DATA ASC')

    # Read noise value
    data = r.ask('CALC:MARK1:FUNC:NOIS:RES?')
    #print 'Data - ' + data

    return float(data)

def setnoise(r, en):
    if en:
        val = 'ON'
    else:
        val = 'OFF'
    r.write('DIAG:SERV:NSO ' + val)

def calcnf(enrdb, offdb, ondb):
    # Not possible but noisy results may result in it happening
    if ondb <= offdb:
        return None
    ydb = ondb - offdb
    y = 10 ** (ydb / 10)
    enr = 10 ** (enrdb / 10)
    nf = 10 * math.log10(enr / (y - 1))
    return nf

def donoisetest(r, enr, sweeps):
    swt = float(r.ask('SWE:TIME?'))
    print('Acquiring with noise off.. (%.1f x %d = %.1f sec)' % (swt, sweeps, swt * sweeps))
    setnoise(r, False)
    off = getnoise(r)
    print('Acquiring with noise on.. (%.1f x %d = %.1f sec)' % (swt, sweeps, swt * sweeps))
    setnoise(r, True)
    on = getnoise(r)
    nf = calcnf(enr, off, on)
    #print(nf)
    return off, on, nf

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description = 'Configures a Rohde Schwarz FSP7 spectrum analyser to do a noise figure test',
                                     epilog = 'Note: video bandwidth is set to 10 times the resolution bandwidth')
    parser.add_argument('-a', '--atten', default = 10, help = 'Input attenuation in dB (default: %(default).0f dB)', type = float)
    parser.add_argument('-b', '--bw', default = 1000, help = 'Resolution bandwidth in Hz (default: %(default).1f Hz)', type = float)
    parser.add_argument('-i', '--input', default = None, help = 'Frequency used to compute ENR (defaults to frequency)', type = float)
    parser.add_argument('-p', '--pause', default = False, action = 'store_true',
                        help = 'Wait between measurements (when not doing N repeats, default: %(default)s)')
    parser.add_argument('-r', '--repeat', help = 'Number of repetitions, if not specified do one and ask to continue', type = int)
    parser.add_argument('-s', '--span', default = 1e6, help = 'Span frequency in Hz (default: %(default).0f Hz)', type = float)
    parser.add_argument('-t', '--time', default = 30, help = 'Sweep time (default: %(default)f sec)', type = float)
    parser.add_argument('-w', '--sweeps', default = 3, help = 'Number of sweeps to average (default: %(default)d)', type = int)
    parser.add_argument('address', help = 'Spectrum analyser address', type = str)
    parser.add_argument('centre', help = 'Centre frequency (Hz)', type = float)

    args = parser.parse_args()

    if args.input == None:
        args.input = args.centre

    if args.time is not None and args.time <= 0:
        parser.error('Sweep time must be >0')

    # Compute ENR at frequency of interest
    enr = findenr(args.input)

    # Connect to the analyser
    r = rsib.RSIBDevice(args.address)

    # ID instrument
    print('Device ID is ' + r.ask('*IDN?').decode('ascii'))

    # Setup parameters
    setup(r, args.centre, args.span, args.sweeps, args.bw, args.atten, args.time)

    nfs = []
    print('Centre: %s, Span %s, Input %s, RBW %s, Sweeps: %d, Sweep time: %.1f sec, ENR %.2f dB' % (
        misc.sifmt(args.centre), misc.sifmt(args.span),
        misc.sifmt(args.input), misc.sifmt(args.bw),
        args.sweeps, args.time, enr))
    while args.repeat == None or args.repeat > 0:
        try:
            off, on, nf = donoisetest(r, enr, args.sweeps)
            if nf is None or nf < 0:
                nfstr = 'Invalid'
            else:
                nfstr = '%.2f dB' % (nf)
                nfs.append(nf)
            print('Off %.3f dBm/Hz, on %.3f dBm/Hz, NF %s' % (off, on, nfstr))
            if args.repeat == None:
                if args.pause:
                    print('Press enter to perform a new measurement')
                    sys.stdin.readline()
            else:
                args.repeat -= 1
        except KeyboardInterrupt:
            print('')
            break

    if len(nfs) > 1:
        nfs = numpy.array(nfs)

        # XXX: not sure mean/std of dBm/Hz values is really meaningful
        print('N: %d, NF min: %.1f dBm/Hz, max: %.1f dBm/Hz, avg: %.1f dBm/Hz, stddev: %.1f' % (
            len(nfs), nfs.min(), nfs.max(), nfs.mean(), numpy.std(nfs)))