import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import re
import telnetlib3
import csv
import time

Z_table = [
[	0.00007362	,	0.00007713	],
[	0.0005028	,	-0.00007668	]
]

Z2_table = [
[	0.000482934	,	-0.000453644	],
[	0.000706542	,	0.000060724	]
]

data_rows = []   # global 2-D list
state = {
    "remaining_receive_lines": 0,
    "freq": 10000.0,
    "freq_step_multiply": 0.85,
    "stop_freq": 6000,
    "initializing": 1
}

PATTERN = re.compile(
    r"Va=(?P<Va>-?\d+\.?\d*)\s+"
    r"Vp=(?P<Vp>-?\d+\.?\d*)\s+\|\s+"
    r"Ia=(?P<Ia>-?\d+\.?\d*)\s+"
    r"Ip=(?P<Ip>-?\d+\.?\d*)\s+\|\s+"
    r"ZR=(?P<ZR>-?\d+\.?\d*)\s+"
    r"ZX=(?P<ZX>-?\d+\.?\d*)"
)

def dump_csv(filename="data.csv"):
    df = pd.DataFrame(
        data_rows,
        columns=["Freq", "Va", "Vp", "Ia", "Ip", "ZR", "ZX"]
    )
    df.to_csv(filename, index=False)

def append_line_to_file(column_nr, filename):
    # Extract freq column (index 0)
    # Extract ZR column (index 5)
    # Extract ZX column (index 6)
    zr_values = [row[column_nr] for row in data_rows if len(row) > column_nr]
    if not zr_values:
        print("No ZR data to write")
        return
    with open(filename, "a", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(zr_values)

def append_Nyquist_run(filename):
    # Flatten ZR/ZX pairs into one row
    row = []
    for r in data_rows:
        if len(r) > 6:
            row.extend([r[5], r[6]])
    if not row:
        print("No ZR/ZX data to write")
        return
    with open(filename, "a", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(row)


def extract_to_dataframe(line):
    # Extract Va, Vp, Ia, Ip, ZR, ZX from a line and append them as a row to the global data_rows list.
    global state
    match = PATTERN.search(line)
    if not match:
        return  # silently ignore malformed lines
    row = [
        float(state["freq"]),
        float(match.group("Va")),
        float(match.group("Vp")),
        float(match.group("Ia")),
        float(match.group("Ip")),
        float(match.group("ZR")),
        float(match.group("ZX")),
    ]
    data_rows.append(row)

def process_line(line):
    extract_to_dataframe(line)
    # Example: print last row
    if data_rows:
        print("Last row:", data_rows[-1])

def get_dataframe():
    return pd.DataFrame(
        data_rows,
        columns=["Va", "Vp", "Ia", "Ip", "ZR", "ZX"]
    )

class TelnetReader:
    def __init__(self, host, port=23, timeout=10):
        self.host = host
        self.port = port
        self.timeout = timeout
        self.tn = None

    def connect(self):
        print(f"Connecting to {self.host}:{self.port} ...")
        self.tn = telnetlib3.Telnet(self.host, self.port, self.timeout)
        print("Connected.")

    def disconnect(self):
        if self.tn:
            self.tn.close()
            self.tn = None
            print("Disconnected.")

    def read_loop(self):
        global state
        # Main loop: reads incoming lines forever
        try:
            while state["freq"] > state["stop_freq"]:
                line = self.tn.read_until(b"\n")
                if not line:
                    break

                decoded = line.decode("utf-8", errors="ignore").strip()
                self.process_line(decoded)

        except KeyboardInterrupt:
            print("Interrupted by user.")
        print(data_rows)
        dump_csv("measurements.csv")
        append_line_to_file(0, "scan_freq.csv")
        append_line_to_file(1, "scan_Va.csv")
        append_line_to_file(2, "scan_Vp.csv")
        append_line_to_file(3, "scan_Ia.csv")
        append_line_to_file(4, "scan_Ip.csv")
        append_line_to_file(5, "scan_ZR.csv")
        append_line_to_file(6, "scan_ZX.csv")
        append_Nyquist_run("scan_Nyquist.csv")

    def process_line(self, line):
        global state
        # Override or extend this method to extract data.
#        print(f"RAW: {line}")
        if state["remaining_receive_lines"] > 0 :
            state["remaining_receive_lines"] -= 1           # we are waiting for settling - just skip the line
        else:
            # prepare new frequency measurement
            if not state["initializing"] == 1:
                extract_to_dataframe(line)  # capture the measurement
                if state["freq"] > 1.0:
                    state["freq"] *= state["freq_step_multiply"]     # calc next freq
                else:
                    if state["freq"] > 0.07:
                        state["freq"] *= state["freq_step_multiply"] ** 2  # skip 1/2 steps to speed up
                    else:
                        state["freq"] *= state["freq_step_multiply"] ** 4  # skip 3/4 steps to speed up
            if state["freq"] <= state["stop_freq"]:
                print("Reached stop frequency")
            else:
                if state["freq"] > 1:
                    response = f"\rf{state["freq"]:.1f}\r"        # new freq
                else:
                    response = f"\rf{state["freq"]:.3f}\r"  # new freq
                self.tn.write(response.encode("utf-8"))
                bandwidth = state["freq"]/20
                if bandwidth > 1 :
                    bandwidth = 1
                response = f"b{bandwidth:.3f}\r"        # new freq
                self.tn.write(response.encode("utf-8"))
                state["remaining_receive_lines"] = 1 + 4/bandwidth
                print(line)
                print(state)
            state["initializing"] = 0

        # Example:
        # value = self.extract_value(line)
        # self.store_value(value)

    # Example placeholder methods
    def extract_value(self, line):
        pass

    def store_value(self, value):
        pass


if __name__ == "__main__":
    reader = TelnetReader(host="10.1.122.152", port=2002)
#    reader = TelnetReader(host="192.168.1.196", port=2002)

    try:
        reader.connect()
        while True:
            data_rows.clear()
            state = {
                "remaining_receive_lines": 0,
                "freq": 10.0,
                "freq_step_multiply": 0.8,
                "stop_freq": 0.1,
                "initializing": 1
            }
            reader.read_loop()
    finally:
        reader.disconnect()

"""
Z_points = [complex(r, i) for r, i in Z_table]
Zp = np.array(Z_points)
Z2_points = [complex(r, i) for r, i in Z2_table]
Z2p = np.array(Z2_points)

plt.plot(
    Zp.real,
    Zp.imag,
    marker='x',
    color='blue',
    label="shunt-300u"
)
plt.plot(
    Z2p.real,
    Z2p.imag,
    marker='x',
    color='red',
    label="clamp"
)

plt.xlabel("Re{Z} [Ω]")
plt.ylabel("-Im{Z} [Ω]")
plt.title("Nyquist Plot")
plt.grid(True)
plt.legend()
plt.show()
"""