"""
This module contains the definition for the high-level Rigol1000z driver.
"""

import numpy as _np
import tqdm as _tqdm
import pyvisa as _visa
from time import sleep
from .commands import *
from typing import List


class Rigol1000z(Rigol1000zCommandMenu):
    """
    The Rigol DS1000z series oscilloscope driver.
    """

    def __init__(self, visa_resource: _visa.Resource):
        # Instantiate The scope as a visa command menu
        super().__init__(visa_resource)

        # Initialize IEEE device identifier command in order to determine the model
        brand, model, serial_number, software_version, *add_args = self._idn_cache.split(",")

        # Ensure a valid model is being used
        assert brand == "RIGOL TECHNOLOGIES"
        assert model in {
            ScopeModel.DS1104Z_S_Plus, ScopeModel.DS1104Z_Plus, ScopeModel.DS1104Z,  # 100MHz models
            ScopeModel.DS1074Z_S_Plus, ScopeModel.DS1074Z_Plus,  # 70MHz models
            ScopeModel.DS1054Z  # 50MHz models
        }

        # Define Channels 1-4
        self.channel_list: List[Channel] = [Channel(self.visa_resource, c) for c in range(1, 5)]
        """
        A four-item list of commands.Channel objects
        """

        # acquire must be able to count enabled channels
        self.acquire = Acquire(self.visa_resource, self.channel_list)
        """
        Hierarchy commands.Acquire object
        """

        self.calibrate = Calibrate(self.visa_resource)
        """
        Hierarchy commands.Calibrate object
        """

        self.cursor = Cursor(self.visa_resource)  # NC
        self.decoder = Decoder(self.visa_resource)  # NC

        self.display = Display(self.visa_resource)
        """
        Hierarchy commands.Display object
        """

        self.event_tables = [EventTable(self.visa_resource, et + 1) for et in range(2)]
        """
        A two-item list of commands.EventTable objects used to detect decode events.
        """

        self.function = Function(self.visa_resource)  # NC

        self.ieee488 = IEEE488(self.visa_resource)
        """
        Hierarchy commands.IEEE488 object
        """

        if self.has_digital:
            self.la = LA(self.visa_resource)  # NC
        self.lan = LAN(self.visa_resource)  # NC
        self.math = Math(self.visa_resource)  # NC
        self.mask = Mask(self.visa_resource)  # NC

        self.measure = Measure(self.visa_resource)
        """
        Hierarchy commands.Measure object
        """

        self.reference = Reference(self.visa_resource)  # NC

        if model in {ScopeModel.DS1104Z_S_Plus, ScopeModel.DS1074Z_S_Plus}:  # Only for "S" models
            self.source = Source(self.visa_resource)  # NC

        self.storage = Storage(self.visa_resource)  # NC
        self.system = System(self.visa_resource)  # NC
        self.trace = Trace(self.visa_resource)  # NC

        self.timebase = Timebase(self.visa_resource)
        """
        Hierarchy commands.Timebase object
        """

        self.trigger = Trigger(self.visa_resource)  # NC

        self.waveform = Waveform(self.visa_resource)
        """
        Hierarchy commands.Waveform object
        """

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.visa_resource.close()
        return False

    def __del__(self):
        self.visa_resource.close()

    def __getitem__(self, i) -> Channel:
        """
        Channels 1 through 4 (or 2 depending on the oscilloscope model) are accessed
        using `[channel_number]`.  e.g. osc[2] for channel 2.  Channel 1 corresponds
        to index 1 (not 0).

        :param i: Channel number to retrieve
        :return:
        """
        # assert i in {c.channel for c in self._channels}
        assert 1 <= i <= 4, 'Not a valid channel.'
        return self.channel_list[i - 1]

    def __len__(self):
        return len(self.channel_list)

    def autoscale(self):
        print("Autoscaling can take several seconds to complete")
        old_timeout = self.visa_resource.timeout
        self.visa_resource.timeout = None
        self.visa_write(':aut')
        wait_for_resp = self.ieee488.operation_complete  # Wait for queued response before moving onto next command
        self.visa_resource.timeout = old_timeout
        print("Autoscaling complete")

    def clear(self):
        self.visa_write(':clear')

    def run(self):
        self.visa_write(':run')

    def stop(self):
        self.visa_write(':stop')

    def set_single_shot(self):
        self.visa_write(':sing')

    def force(self):
        self.visa_write(':tfor')

    def get_channels_enabled(self):
        return [c.enabled() for c in self.channel_list]

    # todo: make this more closely knit with the library
    def get_screenshot(self, filename=None, img_format: str | None = None):
        """
        Downloads a screenshot from the oscilloscope.

        Args:
            filename (str): The name of the image file.  The appropriate
                extension should be included (i.e. jpg, png, bmp or tif).
            img_format (str, optional): Override the on-device image format.
                Supported values are: 'jpeg', 'png', 'bmp8', 'bmp24', 'tiff'.
        """
        # The format image that should be downloaded.
        # Options are 'jpeg', 'png', 'bmp8', 'bmp24' and 'tiff'.
        # Default is 'png'.
        if img_format is None:
            if isinstance(filename, str) and "." in filename:
                img_format = filename.rsplit(".", 1)[-1].lower()
            else:
                img_format = "png"

        # Common alias.
        if img_format == "jpg":
            img_format = "jpeg"

        assert img_format in ("jpeg", "png", "bmp8", "bmp24", "tiff")

        sleep(0.5)  # Wait for display to update

        # Screenshot transfer can take a couple seconds; use a bounded timeout
        # to avoid blocking forever if the instrument doesn't respond.
        old_timeout = self.visa_resource.timeout
        self.visa_resource.timeout = 8000 if old_timeout is None else max(int(old_timeout), 8000)

        try:
            # Query screenshot bytes.
            self.visa_write(f":disp:data? on,off,{img_format}")
            raw = self.visa_resource.read_raw()
        finally:
            self.visa_resource.timeout = old_timeout

        # Extract payload from an IEEE488.2 definite-length block if present.
        raw_bytes: bytes = bytes(raw)
        if raw_bytes[:1] == b"#" and len(raw_bytes) >= 2:
            try:
                digits = int(raw_bytes[1:2].decode("ascii"))
            except Exception:
                digits = -1

            if digits == 0:
                raw_img = raw_bytes[2:]
            elif digits > 0 and len(raw_bytes) >= 2 + digits:
                try:
                    size = int(raw_bytes[2:2 + digits].decode("ascii"))
                except Exception:
                    size = -1

                start = 2 + digits
                end = start + size
                raw_img = raw_bytes[start:end] if size >= 0 and end <= len(raw_bytes) else raw_bytes[start:]
            else:
                raw_img = raw_bytes
        else:
            raw_img = raw_bytes

        if filename:
            try:
                os.remove(filename)
            except OSError:
                pass
            with open(filename, 'wb') as fs:
                fs.write(raw_img)

        return raw_img

    def get_screenshot_safe(self, filename, img_format: str = "png", timeout_ms: int = 8000):
        from .imaging import capture_screenshot_safe

        return capture_screenshot_safe(self, filename, fmt=img_format, timeout_ms=timeout_ms)

    def get_labeled_screenshot(
        self,
        filename,
        labels,
        labeled_filename=None,
        img_format: str = "png",
        timeout_ms: int = 8000,
    ):
        from .imaging import capture_labeled_screenshot

        return capture_labeled_screenshot(
            self,
            filename,
            labels,
            labeled_output_path=labeled_filename,
            fmt=img_format,
            timeout_ms=timeout_ms,
        )

    def get_data(self, mode=EWaveformMode.Normal, filename=None):
        """
        Download the captured voltage points from the oscilloscope.

        Args:
            mode (str): 'norm' if only the points on the screen should be
                downloaded, and 'raw' if all the points the ADC has captured
                should be downloaded.  Default is 'norm'.
            filename (None, str): Filename the data should be saved to.  Default
                is `None`; the data is not saved to a file.

        Returns:
            2-tuple: A tuple of two lists.  The first list is the time values
                and the second list is the voltage values.

        """

        # Stop scope to capture waveform state
        self.stop()

        # Set mode
        assert mode in {EWaveformMode.Normal, EWaveformMode.Raw}
        self.waveform.mode = mode

        # Set transmission format
        self.waveform.read_format = EWaveformReadFormat.Byte

        # Create data structures to populate
        time_series = None
        all_channel_data = []

        # Iterate over possible channels
        for c in range(1, 5):

            # Capture the waveform if the channel is enabled
            if self[c].enabled:

                self.waveform.source = self[c].name

                # retrieve the data preable
                info: PreambleContext = self.waveform.data_premable

                # Generate the time series for the data
                time_series = _np.arange(0, info.points * info.x_increment, info.x_increment)

                max_num_pts: int = 250000
                num_blocks: int = info.points // max_num_pts
                last_block_pts: int = info.points % max_num_pts

                datas = []
                # print(f"Data being gathered for Ch{}")
                for i in _tqdm.tqdm(range(num_blocks + 1), ncols=60):
                    if i < num_blocks:
                        self.waveform.read_start_point = 1 + i * 250000
                        self.waveform.read_end_point = 250000 * (i + 1)
                    else:
                        if last_block_pts:
                            self.waveform.read_start_point = 1 + num_blocks * 250000
                            self.waveform.read_end_point = num_blocks * 250000 + last_block_pts
                            sleep(0.2)
                        else:
                            break
                    data = self.visa_ask_raw(':wav:data?', 250000)

                    data = _np.frombuffer(data[11:-1], 'B')  # Last byte marks the end of the message.
                    datas.append(data)

                # Attach each data packet received into a complete data series
                datas = _np.concatenate(datas)

                # Add the data series to the list of data series
                all_channel_data.append(
                    (datas - info.y_origin - info.y_reference) * info.y_increment
                )

        if filename:
            print(f"writing to: {filename}")
            try:
                os.remove(filename)
            except OSError:
                pass
            _np.savetxt(filename, _np.column_stack((time_series, *all_channel_data)), '%.12e', ', ', '\n')

        return time_series, all_channel_data
