import time

import pyvisa


class EloadError(Exception):
    pass


class EloadConnectionError(EloadError):
    pass


class EloadCommunicationError(EloadError):
    pass


class Eload:
    """ITECH IT6000C bi-directional supply, used here as an electronic load (CV mode with +/- current limit)."""

    def __init__(self, resource_name, baud_rate=9600, timeout_ms=3000, retries=3, retry_delay_s=0.2):
        self._retries = retries
        self._retry_delay_s = retry_delay_s
        try:
            self._rm = pyvisa.ResourceManager()
            self._inst = self._rm.open_resource(resource_name)
        except pyvisa.VisaIOError as exc:
            raise EloadConnectionError(f"failed to open {resource_name}") from exc

        self._inst.baud_rate = baud_rate
        self._inst.write_termination = "\n"
        self._inst.read_termination = "\n"
        self._inst.timeout = timeout_ms

    def _write(self, cmd):
        for attempt in range(self._retries):
            try:
                self._inst.write(cmd)
                return
            except pyvisa.VisaIOError as exc:
                if attempt == self._retries - 1:
                    raise EloadCommunicationError(f"write failed: {cmd!r}") from exc
                time.sleep(self._retry_delay_s)

    def _query(self, cmd):
        for attempt in range(self._retries):
            try:
                return self._inst.query(cmd).strip()
            except pyvisa.VisaIOError as exc:
                if attempt == self._retries - 1:
                    raise EloadCommunicationError(f"query failed: {cmd!r}") from exc
                time.sleep(self._retry_delay_s)

    def set_remote(self):
        self._write("SYST:REM")

    def set_local(self):
        self._write("SYST:LOC")

    def set_mode_cv(self):
        self._write("SOUR:FUNC VOLT")

    def set_mode_cc(self):
        self._write("SOUR:FUNC CURR")

    def set_voltage(self, volts):
        self._write(f"VOLT {volts}")

    def set_current(self, amps):
        self._write(f"CURR {amps}")

    def set_current_limit_positive(self, amps):
        self._write(f"CURR:LIM:POS {amps}")

    def set_current_limit_negative(self, amps):
        self._write(f"CURR:LIM:NEG {amps}")

    def output_on(self):
        self._write("OUTP ON")

    def output_off(self):
        self._write("OUTP OFF")

    def is_output_on(self):
        return self._query("OUTP?") == "1"

    def measure_voltage(self):
        return float(self._query("MEAS:VOLT?"))

    def measure_current(self):
        return float(self._query("MEAS:CURR?"))

    def measure_power(self):
        return float(self._query("MEAS:POW?"))

    def disconnect(self):
        self._inst.close()