# Python 3 code for control of an Arduino driven Antenna Turntable with automated VNA measurent
# Written by Nicholas Caudhil, Applications Engineering, Copper Mountain Technologies
# 317-222-5400 support@coppermountaintech.com
# Used by Brian Walker, SME for Antenna Webinar 3/22/23

import serial
import time
import pyvisa as visa
import numpy as np
import matplotlib.pyplot as plt

# The Arduino might mount on another COM port
# Use the Device Manager to see where it is
COMPORT = 'COM6'
NOP = 1000

arduino = serial.Serial(port=COMPORT, baudrate=115200)
time.sleep(3)

# Connect to the S2VNA program running on the local machine at port 5025
# Socket must be enabled in the S2 program under Misc/Settings
def setup_vna():
    global CMT
    rm = visa.ResourceManager()
    try:
        CMT = rm.open_resource('TCPIP0::127.0.0.1::5025::SOCKET')
    except:
        print("Failure to connect to VNA!")
        print("Check network settings")
    # The VNA ends each line with this. Reads will time out without this
    CMT.read_termination = '\n'
    # Set longer timeout period for slower sweeps
    CMT.timeout = 10000

    CMT.write('SYST:PRESET')

    CMT.write('SENS:SWE:POIN ' + str(NOP))

    CMT.write('SENS:FREQ:STAR ' + str(2.4E9))
    CMT.write('SENS:FREQ:STOP ' + str(2.5E9))

    CMT.write('TRIG:SOUR BUS')
    CMT.query('*OPC?')

    CMT.write('DISP:WIND:SPL 1')
    CMT.write('CALC1:PAR:COUN 1')
    CMT.write('CALC1:PAR1:DEF S21')  # Choose S11 for trace 1
    CMT.write('CALC1:PAR1:SEL')
    CMT.write('CALC1:FORM MLOG')  # log Mag format


def measure():
    CMT.write('TRIG:SING')  # Trigger a single sweep
    CMT.query('*OPC?')  # Wait for measurement to complete
    CMT.write('CALC1:PAR1:SEL')  # Read log mag data

    M = CMT.query("CALC1:DATA:FDAT?")  # Get data as string
    M = M.split(",")
    M = M[::2]
    M = [float(m) for m in M]

    return M[len(M)//2]


def rotate(direction):
    if direction == 1:
        msg = "1\r"
        print("Turning CCW 10°")
    elif direction == -1:
        msg = "2\r"
        print("Turning CW 10°")
    elif direction == 2:
        msg = "3\r"
        print("Turning CCW 180°")
    elif direction == -2:
        msg = "4\r"
        print("Turning CW 180°")

    arduino.write(msg.encode())
    while (arduino.in_waiting == 0):
        pass
    data = str(arduino.readline(), 'utf-8').strip('\r\n')
    while data != "done":
        pass


def plot(mag):
    # angles = [0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 110, 120, 130, 140, 150, 160, 170, 180,
    #           350, 340, 330, 320, 310, 300, 290, 280, 270, 260, 250, 240, 230, 220, 210, 200, 190]
    angles = np.arange(0, 370, 10)
    angles = angles / 180 * np.pi
    # ang = [i / 180 * np.pi for i in angles]

    ax = plt.subplot(111, projection='polar')
    plt.polar(angles, mag)
    plt.show()


def main_loop():
    data = []

    setup_vna()

    for rotation in range(18):
        data.append(measure())
        rotate(1)

    time.sleep(1)
    rotate(-2)
    time.sleep(1)

    for rotation in range(18):
        rotate(-1)
        data.append(measure())

    data.append(data[0])
    time.sleep(1)
    rotate(2)
    time.sleep(1)

    plot(data)

    arduino.close()


main_loop()