# -*- coding: utf-8 -*-
"""
MTB Calibration + Qualification Script (v3)
-----------------------------------------
Auf Wunsch angepasste Zielgrenzen und Iterationen :
* **Kalibrierung**: |ppm| < 1000 als Ziel, max 5 Iterationen.
* **Qualifikation**: ``OK`` ≤ 10 000 ppm (≈1 nA @ 100 µA FS),
  ``COOL`` ≤ 5 000 ppm ; sonst ``FAIL``.
Weitere Features aus v2 bleiben unverändert.
"""

from __future__ import annotations
import sys, time, numpy as np, pyvisa, serial.tools.list_ports
from typing import Sequence
from pymeasure.instruments.keithley import Keithley2400, KeithleyDMM6500
from ModuleTestBox import ModuleTestBox
from sklearn.linear_model import LinearRegression
import json

# ---------- Gerätesuche -------------------------------------------------------
rm = pyvisa.ResourceManager()

def find_visa(hint: str) -> str:
    for res in rm.list_resources():
        if hint.lower() in res.lower():
            return res
    raise RuntimeError(f'VISA-Resource "{hint}" nicht gefunden')

def find_serial(hint: str) -> str:
    for port, desc, _ in serial.tools.list_ports.comports():
        if hint.lower() in desc.lower():
            return port
    raise RuntimeError(f'Serial-Port "{hint}" nicht gefunden')

# ---------- User-Prompt -------------------------------------------------------

def ask_user(prompt: str = "Last umgesteckt?  <Enter>=OK, n=repeat, q=quit → ") -> str:
    while True:
        ans = input(prompt).strip().lower()
        if ans in ("", "y", "yes"):
            return "ok"
        if ans in ("n", "no"):
            return "repeat"
        if ans == "q":
            sys.exit("Abbruch durch Benutzer")
        print("Bitte nur Enter, n oder q eingeben …")

# ---------- Mess-Hilfsfunktionen ---------------------------------------------

VOLTAGES          = np.arange(10, 201, 10)  # 10 … 200 V, 10 V-Step
VOLTAGES_FULLRANGE = np.arange(10, 1001, 10)  # 1 … 100 V, 10 V-Step
CAL_TARGET        = 100                     # ppm-Band für Kalibrierung
CAL_MAX_ITERS     = 5                       # max. Iterationen
QUAL_LIMIT_COOL   = 100                     # ppm für "cool"
QUAL_LIMIT_OK     = 500                     # ppm für "ok" -> 1nA @ 2uA
SHUNT_OHM         = 102.9807197e6           # Referenz-Widerstand (Ω)
SETTLING_S        = 0.2                     # Settling-Zeit pro Step


def sweep_currents(kei: Keithley2400, mtb: ModuleTestBox):
    """Sweep über VOLTAGES - liefert (I_ref, I_dut)."""
    I_ref, I_dut = [], []
    kei.enable_source()
    for v in VOLTAGES:
        kei.source_voltage = v
        time.sleep(SETTLING_S)
        I_ref.append(v / SHUNT_OHM)
        I_dut.append(mtb.GetI() * -1e-12)
    kei.disable_source()
    return np.asarray(I_ref), np.asarray(I_dut)

def sweep_currents_fullRange(kei: Keithley2400, mtb: ModuleTestBox):
    """Sweep über VOLTAGES - liefert (I_ref, I_dut)."""
    I_ref, I_dut = [], []
    kei.enable_source()
    for v in VOLTAGES_FULLRANGE:
        kei.source_voltage = v
        time.sleep(SETTLING_S)
        I_ref.append(v / SHUNT_OHM)
        I_dut.append(mtb.GetI() * -1e-12)
    kei.disable_source()
    return np.asarray(I_ref), np.asarray(I_dut)

def lineaer_fit(x: Sequence[float], y: Sequence[float]) -> tuple[float, float]:
    """Lineare Regression: y = m * x + b"""
    x = np.asarray(x).reshape(-1, 1)
    y = np.asarray(y).reshape(-1, 1)
    model = LinearRegression()
    model.fit(x, y)
    m = model.coef_[0][0]
    b = model.intercept_[0]
    return m, b


def gain_ppm(I_ref: Sequence[float], I_dut: Sequence[float]) -> float:
    slope, offset = np.polyfit(I_ref, I_dut, 1)
    # print(f"  Fit: {curve:.3e} * x² + {slope:.3e} * x + {offset:.3e}")
    return (1.0 / slope - 1.0) * 1e6

# ---------- Hauptprogramm -----------------------------------------------------

dataToSave = []

if __name__ == "__main__":
    serial_id = int(input("Enter the MTB Serial ID [default 100]: ") or "100")
    print(f"Serial ID = {serial_id}")

    with Keithley2400(find_serial("Prolific")) as kei, \
         KeithleyDMM6500(find_visa("0x05E6::0x6500")) as dmm, \
         ModuleTestBox(hints=["VID:PID=CAFE:4001"], verbose=False) as mtb:

        print("Verbunden.  FW:", mtb.GetFirmwareVersion())
        if mtb.HV_IsLocked():
            sys.exit("HV ist gesperrt - Abbruch")

        mtb.SetBoardID(serial_id)
        mtb.SetMeanCount(60)
        dataToSave.append({"serial_id": serial_id, "firmware": mtb.GetFirmwareVersion()}) 
        mtb.HV_SetAll(0)

        kei.apply_voltage(compliance_current=11e-6)  # 11 µA-Limit

        dataChannels = []

        for ch in range(1, 9):
            print(f"\n=== Channel {ch} ===")
            mtb.SelectChannel(ch)
            mtb.HV_Enable(True)
            mtb.AdjustScaleI(0)

            # ------ Kalibrier-Loop ------
            corr_ppm = 0.0
            for it in range(CAL_MAX_ITERS):
                try:
                    mtb.AdjustScaleI(round(corr_ppm))
                except Exception as e:
                    print(f"  Fehler bei der Kalibrierung: {e}")
                    break
                I_ref, I_dut = sweep_currents(kei, mtb)
                ppm = gain_ppm(I_ref, I_dut)
                print(f"  Iter {it}: gain error = {ppm:+.0f} ppm")
                if abs(ppm) < CAL_TARGET:
                    break
                corr_ppm += ppm
            print(f"→ final gain error for channel {ch}: {corr_ppm:+.0f} ppm")

            # ------ Qualifikations-Sweep ------
            I_ref, I_dut = sweep_currents(kei, mtb)
            ppm_chk = abs(gain_ppm(I_ref, I_dut))
            if ppm_chk <= QUAL_LIMIT_COOL:
                status = "COOL"  # ≤5 k
            elif ppm_chk <= QUAL_LIMIT_OK:
                status = "OK"    # ≤10 k
            else:
                status = "FAIL"
            print(f"Qualification: {ppm_chk:.0f} ppm → {status}")

            print("Doing full range sweep …")
            I_ref, I_dut = sweep_currents_fullRange(kei, mtb)

            kei.beep(1000, 0.5)

            # ------ Benutzer Prompt ------
            if ask_user() == "repeat":
                print("Wiederhole Kanal …")
                ch -= 1
                continue

            data = {
                'v': VOLTAGES_FULLRANGE.tolist(),
                'i_ref': I_ref.tolist(),
                'i_dut': I_dut.tolist()
            }

            dataChannels.append({
                "channel": ch,
                "gain_error": corr_ppm,
                "qualification": status,
                "data": data
            })

        mtb.HV_SetAll(0)

        print("\nAlle Kanäle bearbeitet. Kalibrierung beendet.")
        dataToSave.append({"HVchannels": dataChannels})
        print("Tests starten …")
        print("Interlock aktivieren …")

        while True:
            if mtb.HV_IsLocked():
                print("Interlock detected")
                break
            time.sleep(0.1)

        print("Interlock Reset Button drücken")
        print(mtb.HV_GetAll())

        while True:
            if mtb.HV_GetAll():
                print("Interlock released")
                break
            time.sleep(0.1)

        input("Interlock Reset Button gedrückt. <Enter> zum Fortfahren …")
        print("Bias Test starten …")

        dataChannels = []

        dmm.measure_current(max_current=0.001)
        dmm.current_nplc = 12
        
        for ch in range(1, 9):
            print(f"\n=== Channel {ch} ===")
            mtb.SelectChannel(ch)
            mtb.Bias_Enable(True)

            voltages = range(18, 901, 9)
            currents = []
            dmmcurrent = []
            failed = False
            for v in voltages:
                mtb.Bias_SetV(v / 1000)
                time.sleep(SETTLING_S/2)
                dmmcur = dmm.current
                while dmmcur > 1:
                    dmmcur = dmm.current
                dmmcurrent.append(dmmcur * 1e6)
                current = mtb.Bias_GetI()
                currents.append(current * 1e6)

                if abs(current - dmmcur) >= (2e-6 + (v/300 * 1e-6)):
                    print(f"  Bias-Test für Kanal {ch} fehlgeschlagen: {current:.3e} != {dmmcur:.3e}")
                    failed = True

            mtb.Bias_SetV(0)
            mtb.Bias_Enable(False)

            if not failed:
                print(f"Bias-Test für Kanal {ch} erfolgreich")
                status = "OK"
            else:
                status = "FAIL"

            kei.beep(1000, 0.5)

            # ------ Benutzer Prompt ------
            if ask_user() == "repeat":
                print("Wiederhole Kanal …")
                ch -= 1
                continue

            dataChannels.append({
                "channel": ch,
                "voltages": list(voltages),
                "currents": currents,
                "dmmcurrent": dmmcurrent,
                "status": status
            })

        print("\nAlle Kanäle bearbeitet. Bias-Test beendet.")
        dataToSave.append({"Biaschannels": dataChannels})
        print("Tests abgeschlossen. Daten speichern …")
        # check if file exists, if yes, append number to filename
        i = 1
        #save json data to file
        filename = f"TestResults/MTB_Test_{serial_id}.json"
        while True:
            try:
                with open(filename, "x") as f:
                    json.dump(dataToSave, f, indent=4)
                break
            except FileExistsError:
                filename = f"TestResults/MTB_Test_{serial_id}_{i}.json"
                i += 1
        print("Daten gespeichert.")
        print("Programm beendet.")