# -*- coding: utf-8 -*-

#----------
# ※ Tektronix provides the following example "AS IS" with no support or warranty.
# 本スクリプトは無償のサンプル・ファイルとしてお渡しするものであり、本スクリプトに対して、
# 今後弊社は一切のサポートを行うことはできません。ご了承の上ご利用をお願い致します。
# 本パッケージで示したプログラムは、正しく稼働することを保証するものではありません。
# 
# Name： AWG simple sequence builder
# Version: 1.1.1
# 
# 概要：
# 本スクリプトは、AWG5200シリーズ、AWG70000シリーズ、およびSourceXpressソフトウェアにて、
# シーケンス編集を自動化するためのものです。
# 
# 注意：
# シーケンス機能のご利用には、AWG5200/70000シリーズ Opt. SEQの追加が必要となります。
# 
# 推奨機器
# - AWG5200/70000シリーズ、もしくはSourceXpressソフトウェア
#
# 開発言語
# Python3.x
# 
# 必要なPythonモジュール：
# - PyVISA
# - Numpy
#
# 動作確認済環境
# - AWG70001B firmware version 8.3.0032
# - TekVISA Resource Manager v4.3.1.32
# = VISA Shared Components v8.0.7331
# - Python 3.12.8
# - PyVISA 1.15.0
# - Numpy 2.3.4
#
# Author: Shunsuke Katori (Field Application Engineer, Tektronix Japan)
# Original version created on May 27th, 2026
# 
#----------

import sys
import pyvisa
import numpy as np

def seq_build_main():
    # 起動メッセージを表示します
    print('AWG simple sequence builder (version 1.1.0)\n')
    print('※ Tektronix provides the following example "AS IS" with no support or warranty.')
    print('注意：本プログラムは無償のサンプル・ファイルとしてお渡しするものであり、本プログラムに対して、')
    print('今後弊社は一切のサポートを行うことはできません。ご了承の上ご利用をお願い致します。')
    print('本パッケージで示したプログラムは、正しく稼働することを保証するものではありません。\n')

    # PyVISA Resource Manager インスタンスを作成します
    rm = pyvisa.ResourceManager()

    # AWGもしくはSourceXpressソフトウェアとの接続を行います
    # visa_nameは、適宜変更してご利用ください。
    # TCPIP::10.235.36.132::INSTRの部分は使用する計測器のアドレスに合わせて変更ください。
    # 下記の例はLANを使用したアドレス設定のケースです。

    # visa_name = 'GPIB8::1::INSTR'
    # visa_name = 'TCPIP::127.0.0.1::INSTR'
    # visa_name = 'TCPIP::10.235.36.240::4000::SOCKET'
    visa_name = 'TCPIP::10.235.36.132::INSTR'

    try:
        awg = rm.open_resource(visa_name)
    except pyvisa.VisaIOError as e:
        print('AWGとの接続ができません。')
        print('AWGおよびTekVISAが起動していることをご確認ください。')
        input('Enterキーを押してプログラムを終了します。')
        sys.exit(0)
    else:
        awg.timeout = 25000
        devicename = awg.query('*idn?').rstrip('\n')
        if devicename.startswith('TEKTRONIX,AWG520') or devicename.startswith('TEKTRONIX,AWG7000'):
            print('接続が確認できました。\n 機器ID: ' + devicename + '\n')
        else:
            print('接続した機器は、AWG5200シリーズ、AWG70000シリーズ、もしくはSourceXpressソフトウェアではありません。')
            input('Enterキーを押してプログラムを終了します。')
            sys.exit(0)

    # 各パラメータの初期値を入力します。適宜変更してご利用ください。
    sample_rate_gsps = 25 # サンプル・レート[GS/s]
    fc_start_ghz = 0.1 # 開始周波数[GHz]
    fc_stop_ghz = 1 # 終了周波数[GHz]
    num_wfms = 10 # シーケンス・ステップの数
    sweep_time_ms = 1000 # シーケンス全体の出力時間 [ms]

    print('本プログラムでは、入力された開始周波数から停止周波数まで、指定されたステップ数で周波数が変化するシーケンスを作成します。')
    print('以下のパラメータを入力してください。無効な入力の場合は、初期値が適用されます。\n')

    # input start frequency
    str_fc_start = input('入力周波数を GHz 単位で入力してください: ')
    try:
        fc_start_ghz = float(str_fc_start)
    except:
        print('無効な入力です。開始周波数設定の初期値を適用します。\n')
    print('開始周波数は ' + str(fc_start_ghz) + ' GHzです。\n')

    # input stop frequency
    str_fc_stop = input('終了周波数を GHz 単位で入力してください: ')
    try:
        fc_stop_ghz = float(str_fc_stop)
    except:
        print('無効な入力です。終了周波数設定の初期値を適用します。\n')
    print('終了周波数は ' + str(fc_stop_ghz) + ' GHzです。\n')

    # input sample rate
    str_sample_rate = input('サンプル・レートを GS/s 単位で入力してください: ')
    try:
        sample_rate_gsps = float(str_sample_rate)
    except:
        print('無効な入力です。サンプル・レート設定の初期値を適用します。\n')
    print('サンプル・レートは ' + str(sample_rate_gsps) + ' GS/sです。\n')

    # input number of pulses
    str_num_wfms = input('入力波形数を入力してください: ')
    try:
        num_wfms = int(str_num_wfms)
    except:
        print('無効な入力です。波形数の初期値を適用します。\n')
    print('波形数は ' + str(num_wfms) + ' です。\n')

    # input sweep time
    str_sweep_time = input('シーケンス全体の出力時間を ms 単位で入力してください: ')
    try:
        sweep_time_ms = float(str_sweep_time)
    except:
        print('無効な入力です。シーケンス全体の出力時間の初期値を適用します。\n')
    print('シーケンス全体の出力時間は ' + str(sweep_time_ms) + ' msです。\n')

    # 入力されたパラメータをもとに、周波数ステップと1波形あたりのサンプル数を計算します
    freq_step_ghz = (fc_stop_ghz - fc_start_ghz) / (num_wfms - 1)
    size = int(sample_rate_gsps / freq_step_ghz * 80)

    print('周波数ステップは ' + str(freq_step_ghz) + ' GHzです。')
    print('1波形あたりのサンプル数は ' + str(size) + ' です。\n')

    # 波形データ演算の為に、時間軸情報をNumPy配列で作成します。単位はnsとします。
    time_ns = np.arange(size) / sample_rate_gsps

    # 各波形データを作成します
    for i in range(num_wfms):
        fc = fc_start_ghz + (fc_stop_ghz - fc_start_ghz) / (num_wfms - 1) * i # 周波数ステップに従って、各波形の周波数を計算します
        wfm_sin = np.sin(2*np.pi*fc*time_ns).astype(np.float32) # サイン波を作成します。データ形式は32ビット浮動小数点数とします

        awg.write('WLISt:WAVeform:DELete \"test_' + str(i) + '\"') # 同名の波形が存在する場合は削除します
        awg.query('*OPC?')

        awg.write('WLISt:WAVeform:NEW \"test_' + str(i) + '\",' + str(size)) # 波形を新規作成します
        awg.write_binary_values('WLISt:WAVeform:DATA \"test_' + str(i) + '\",0,' + str(size) + ',', wfm_sin) # 波形データを転送します
        awg.query('*OPC?')

        print('test_' + str(i) + ' 波形データが作成されました。')
    
    # シーケンスを作成します

    rcount = int (sweep_time_ms * 1000000 / num_wfms / size * sample_rate_gsps) # 各ステップのリピート回数を計算します

    awg.write('SLISt:SEQuence:DELete \"Sin_Sequence\"') # 同名のシーケンスが存在する場合は削除します。
    awg.query('*OPC?')
    
    awg.write('SLISt:SEQuence:NEW \"Sin_Sequence\",' + str(num_wfms) + ',2') # シーケンスを新規作成します。ステップ数はnum_wfms、トラック数は2とします。
    
    for i in range(num_wfms):
        awg.write('SLISt:SEQuence:STEP' + str(i + 1) + ':TASSet1:WAVeform \"Sin_Sequence\",\"test_' + str(i) + '\"') # 各ステップのトラック1に、作成した波形をセットします
        awg.write('SLISt:SEQuence:STEP' + str(i + 1) + ':TASSet2:WAVeform \"Sin_Sequence\",\"test_' + str(i) + '\"') # 各ステップのトラック2に、作成した波形をセットします
        awg.write('SLISt:SEQuence:STEP' + str(i + 1) + ':RCOunt \"Sin_Sequence\", ' + str(rcount)) # 各ステップのリピート回数を設定します
        # awg.write('SLISt:SEQuence:STEP' + str(i + 1) + ':EJINPut \"Sin_Sequence\", ATRigger') # 各ステップにトリガ待ちを加える場合は、ここを非コメント化してください
        # awg.write('SLISt:SEQuence:STEP' + str(i + 1) + ':GOTO \"Sin_Sequence\", NEXT') # 各ステップの遷移先を設定します。ここでは、次のステップに遷移するように設定しています
        awg.query('*OPC?')

    awg.write('SLISt:SEQuence:STEP' + str(i + 1) + ':GOTO \"Sin_Sequence\", FIRSt') # 最終ステップの遷移先を設定します。ここでは、最初のステップに遷移するように設定しています
    awg.write('SOURce1:CASSet:SEQuence \"Sin_Sequence\",1') # Ch1にシーケンスをセットします
    # awg.write('SOURce2:CASSet:SEQuence \"Sin_Sequence\",2') # Ch2にシーケンスをセットする場合は、ここを非コメント化してください
    
    print('波形データとシーケンスの生成が完了しました。\n')

    # 最後に、AWGのサンプル・レートを、最初に入力したサンプル・レートの値に設定します。
    # ただし、お使いのAWGの最高サンプル・レートを超える値に設定することはできません。
    # プログラム実行後に、実際にサンプル・レート設定がいくつになっているか、ご確認をお願いします。
    awg.write('CLOCk:SRATe ' + str(sample_rate_gsps) + 'e9')
    print('波形を出力の際には、AWG サンプル・レート設定をご確認ください。\n')

    awg.close()

if __name__ == '__main__':
    seq_build_main()
