"""
keithley2400 — Keithley 2400 Source Measure Unit (SMU) control library

Controls the Keithley 2400 over RS-232 (USB-serial adapter) using PyVISA
with the pure-Python backend (pyvisa-py). No National Instruments VISA
installation required — works natively on macOS, Linux, and Windows,
including Apple Silicon (M1/M2/M3/M4).

Typical usage:
    from keithley2400 import Keithley2400, find_serial_port, quick_iv_sweep

Requirements:
    pyvisa >= 1.11.0
    pyvisa-py >= 0.5.2
    pyserial >= 3.4
"""

from __future__ import annotations

import csv
import statistics
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Iterable, Iterator, Literal, Optional, Tuple

import pyvisa
from pyvisa import constants


# ---------------------------------------------------------------------------
# Type aliases
# ---------------------------------------------------------------------------

Param = Literal["V", "I", "R"]
RMode = Literal["calc", "ohms"]


# ---------------------------------------------------------------------------
# Serial settings dataclass
# ---------------------------------------------------------------------------

@dataclass
class SerialSettings:
    """RS-232 communication parameters for the Keithley 2400.

    The factory defaults that match the instrument's out-of-box GPIB/RS-232
    configuration are used here — change only if you have manually altered
    the instrument's communication settings.
    """
    baud_rate: int = 9600
    data_bits: int = 8
    parity: constants.Parity = constants.Parity.none
    stop_bits: constants.StopBits = constants.StopBits.one
    terminator: str = "\r"
    timeout_ms: int = 5000
    xonxoff: bool = False


# ---------------------------------------------------------------------------
# Main instrument class
# ---------------------------------------------------------------------------

class Keithley2400:
    """
    Keithley 2400 Source Measure Unit — RS-232 (USB-serial) control via PyVISA.

    Communicates using SCPI commands:
      - SOUR subsystem  : set voltage / current source
      - SENS subsystem  : configure measurement function and range
      - OUTP            : enable / disable output
      - :READ?          : acquire a single reading
      - FORM:ELEM       : select which values :READ? returns (V, I, R)
      - SYST:ERR?       : query instrument error queue

    Connection
    ----------
    Connects over RS-232 via a USB-serial adapter using the ``@py`` (pyvisa-py)
    backend.  Find your port with :func:`find_serial_port`.

    Example
    -------
    >>> from keithley2400 import Keithley2400
    >>> k = Keithley2400(resource="ASRL/dev/cu.usbserial-FTRTKC3X::INSTR")
    >>> k.set_source_v(5.0)
    >>> k.set_current_compliance(0.01)
    >>> k.configure_measure_vi()
    >>> k.output_on()
    >>> v, i = k.read_v_i()
    >>> print(f"V={v:.4f} V  I={i:.6f} A")
    >>> k.output_off()
    >>> k.close()
    """

    def __init__(
        self,
        resource: str = "ASRL/dev/cu.usbserial-FTRTKC3X::INSTR",
        backend: str = "@py",
        settings: Optional[SerialSettings] = None,
    ) -> None:
        """
        Open connection to the Keithley 2400.

        Parameters
        ----------
        resource : str
            PyVISA resource string for the serial port.
            Use :func:`find_serial_port` to discover available ports.
        backend : str
            PyVISA backend specifier.  ``"@py"`` uses pyvisa-py (no NI-VISA
            needed).  Pass ``""`` to use the NI-VISA backend if installed.
        settings : SerialSettings, optional
            RS-232 parameters.  Defaults to factory settings (9600 8N1).

        Raises
        ------
        pyvisa.errors.VisaIOError
            If the resource cannot be opened (wrong port, device off, etc.).
        """
        self.resource = resource
        self.backend = backend
        self.settings = settings or SerialSettings()

        self.rm = pyvisa.ResourceManager(self.backend)
        self.inst = self.rm.open_resource(self.resource)

        self._apply_serial_settings()
        self.initialize()

    # ------------------------------------------------------------------
    # Internal setup
    # ------------------------------------------------------------------

    def _apply_serial_settings(self) -> None:
        s = self.settings
        self.inst.timeout = s.timeout_ms
        self.inst.baud_rate = s.baud_rate
        self.inst.data_bits = s.data_bits
        self.inst.parity = s.parity
        self.inst.stop_bits = s.stop_bits
        self.inst.write_termination = s.terminator
        self.inst.read_termination = s.terminator

        try:
            self.inst.flow_control = (
                constants.VI_ASRL_FLOW_XON_XOFF if s.xonxoff
                else constants.VI_ASRL_FLOW_NONE
            )
        except Exception:
            try:
                self.inst.xonxoff = bool(s.xonxoff)
            except Exception:
                pass

    def initialize(self) -> None:
        """Reset instrument to known state and configure output format.

        Called automatically by ``__init__``.  Can be called again to
        restore a clean state without reconnecting.
        """
        self.write("*RST")
        self.write("*CLS")
        self.write("OUTP OFF")
        # Default: :READ? returns voltage and current only — deterministic parsing
        self.write("FORM:ELEM VOLT,CURR")

    # ------------------------------------------------------------------
    # Connection lifecycle
    # ------------------------------------------------------------------

    def close(self) -> None:
        """Close the VISA session and release resources."""
        try:
            self.inst.close()
        finally:
            self.rm.close()

    def __enter__(self) -> "Keithley2400":
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        self.output_off()
        self.close()

    # ------------------------------------------------------------------
    # Low-level SCPI I/O
    # ------------------------------------------------------------------

    def write(self, cmd: str) -> None:
        """Send a SCPI command (no response expected)."""
        self.inst.write(cmd)

    def query(self, cmd: str) -> str:
        """Send a SCPI command and return the stripped response."""
        return self.inst.query(cmd).strip()

    def idn(self) -> str:
        """Query and return the instrument identification string (*IDN?)."""
        return self.query("*IDN?")

    def get_error(self) -> str:
        """Query the error queue (SYST:ERR?).  Returns '0,"No error"' if clear."""
        return self.query("SYST:ERR?")

    # ------------------------------------------------------------------
    # Output control
    # ------------------------------------------------------------------

    def output_on(self) -> None:
        """Enable the source output (OUTP ON)."""
        self.write("OUTP ON")

    def output_off(self) -> None:
        """Disable the source output (OUTP OFF)."""
        self.write("OUTP OFF")

    # ------------------------------------------------------------------
    # Source configuration
    # ------------------------------------------------------------------

    def set_source_v(self, volts: float) -> None:
        """Configure instrument to source voltage.

        Parameters
        ----------
        volts : float
            Source voltage in volts.
        """
        self.write("SOUR:FUNC VOLT")
        self.write(f"SOUR:VOLT {volts}")

    def set_source_i(self, amps: float) -> None:
        """Configure instrument to source current.

        Parameters
        ----------
        amps : float
            Source current in amperes.
        """
        self.write("SOUR:FUNC CURR")
        self.write(f"SOUR:CURR {amps}")

    def set_current_compliance(self, amps: float) -> None:
        """Set current compliance limit (used when sourcing voltage).

        Parameters
        ----------
        amps : float
            Maximum allowed current in amperes.
        """
        self.write(f"SENS:CURR:PROT {amps}")

    def set_voltage_compliance(self, volts: float) -> None:
        """Set voltage compliance limit (used when sourcing current).

        Parameters
        ----------
        volts : float
            Maximum allowed voltage in volts.
        """
        self.write(f"SENS:VOLT:PROT {volts}")

    # ------------------------------------------------------------------
    # Measurement configuration
    # ------------------------------------------------------------------

    def configure_measure_vi(self, auto_range: bool = True) -> None:
        """Configure simultaneous voltage and current measurement.

        Sets FORM:ELEM to VOLT,CURR so that :READ? returns exactly two
        comma-separated values: voltage then current.

        Parameters
        ----------
        auto_range : bool
            Use auto-ranging for both V and I measurements (default True).
        """
        self.write("FORM:ELEM VOLT,CURR")
        self.write("SENS:FUNC 'VOLT','CURR'")
        self.write(f"SENS:VOLT:RANG:AUTO {'ON' if auto_range else 'OFF'}")
        self.write(f"SENS:CURR:RANG:AUTO {'ON' if auto_range else 'OFF'}")

    def configure_measure_ohms(self, method: Literal["2W", "4W"] = "2W") -> None:
        """Configure true resistance measurement mode.

        Sets FORM:ELEM to RES so that :READ? returns a single resistance value.
        Call :meth:`set_ohms_source_current` or :meth:`set_ohms_source_voltage`
        to configure the ohms source before measuring.

        Parameters
        ----------
        method : "2W" or "4W"
            Two-wire (default) or four-wire (Kelvin) sense connection.
        """
        self.write("SENS:FUNC 'RES'")
        self.write("SYST:RSEN ON" if method == "4W" else "SYST:RSEN OFF")
        self.write("FORM:ELEM RES")

    def set_current_range(self, amps: float) -> None:
        """Manually set the current measurement range.

        Parameters
        ----------
        amps : float
            Full-scale current range in amperes (e.g., 0.01 for 10 mA range).
        """
        self.write(f"SENS:CURR:RANG {amps}")

    def set_ohms_source_current(self, amps: float) -> None:
        """Use current source for ohms measurement.

        Parameters
        ----------
        amps : float
            Test current in amperes.
        """
        self.write("SOUR:FUNC CURR")
        self.write(f"SOUR:CURR {amps}")

    def set_ohms_source_voltage(self, volts: float) -> None:
        """Use voltage source for ohms measurement.

        Parameters
        ----------
        volts : float
            Test voltage in volts.
        """
        self.write("SOUR:FUNC VOLT")
        self.write(f"SOUR:VOLT {volts}")

    # ------------------------------------------------------------------
    # Readings
    # ------------------------------------------------------------------

    def _read_str(self) -> str:
        return self.query(":READ?")

    @staticmethod
    def _parse_csv_floats(s: str) -> Tuple[float, ...]:
        parts = [p.strip() for p in s.split(",") if p.strip()]
        return tuple(float(p) for p in parts)

    def read_v_i(self) -> Tuple[float, float]:
        """Read voltage and current simultaneously.

        Requires ``configure_measure_vi()`` to have been called (or FORM:ELEM
        to be set to VOLT,CURR).

        Returns
        -------
        (voltage, current) : Tuple[float, float]
            Measured voltage in volts and current in amperes.

        Raises
        ------
        ValueError
            If the instrument does not return exactly 2 fields.
        """
        vals = self._parse_csv_floats(self._read_str())
        if len(vals) != 2:
            raise ValueError(
                f"Expected 2 fields (V, I), got {len(vals)} from: {vals!r}. "
                "Ensure configure_measure_vi() was called."
            )
        return vals[0], vals[1]

    def read_v(self) -> float:
        """Read and return the measured voltage in volts."""
        v, _ = self.read_v_i()
        return v

    def read_i(self) -> float:
        """Read and return the measured current in amperes."""
        _, i = self.read_v_i()
        return i

    def read_r(self, mode: RMode = "calc", i_min_abs: float = 1e-9) -> float:
        """Read resistance.

        Parameters
        ----------
        mode : "calc" or "ohms"
            ``"calc"``  — compute R = V/I from a simultaneous V+I reading.
            ``"ohms"``  — use the instrument's built-in resistance measurement
                         (requires ``configure_measure_ohms()`` beforehand).
        i_min_abs : float
            Guard threshold for calc mode: raises if |I| < i_min_abs.

        Returns
        -------
        float
            Resistance in ohms.
        """
        if mode == "ohms":
            vals = self._parse_csv_floats(self._read_str())
            if len(vals) != 1:
                raise ValueError(
                    f"Expected 1 field (RES), got {len(vals)} from: {vals!r}. "
                    "Ensure configure_measure_ohms() was called."
                )
            return vals[0]

        v, i = self.read_v_i()
        if abs(i) < i_min_abs:
            raise ValueError(
                f"Current too small for reliable R=V/I "
                f"(|I|={abs(i):.2e} < threshold {i_min_abs:.2e} A). "
                f"V={v:.6f} V, I={i:.6e} A"
            )
        return v / i

    # ------------------------------------------------------------------
    # Statistical sampling
    # ------------------------------------------------------------------

    def sample_stats(
        self,
        param: Param,
        n: int = 50,
        delay_s: float = 0.01,
        stabilize_s: float = 0.0,
        r_mode: RMode = "calc",
        i_min_abs_for_r: float = 1e-9,
    ) -> Dict[str, float]:
        """Sample a parameter N times and return mean and standard deviation.

        Parameters
        ----------
        param : "V", "I", or "R"
            Quantity to sample.
        n : int
            Number of samples (must be >= 2).
        delay_s : float
            Inter-sample delay in seconds (default 0.01).
        stabilize_s : float
            Wait time before sampling begins, to allow the circuit to settle.
        r_mode : "calc" or "ohms"
            Resistance calculation mode (only used when param="R").
        i_min_abs_for_r : float
            Current guard threshold for R=V/I mode.

        Returns
        -------
        dict
            ``{"mean": float, "stdev": float, "n": float}``
        """
        if n < 2:
            raise ValueError("n must be >= 2 to compute standard deviation")

        if stabilize_s > 0:
            time.sleep(stabilize_s)

        samples = []
        for _ in range(n):
            if param == "V":
                samples.append(self.read_v())
            elif param == "I":
                samples.append(self.read_i())
            elif param == "R":
                samples.append(self.read_r(mode=r_mode, i_min_abs=i_min_abs_for_r))
            else:
                raise ValueError(f"Unknown param: {param!r}. Must be 'V', 'I', or 'R'.")

            if delay_s > 0:
                time.sleep(delay_s)

        return {
            "mean": float(statistics.mean(samples)),
            "stdev": float(statistics.stdev(samples)),
            "n": float(n),
        }

    # ------------------------------------------------------------------
    # IV sweep
    # ------------------------------------------------------------------

    def sweep_iv(
        self,
        voltages: Iterable[float],
        compliance_amps: float = 0.01,
        stabilize_s: float = 0.2,
        n_samples: int = 50,
        sample_delay_s: float = 0.01,
        current_range_amps: Optional[float] = None,
    ) -> Iterator[Tuple[float, float, float]]:
        """Sweep a list of voltages and yield (V_set, I_mean, I_stdev) tuples.

        The output must be enabled *before* calling this method.  The method
        does not enable or disable the output itself, so the caller controls
        when current flows.

        Parameters
        ----------
        voltages : Iterable[float]
            Ordered sequence of source voltages in volts.
        compliance_amps : float
            Current compliance limit in amperes.
        stabilize_s : float
            Wait time after changing voltage before measuring (default 0.2 s).
        n_samples : int
            Number of current samples per voltage step.
        sample_delay_s : float
            Inter-sample delay in seconds.
        current_range_amps : float, optional
            Manual current range (e.g., 0.01 for 10 mA). If None, auto-range is used.

        Yields
        ------
        (v_set, i_mean, i_stdev) : Tuple[float, float, float]
            Set voltage, mean measured current, and standard deviation.
        """
        self.set_current_compliance(compliance_amps)
        self.configure_measure_vi(auto_range=(current_range_amps is None))
        if current_range_amps is not None:
            self.set_current_range(current_range_amps)

        for v in voltages:
            self.set_source_v(float(v))
            stats = self.sample_stats(
                param="I",
                n=n_samples,
                delay_s=sample_delay_s,
                stabilize_s=stabilize_s,
            )
            yield float(v), stats["mean"], stats["stdev"]


# ---------------------------------------------------------------------------
# Convenience functions
# ---------------------------------------------------------------------------

def find_serial_port(backend: str = "@py") -> list:
    """List all ASRL (serial) resources visible to PyVISA.

    Use this to find the correct resource string for your USB-serial adapter.

    Parameters
    ----------
    backend : str
        PyVISA backend (default ``"@py"``).

    Returns
    -------
    list of str
        Resource strings for all detected serial ports.

    Example
    -------
    >>> from keithley2400 import find_serial_port
    >>> ports = find_serial_port()
    >>> print(ports)
    ['ASRL/dev/cu.usbserial-FTRTKC3X::INSTR', ...]
    """
    rm = pyvisa.ResourceManager(backend)
    resources = rm.list_resources()
    rm.close()
    serial = [r for r in resources if r.startswith("ASRL")]
    if serial:
        print("Detected serial ports:")
        for r in serial:
            print(f"  {r}")
    else:
        print("No serial (ASRL) ports detected.")
    return serial


def quick_iv_sweep(
    voltages: Iterable[float],
    out_csv: Optional[str] = None,
    resource: str = "ASRL/dev/cu.usbserial-FTRTKC3X::INSTR",
    compliance_amps: float = 0.01,
    stabilize_s: float = 0.2,
    n_samples: int = 50,
    sample_delay_s: float = 0.01,
    current_range_amps: Optional[float] = None,
) -> Optional[str]:
    """Connect, run a voltage sweep, save results to CSV, and disconnect.

    Parameters
    ----------
    voltages : Iterable[float]
        Sequence of source voltages in volts.
    out_csv : str, optional
        Output CSV file path (auto-generated from timestamp if None).
    resource : str
        PyVISA resource string for the Keithley 2400.
    compliance_amps : float
        Current compliance limit in amperes.
    stabilize_s : float
        Stabilization wait after each voltage step.
    n_samples : int
        Number of current samples per voltage step.
    sample_delay_s : float
        Inter-sample delay in seconds.
    current_range_amps : float, optional
        Manual current range. If None, auto-range is used.

    Returns
    -------
    str or None
        Path to the saved CSV file, or None if the sweep failed.

    Example
    -------
    >>> from keithley2400 import quick_iv_sweep
    >>> voltages = [v * 0.5 for v in range(0, 21)]  # 0 to 10 V in 0.5 V steps
    >>> csv_path = quick_iv_sweep(voltages, out_csv="iv_sweep.csv", compliance_amps=0.01)
    >>> print(f"Saved: {csv_path}")
    """
    if out_csv is None:
        from datetime import datetime
        out_csv = f"iv_sweep_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"

    print(f"Connecting to {resource} ...")
    try:
        k = Keithley2400(resource=resource)
    except Exception as e:
        print(f"Connection failed: {e}")
        return None

    print(f"Connected: {k.idn()}")
    print(f"Output: {out_csv}")

    rows = []
    k.output_on()
    try:
        for v_set, i_mean, i_stdev in k.sweep_iv(
            voltages=voltages,
            compliance_amps=compliance_amps,
            stabilize_s=stabilize_s,
            n_samples=n_samples,
            sample_delay_s=sample_delay_s,
            current_range_amps=current_range_amps,
        ):
            rows.append((v_set, i_mean, i_stdev, n_samples))
            print(f"  V={v_set:8.3f} V   I={i_mean:+.6e} A   stdev={i_stdev:.2e} A")
    except KeyboardInterrupt:
        print("\nSweep interrupted by user.")
    finally:
        k.output_off()
        k.close()

    if not rows:
        print("No data collected.")
        return None

    out_path = Path(out_csv)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    with out_path.open("w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["voltage_V", "current_mean_A", "current_stdev_A", "n_samples"])
        w.writerows(rows)

    print(f"\nSaved {len(rows)} points to: {out_csv}")
    return out_csv


if __name__ == "__main__":
    print("=" * 60)
    print("Keithley 2400 Library — Quick Test")
    print("=" * 60)

    ports = find_serial_port()
    if not ports:
        print("No serial ports found. Is the USB-serial adapter plugged in?")
    else:
        k = Keithley2400(resource=ports[0])
        print(f"IDN: {k.idn()}")

        k.set_source_v(0.0)
        k.set_current_compliance(0.001)
        k.configure_measure_vi()
        k.output_on()

        v, i = k.read_v_i()
        print(f"Reading: V={v:.6f} V, I={i:.6e} A")

        k.output_off()
        k.close()
        print("Done.")
