Ad20p 1230e

class CS_IO404_Based_AD20P_1230E(Pump, MODBUSInterface, HasSetupFunctionsMixin):
    """AD20P-1230E pump controlled via a CS-IO404 Modbus relay.

    Converts dosage requests (ml) to timed relay activations using linear
    calibration parameters. The relay module is addressed via MODBUS RTU.

    Attributes:
        relay_address: MODBUS device address of the CS-IO404 module (1-247).
        coil_number: Output coil index on the relay module (0-3).
        calibration_save_function: Callback to persist new calibration data.
    """
    def __init__(
        self,
        bus: MODBUS,
        db_save_function: Callable[[Datapoint | list[Datapoint]], None],
        calibration_parameters: dict[str, Any],
        calibration_save_function: Callable,
        relay_address: int,
        coil_number: int
    ):
        """Initialize the pump with relay addressing and calibration.

        Args:
            bus: MODBUS bus instance.
            db_save_function: Async function to persist watering events.
            calibration_parameters: Dict with 'slope' and 'offset' keys.
            calibration_save_function: Callable(slope, offset) to persist
                new calibration values after the calibration procedure.
            relay_address: MODBUS address of the CS-IO404 (1-247).
            coil_number: Relay output index (0-3).

        Raises:
            ValueError: If coil_number or relay_address is out of range.
        """
        if not (0 <= coil_number < 4):
            raise ValueError(f"Invalid coil number {coil_number} for pump. Coil number must be between 0 and 3 inclusive, corresponding to the 4 outputs of the CS-IO404 relay module.")
        if not (1 <= relay_address <= 247):
            raise ValueError(f"Invalid relay address {relay_address} for pump. MODBUS relay addresses must be between 1 and 247 inclusive.")
        if relay_address == 1:
            _logger.warning(f"A pump is configured with relay address 1, which is typically reserved for the controller itself. If your relay module is configured to use address 1, consider changing it to avoid potential conflicts on the MODBUS network.")
        self.bus = bus
        self.db_save_function = db_save_function
        self.calibration_parameters = calibration_parameters
        self.calibration_save_function = calibration_save_function
        self.relay_address = relay_address
        self.coil_number = coil_number

    def doseage_to_time(self, dosage: int) -> float:
        """Convert a dosage in ml to pumping duration in seconds.

        Uses the linear model: time = slope * dosage + offset.

        Args:
            dosage: Desired water amount in milliliters.

        Returns:
            Pumping duration in seconds.
        """
        slope = self.calibration_parameters["slope"]
        offset = self.calibration_parameters["offset"]
        return slope * dosage + offset

    async def pumping_callback(self, dosage: int):
        """Pump the specified dosage and record the watering event.

        Args:
            dosage: Amount to pump in milliliters.
        """
        pump_time = self.doseage_to_time(dosage)
        _logger.debug(f"Starting pump {self.relay_address}-{self.coil_number} for {pump_time} seconds, corresponding to a dosage of {dosage} ml")
        await self._toggle_pump_on_for_duration(pump_time)
        self.db_save_function(WateringEvent(dosage=dosage))

    async def _toggle_pump_on_for_duration(self, duration: float):
        """Activate the relay for a precise duration, then deactivate.

        Runs the blocking relay toggle in a worker thread to avoid
        blocking the async event loop during the sleep.

        Args:
            duration: Time in seconds to keep the pump running.
        """
        def _blocking_pump():
            self.bus.write_coil(
                address=self.coil_number,
                value=True,
                device_id=self.relay_address
            )
            sleep(duration)
            self.bus.write_coil(
                address=self.coil_number,
                value=False,
                device_id=self.relay_address
            )
        await anyio.to_thread.run_sync(_blocking_pump)
        _logger.debug(f"Stopped pump {self.relay_address}-{self.coil_number}")

    async def calibrate(self):
        """Interactive calibration procedure for determining pump flow rate.

        Guides the user through multiple pump cycles at different durations,
        collecting volume measurements to compute a linear fit (slope and
        offset) for the dosage-to-time relationship. Results are saved via
        the calibration_save_function.
        """
        import numpy
        import datetime
        clear_screen()
        print("You will now be guided through calibration of the pump step by step.")
        print("For a robust calibration, the testing setup must be as close as possible to")
        print("the real system setup.")
        print("To that end, first ensure that the following is true:")
        print("")
        print("  - The water tank is filled with water.")
        print("")
        print("  - The pump is submerged in the tank.")
        print("")
        print("  - The pump is connected to the controller.")
        print(f"    (The pump should be connected to output DO{self.coil_number + 1} of the IO404 module, which")
        print(f"    should in turn have the MODBUS address {self.relay_address} and be connected to the RPI.)")
        print("")
        print("  - The outlet tube from the pump is positioned such that it pumps into a")
        print("    measuring vessel at the same height relative to the water tank as it")
        print("    would be were it pumping into the pot of its respective plant.")
        print("")
        print("")
        print("When you've made sure that the above is true, then press enter to continue.")
        print("(The calibration procedure can be cancelled whenever you are prompted for")
        print("input by responding with 'cancel' and pressing enter.)")
        response = input()
        clear_screen()
        match response:
            case 'stop' | 'cancel' | 'quit':
                print("Calibration Aborted.")
                return

        while True:
            print("Measure the height from the water surface in the tank to the outlet of the pump.")
            print("Enter this height in centimeters (rounding up) and press enter.")
            response = input()
            match response:
                case 'stop' | 'cancel' | 'quit':
                    print("Calibration Aborted.")
                    return
            try:
                response = int(response)
                break
            except ValueError:
                clear_screen()
                print(f"The given input '{response}' wasn't a whole number!")

        # Time to wait on water in watering tube to recede back into watertank
        rest_time = response / 10 if response > 0 else 0
        clear_screen()

        print("The controller will now go through multiple cycles of turning on the pump for")
        print("different amounts of time. After each cycle you'll have to measure and record")
        print("how much water was actually pumped.")
        print("")
        print("Before continuing ensure that you have a measuring cup with a capacity of")
        print("at least 1 liter, and a way to measure how much water is in that cup down to the")
        print("nearest 10 milliliters.")
        print("")
        print("(Press enter to continue)")
        response = input()
        match response:
            case 'stop' | 'cancel' | 'quit':
                print("Calibration Aborted.")
                return
        clear_screen()

        pumped_durations = []
        pumped_amounts = []
        for pumping_time in [time / 2 for time in range(2, 17)]:
            print("Place the measuring cup under the outlet of the pump.")
            print("(Press enter to continue)")
            response = input()
            clear_screen()
            match response:
                case 'stop' | 'cancel' | 'quit':
                    print("Calibration Aborted.")
                    return
            print(f"Pumping... ({pumping_time} seconds)")
            await self._toggle_pump_on_for_duration(pumping_time)
            clear_screen()
            print("Pumping... Done!")
            print("")
            rest_period_start_time = datetime.datetime.now()
            rest_period_end_time = rest_period_start_time + datetime.timedelta(seconds=rest_time)
            while True:
                print("How many milliliters of water was pumped into the measuring cup?.")
                response = input()
                clear_screen()
                match response:
                    case 'stop' | 'cancel' | 'quit':
                        print("Calibration Aborted.")
                        return
                try:
                    response = int(response)
                    if response >= 0:
                        break
                    print("The given input was negative! If the pump is sucking in water, then you've installed it wrong.")
                except Exception:
                    print(f"The given input '{response}' wasn't a whole number!")
            if response > 0:
                _logger.debug(f"Recorded pumped amount: {response} ml")
                _logger.debug(f"Recorded pumping time: {pumping_time} seconds")
                pumped_amounts.append(response)
                pumped_durations.append(pumping_time)
            else:
                _logger.debug("Calibrater recorded a pumped amount of 0ml, discarding this data point.")

            print("Empty the measuring cup.")
            print("(Press enter to continue)")
            response = input()
            clear_screen()
            match response:
                case 'stop' | 'cancel' | 'quit':
                    print("Calibration Aborted.")
                    return
            if (
                remaining_rest_time := (
                    rest_period_end_time - datetime.datetime.now()
                ).total_seconds()
            ) > 0:
                print(f"Waiting for {remaining_rest_time} seconds to let the water in the tube recede back into the tank...")
                await anyio.sleep(remaining_rest_time)
                clear_screen()

        print("Calculating calibration...")

        slope, offset = numpy.polyfit(pumped_amounts, pumped_durations, 1)
        old_slope = self.calibration_parameters["slope"]
        old_offset = self.calibration_parameters["offset"]
        self.calibration_parameters["slope"] = slope
        self.calibration_parameters["offset"] = offset
        print("Calibration complete!")
        print("")
        print(f"Old slope: {old_slope}, Old offset: {old_offset}")
        print(f"New slope: {slope}, New offset: {offset}")
        _logger.debug(f"Pump {self.relay_address}-{self.coil_number} calibrated with slope {slope} and offset {offset}. Old slope was {old_slope} and old offset was {old_offset}. Calibration data points were: {list(zip(pumped_amounts, pumped_durations))}")
        self.calibration_save_function(slope, offset)
        print("")
        print("The pump is now calibrated. You can test the pump with its new calibration by running the 'test' setup function.")
        print("(Press enter to continue)")
        input()

    async def test_pump(self):
        """Interactive test: pump a user-specified dosage to verify calibration."""
        print("How many ml do you want to pump for the test?")
        while True:
            try:
                dosage = int(input("Enter dosage in ml: "))
                break
            except ValueError:
                print("Invalid input. Please enter an integer value for the dosage.")
        pump_time = self.doseage_to_time(dosage)
        _logger.info(f"Starting pump {self.relay_address}-{self.coil_number} for {pump_time} seconds, corresponding to a dosage of {dosage} ml")
        await self._toggle_pump_on_for_duration(pump_time)

    def setup_functions(self) -> dict[str, dict[str, any]]:
        """Return available setup actions for this pump.

        Returns:
            Dict with 'calibrate' and 'test' entries.
        """
        return {
            "calibrate": {
                "description": "Run the calibration procedure for the pump.",
                "function": self.calibrate
            },
            "test": {
                "description": "Test the pump by running it for a short time.",
                "function": self.test_pump
            }
        }