import time
import os
import sys
import serial
import csv
from tsapython import tinySA
import datetime
import numpy as np
import matplotlib.pyplot as plt

ser = serial.Serial("/dev/ttyACM2", 115200)

max_x = 250
max_y = 200
max_z = 150

min_z = 50
z_search = 100


prboffs_x = 40
prboffs_y = 40
prboffs_z = 10


feedrate = 5000

now = datetime.datetime.now()


def convert_data_to_arrays(start, stop, pts, data):
    # using the start and stop frequencies, and the number of points,

    freq_arr = np.linspace(
        start, stop, pts
    )  # note that the decimals might go out to many places.
    # you can truncate this because it’s only used
    # for plotting in this example

    # As of the Jan. 2024 build in some data returned with SWEEP or SCAN calls there is error data.
    # https://groups.io/g/tinysa/topic/tinasa_ultra_sweep_command/104194367
    # this shows up as "-:.000000e+01".
    # TEMP fix - replace the colon character with a -10. This puts the 'filled in' points around the noise floor.
    # more advanced filtering should be applied for actual analysis.
    data1 = bytearray(data.replace(b"-:.0", b"-10.0"))

    # get both values in each row returned (for reference)
    # data_arr = [list(map(float, line.split())) for line in data.decode('utf-8').split('\n') if line.strip()]

    # get first value in each returned row
    data_arr = [
        float(line.split()[0])
        for line in data1.decode("utf-8").split("\n")
        if line.strip()
    ]

    return freq_arr, data_arr


def writeCSV(freq_array, data, csv_filename):
    try:
        with open(csv_filename, "w", newline="") as csvfile:
            writer = csv.writer(csvfile, delimiter=";")
            writer.writerow(
                f'# Nearfield Scan {now.strftime("%Y-%m-%d %H:%M")}, {desc}'
            )
            row = []
            row.append("# ;;")
            for elem in freq_array:
                row.append(str(elem))
            writer.writerow(row)

            for key in data:
                row = []
                row.append(key, data[key])
                writer.writerow(row)
    except IOError:
        print("I/O error")
    return


def initXYZ():
    ser.write(str.encode("M17 \r\n"))
    ser.write(str.encode("M84 X Y Z S12000\r\n"))
    ser.write(str.encode("M84 E\r\n"))
    ser.write(str.encode("G28\r\n"))

    # set to absolute movement
    ser.write(str.encode("G90\r\n"))
    time.sleep(20)


def main():
    import argparse
    import configparser

    parser = argparse.ArgumentParser(
        prog="3D NearField Scanner",
        description="Uses a 3D Printer and a Tiny SA for measureing the near field of a PCB",
        epilog="=^.^=",
    )

    parser.add_argument("--fstart", type=float, default=1e6)

    parser.add_argument("--fstop", type=float, default=800e6)

    parser.add_argument("--pts", type=int, default=450)

    parser.add_argument(
        "--rbw",
        help="RBW in kHz",
        default=auto,
        choices=[200, 1, 3, 10, 30, 100, 300, 600, 850],
    )

    parser.add_argument("--xstart", type=int, default=0, help="x start in mm")

    parser.add_argument("--xend", type=int, default=250, help="x end in mm")

    parser.add_argument("--ystart", type=int, default=0, help="y start in mm")

    parser.add_argument("--yend", type=int, default=200, help="y end in mm")
    parser.add_argument("--xstep", type=int, default=10, help="x step in mm")

    parser.add_argument("--ystep", type=int, default=10, help="y step in mm")

    parser.add_argument("--desc", type=str, default="", help="EUT Description")

    args = parser.parse_args()

    # scan surface
    ser.write(str.encode(f"G1 Z{min_z}\r\n"))
    input("Attach Probe and press Enter to continue...")
    print("Starting Measurement")

    # create a new tinySA object
    tsa = tinySA()

    # set the return message preferences
    tsa.set_verbose(True)  # detailed messages
    tsa.set_error_byte_return(True)  # get explicit b'ERROR' if error thrown

    # attempt to autoconnect
    found_bool, connected_bool = tsa.autoconnect()

    # if port closed, then return error message
    if connected_bool == False:
        print("ERROR: could not connect to port")
    else:  # if port found and connected, then complete task(s) and disconnect

        if args.rbw == "auto":
            tsa.set_rbw_auto()
        else:
            tsa.rbw(int(args.rbw))

        datadict = {}

        for y in range(args.ystart, args.yend, args.ystep):
            if y >= max_y - args.ystep:
                break
            else:
                ser.write(str.encode(f"G1 X{args.xstart} Y{y} F{feedrate}\r\n"))
                time.sleep(2)

            for x in range(args.xstart, args.xend, args.xstep):
                if x >= max_x - args.xstep:
                    break
                else:

                    ser.write(str.encode(f"G1 X{x} Z{z_search} F{feedrate}\r\n"))
                    time.sleep(1)
                    ser.write(str.encode(f"G1 Z{min_z} F{feedrate}\r\n"))
                    ser.write(str.encode("M18 \r\n"))  # disable steppers
                    tsa.resume()  # scan
                    data_bytes = tsa.scan(args.fstart, args.fstop, args.pts, 2)
                    tsa.wait()
                    ser.write(str.encode("M17 X Y Z\r\n"))  # enable steppers
                    ser.write(str.encode(f"G1 Z{z_search} F{feedrate}\r\n"))

                    freq_arr, data_arr = convert_data_to_arrays(
                        start, stop, pts, data_bytes
                    )
                    datadict[f"{str(x)};{str(y)}"] = data_arr
                    time.sleep(1)
        writeCSV(freq_arr, datadict, "test.csv")
        tsa.resume()  # resume so screen isn't still frozen
        tsa.disconnect()


if __name__ == "__main__":
    sys.exit(main())