import os,sys,time
import casperfpga
import numpy as np
import matplotlib.pyplot as plt
from optparse import OptionParser

if __name__=="__main__":
  p = OptionParser()
  p.set_usage('dp_loopback.py <HOSTNAME_or_IP> [options]')
  p.add_option('-s', '--skip', dest='skip', action='store_true',
      help='Skip programming and begin to plot data')

  # get board hostname or ip address
  opts, args = p.parse_args(sys.argv[1:])
  if len(args) < 1:
    print('Specify a hostname or IP for your casper platform.\n'
          'Run with the -h flag to options.')

    exit()
  else:
    hostname = args[0]

  # ip/mac/port variables
  ip_base  = 10*(2**24) + 17*(2**16) + 16*(2**8)
  mac_base = 10*(2**24) + 17*(2**16)
  fabric_port = 60000

  # packet simulator variables
  pkt_period = 16384
  pkt_payload_words = 32

  # check for fpg file
  fpgfile = './zcu216_dp_example_2025-09-03_1308.fpg'
  if not os.path.exists(fpgfile):
    print("This script must run from the same directory where {:s} is located".format(fpgfile))
    sys.exit(1)

  # connect and program
  print('Connecting to {:s}'.format((hostname)))
  z = casperfpga.CasperFpga(hostname)
  if not opts.skip:
    print("Programming fpga...")
    z.upload_to_ram_and_program(fpgfile)
    time.sleep(0.2)
  else:
    print("Skipped programming...")
    z.get_system_information()

  if not opts.skip:
    # get gbe objects
    gbe_port1 = z.gbes.onehundred_gbe_port1
    gbe_port2 = z.gbes.onehundred_gbe_port2

    # setup up the transmitter
    print("Configuring gbe cores...")
    z.registers.port1_dest_port.write(reg=fabric_port)
    z.registers.port1_dest_ip.write(reg=ip_base+17)

    gbe_port1.set_single_arp_entry('10.17.16.17', mac_base+17)
    gbe_port1.configure_core(mac_base+16, ip_base+16, fabric_port)

    # setup receiver
    gbe_port2.configure_core(mac_base+17, ip_base+17, fabric_port)

    # setup packet simulator
    z.registers.pkt_sim_port1_period.write(reg=pkt_period)
    z.registers.pkt_sim_port1_payload_len.write(reg=pkt_payload_words)

    print("waiting for core to startup...")
    time.sleep(6)

    # start transmitting packets from port 1 to port port 2
    print("enable traffic generator...")
    z.registers.pkt_sim_port1_enable.write(reg=1)

  # check results
  print("verifying tx operation...")
 
  # read the tx packet counter, make sure it is increasing
  tic = z.registers.gbe_tx_cnt_port1.read()
  time.sleep(1)
  toc = z.registers.gbe_tx_cnt_port1.read()

  delta = toc['data']['reg'] - tic['data']['reg']
  if delta <=0:
    print("Not detecting transmitted packets...")
    print(tic)
    print(toc)
    sys.exit(1)

  print("Verified TX operation!")

  # read the rx packet counter
  print("verifying rx operation")
  tic = z.registers.gbe_rx_frame_cnt_port2.read()
  time.sleep(1)
  toc = z.registers.gbe_rx_frame_cnt_port2.read()

  delta = toc['data']['reg'] - tic['data']['reg']
  if delta <=0:
    print(tic)
    print(toc)
    print("Not detecting received packets...")
    sys.exit(1)

  print("Verified RX operation!")

  print("Plotting packet data...")
  # extract snapshot data for received data on port
  z.snapshots.bitfield_snapshot2_ss.arm()
  d = z.snapshots.bitfield_snapshot2_ss.read(arm=False)['data']

  # `d` has 'data', 'vld', 'eof' fields. The packet simulator sends
  # sequential counter data. We can extract the data and plot the data
  # and show the counts are ascending in order
  pkt_data = d['data']
  plt.plot(np.arange(0,len(pkt_data)), pkt_data)
  plt.show()