# -*- coding: utf-8 -*-
"""
Created July 04, 2026 with support from ChatGPT

@author: Christian Monstein, HB9SCT
"""
#---------------------------------------------------------------------------

import time
from datetime import datetime
import serial
import numpy as np
from matplotlib import pyplot as plt
from scipy.interpolate import interp1d
import subprocess
import sys
import struct
import SoapySDR
from SoapySDR import *


# ==========================================================
# Parameters
# ==========================================================

MyTitle = 'NS16-BZX384' # Name of the source to be tested(DUT)

DUTvoltage = 15 # Be careful for self produced noise sources!
REFvoltage = 28 # Supply voltage of the calibration (reference) source

# Calibration sourcew AILTECH 7618E 10 MHz - 18 GHz
x = [10e6,  30e6,  100e6, 300e6, 600e6, 1000e6, 2000e6, 3000e6, 4000e6] # Frequency table noise source
y = [14.36, 14.84, 15.07, 15.07, 15.05, 15.01,  15.06,  14.91,  14.92]  # ENR dB
intfunc = interp1d(x,y)

# Frequency table which we are interested in
FREQ_LIST =   [11e6, 15e6, 20e6, 30e6, 40e6, 50e6, 70e6, 100e6, 150e6, 200e6, 300e6, 400e6, 500e6, 600e6, 700e6, 800e6, 900e6] # Hz

ENRcal = intfunc(FREQ_LIST)

CALcold = []
CALhot  = []
DUTcold = []
DUThot  = []

SAMPLE_RATE = 20e6

FFT_BINS = 32
READ_SIZE = 4096
N_AVG = 400

GAIN_LNA = 15
GAIN_VGA = 30
GAIN_AMP = 14

# ==========================================================
# HackRF Initialisation
# ==========================================================

print("Opening HackRF")

devices = SoapySDR.Device.enumerate()

if len(devices) == 0:
    print("No HackRF found.")
    sys.exit()

print(dict(devices[0]))

sdr = SoapySDR.Device(devices[0])

print("HackRF opened successfully")

sdr.setSampleRate(
    SOAPY_SDR_RX,
    0,
    SAMPLE_RATE
)

sdr.setGain(
    SOAPY_SDR_RX,
    0,
    "LNA",
    GAIN_LNA
)

sdr.setGain(
    SOAPY_SDR_RX,
    0,
    "VGA",
    GAIN_VGA
)

sdr.setGain(
    SOAPY_SDR_RX,
    0,
    "AMP",
    GAIN_AMP
)

rxStream = sdr.setupStream(
    SOAPY_SDR_RX,
    SOAPY_SDR_CF32
)

sdr.activateStream(
    rxStream
)

print("HackRF ready")
print()

# ==========================================================
# Relay control
# ==========================================================

def relay(channel, state):
    name = f"A0001_{channel}"
    subprocess.run(
        ["usbrelay", f"{name}={state}"],
        capture_output=True,
        text=True,
        check=True
    )
    time.sleep(0.5)

def SwitchDUT():
    relay(2, 1)
    time.sleep(0.5)
    relay(1, 1)

def SwitchREF():
    relay(1, 0)
    time.sleep(0.5)
    relay(2, 0)

# ==========================================================
# Control laboratory power supply JT-PSM01
# ==========================================================


def set_psm01_voltage(voltage):
    """
    Set voltage of module Joy-IT JT-PSM01

    Range: 0 ... 30 V
    """

    PORT = "/dev/ttyACM0"
    BAUD = 9600

    # Spannung prüfen
    if not 0.0 <= voltage <= 30.0:
        raise ValueError("Voltage must be in range 0 and 30 V.")

    # Checksum
    def checksum(data):
        return sum(data) & 0xFF

    # Open COM-Port
    with serial.Serial(
        PORT,
        BAUD,
        bytesize=8,
        parity="N",
        stopbits=1,
        timeout=1
    ) as ser:

        # ---------------------------------------------
        # power supply ONLINE
        # ---------------------------------------------

        packet = bytes.fromhex(
            "F1 C1 00 01 01 02"
        )

        ser.write(packet)
        time.sleep(0.2)

        # ---------------------------------------------
        # Set voltage
        # ---------------------------------------------

        value = struct.pack("<f", voltage)

        packet = bytearray([
            0xF1,
            0xB1,
            0xC1,
            0x04
        ])

        packet.extend(value)
        packet.append(checksum(packet[2:]))

        print("SET:", packet.hex(" "))

        ser.write(packet)
        time.sleep(0.2)

        # ---------------------------------------------
        # Activate output
        # ---------------------------------------------

        packet = bytes.fromhex(
            "F1 B1 DB 01 01 DD"
        )

        print("OUTPUT ON:", packet.hex(" "))

        ser.write(packet)
        time.sleep(0.2)

        print(f"JT-PSM01: {voltage:.3f} V rated.")

        # -------------------------------------------------
        # Search for temperatur-package
        # -------------------------------------------------
        packet = bytes.fromhex("F1 A1 C4 01 00 C5")
    
        ser.write(packet)
        time.sleep(0.2)
    
        # Read response
        data = ser.read(200)
        
        pattern = bytes.fromhex("F0 A1 C4 04")
    
        pos = data.find(pattern)
    
        if pos < 0:
            raise RuntimeError(
                "Did not get temperature from JT-PSM01."
            )
    
        # 4 Byte Float, Little Endian
        temperature = struct.unpack(
            "<f",
            data[pos + 4:pos + 8]
        )[0]
    
        print(
            f"JT-PSM01: {voltage:.3f} V, "
            f"Temperature: {temperature:.2f} °C"
        )
        return(temperature)

# ==========================================================
# Measure one single frequency
# ==========================================================

def measure(fc):

    IF = fc

    sdr.setFrequency(
        SOAPY_SDR_RX,
        0,
        IF
    )

    time.sleep(0.05)

    avg_spec = np.zeros(
        FFT_BINS,
        dtype=np.float64
    )

    n_avg = 0

    # Flush one receive buffer
    buff = np.empty(
        READ_SIZE,
        dtype=np.complex64
    )

    try:
        sdr.readStream(
            rxStream,
            [buff],
            READ_SIZE
        )
    except Exception:
        pass

    for _ in range(N_AVG):

        buff = np.empty(
            READ_SIZE,
            dtype=np.complex64
        )

        sr = sdr.readStream(
            rxStream,
            [buff],
            READ_SIZE
        )

        if sr.ret <= 0:
            continue

        iq = buff[:sr.ret]

        fft = np.fft.fftshift(
            np.fft.fft(iq)
        )

        # nur positive Frequenzen
        fft = fft[len(fft)//2:]

        power = np.abs(fft)**2

        power = power.reshape(
            FFT_BINS,
            -1
        ).mean(axis=1)

        power = 10*np.log10(
            power + 1e-12
        )

        avg_spec += power

        n_avg += 1

    if n_avg == 0:
        return np.nan

    avg_spec /= n_avg

    return np.mean(avg_spec)

# ==========================================================
# Main program
# ==========================================================

def Sweep(): 
    temp = []
    p = measure(FREQ_LIST[0]) # forget
    for f in FREQ_LIST:
        p = measure(f)
        print ('Frequency: {:5.1f} MHz  Power: {:5.1f} dB'.format(f/1e6,p))
        temp.append(p)
        time.sleep(0.01)
    return (temp)
            
#---------------------------------------------------------------------------

filename = 'Reference-Cold.prn'
print (filename)
SwitchREF()
set_psm01_voltage(0)
time.sleep(1)
CALcold = Sweep() # CAL cold
print()

#---------------------------------------------------------------------------

filename = 'Reference-Hot.prn'
print (filename)
SwitchREF()
set_psm01_voltage(REFvoltage)
time.sleep(1)
CALhot = Sweep() # CAL hot
set_psm01_voltage(0)
time.sleep(1)
print()

#---------------------------------------------------------------------------

filename = MyTitle+'-Cold.prn'
print (filename)
SwitchDUT()
time.sleep(5)
DUTcold = Sweep() # DUT cold
print()

#---------------------------------------------------------------------------

filename = MyTitle+'-Hot.prn'
print (filename)
SwitchDUT()
set_psm01_voltage(DUTvoltage)
time.sleep(1)
DUThot = Sweep() # DUT hot
print()

#---------------------------------------------------------------------------

temperature = set_psm01_voltage(0)
SwitchREF()
print("Swich off power and select REF")
print("Closing HackRF")

sdr.deactivateStream(rxStream)
sdr.closeStream(rxStream)

print("Finished")

#---------------------------------------------------------------------------

Ycal = np.array(CALhot) - np.array(CALcold)
Ydut = np.array(DUThot) - np.array(DUTcold)
ENRdut = ENRcal - Ycal + Ydut
print ('ENR DUT: ')
print (ENRdut)

freq_mhz = [f / 1e6 for f in FREQ_LIST]
#---------------------------------------------------------------------------

timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
prn_filename = f"{MyTitle}_{timestamp}.prn"

with open(prn_filename, "w") as f:
    #f.write("Frequency [MHz];ENR [dB]\n")
    f.write(f"Frequency [MHz];ENR [dB]; Temperature: {temperature:.1f} °C\n")
    for freq, enr in zip(freq_mhz, ENRdut):
        f.write(f"{freq:.3f};{enr:.2f}\n")
    #f.write(f"Temperature: {temperature:.1f} °C")

print(f"Results written to {filename}")

#---------------------------------------------------------------------------
            
plt.figure()
plt.plot(freq_mhz,ENRdut,'-*b',label='DUT ('+MyTitle+')')
plt.plot(freq_mhz,ENRcal,'-+r',linewidth=2,label='Reference (Ailtech 7618E)')
plt.title('Measured ENR of Noise Source '+MyTitle)
plt.xlabel('Frequency [MHz]')
plt.ylabel('Excess Noise ENR [dB]')
plt.legend(loc='upper right')
plt.ylim(0,25)
plt.xscale('log')
plt.grid(True,which='both',axis='both')
png_filename = f"{MyTitle}_{timestamp}.png"
plt.savefig(png_filename,bbox_inches="tight")
plt.show()

#---------------------------------------------------------------------------