import logging
import time
from datetime import datetime
from os import PathLike
from os.path import exists

import h5py
import numpy as np

from naludaq.board import Board
from naludaq.tools.dac_sweep.dac_sweep_controller import DACSweepController
from naludaq.tools.lookup_table import LookupTable

LOGGER = logging.getLogger(__name__)


class LookupTableGenerator:
    """
    Class for generating/continuing a lookup table for ADC2mV

    Example:
        LookupTableGenerator.initialize_lookuptable('file', board)
        with LookupTable('file', 'r+') as lookup_table:
            ltg = LookupTableGenerator(lookup_table, board)
            ltg.dac_sweep(
                min_counts=1000,
                max_counts=3000,
                step_size=1000,
                events_per_datapt=1,
                num_warmup_events=1
                )
            ltg.interpolate_data()
    """

    def __init__(
        self,
        lookup_table: LookupTable,
        board: Board,
    ):
        self.lookup_table = lookup_table
        self.board = board
        self._dac_sweep_ctrl = None

    def dac_sweep(
        self,
        min_counts: int = 0,
        max_counts: int = None,
        step_size: int = 50,
        events_per_datapt: int = 3,
        num_warmup_events: int = 3,
        channels: list = None,
        mode: str = "overwrite",
    ):
        """
        Perform a DAC sweep to collect data for the lookup table, follows
        np.linspace when determining range

        Args:
            min_counts (int): Min count to start the sweep
            max_counts (int): Max count that the sweep will end at
            step_size (int): Step size to increment by
            events_per_datapt (int): Number of events to capture per DAC Value
            num_warmup_events (int): Number of events to capture/discard before capturing data
            mode (str): 'overwrite' or 'average' existing data points
        """
        validate_table_params_or_raise(self.lookup_table, self.board)

        if max_counts == None:
            max_counts = self.board.params["ext_dac"]["max_counts"]
        if channels is None:
            channels = list(range(self.board.channels))
        self._validate_sweep_params(step_size, max_counts, min_counts, mode)

        self._dac_sweep_ctrl = DACSweepController(self.board)
        self._dac_sweep_ctrl._backup_settings()
        try:
            for dac_val in range(min_counts, max_counts + 1, step_size):
                LOGGER.info("Running dac value %d", dac_val)
                dac_val = int(dac_val)
                self.run_dac_value(
                    dac_val, channels, events_per_datapt, num_warmup_events, mode
                )
        finally:
            self._dac_sweep_ctrl._restore_settings()

    def run_dac_value(
        self,
        dac_val,
        channels,
        events_per_datapt: int = 3,
        num_warmup_events: int = 3,
        mode: str = "overwrite",
    ):
        """Runs a dac capture at the given dac value and stores the
        data into the captured_data table then updates the lookup table.

        Args:
            dac_data (list): entry of data captured using the dac sweep controller
            dac_val (int): dac value used to capture the data
            channels (list): channels to write
            events_per_datapt (int): Number of events to capture for pedestals
            num_warmup_events (int): Number of warmup events for pedestals
            mode (str): 'overwrite' or 'average' how to handle repeat data
        """
        if self._dac_sweep_ctrl is None:
            self._dac_sweep_ctrl = DACSweepController(self.board)

        dac_pedestals = self._dac_sweep_ctrl._run_dac_value(
            dac_val=dac_val,
            channels=channels,
            num_captures=events_per_datapt,
            num_warmup_events=num_warmup_events,
        )
        dac_data = dac_pedestals["data"]
        LOGGER.debug("Dac value: %d total events: %d", dac_val, len(dac_data))
        self.write_dac_data(dac_data, dac_val, channels, mode)

    def write_dac_data(self, dac_data, dac_val, channels, mode: str = "overwrite"):
        """Writes an entry of the data generated by the dac sweep controller
        to captured_data table, converting dac_val to mV. Updates the times
        captured table accordingly, and updates the lookup table to the
        captured_data table.

        Args:
            dac_data (list): pedestals data captured using the dac sweep controller
            dac_val (int): dac value used to capture the data
            channels (list): channels to write
            mode (str): 'overwrite' or 'average' how to handle repeat data
        """
        max_mv = self.board.params["ext_dac"]["max_mv"]
        mV = max_mv * dac_val / self.board.params["ext_dac"]["max_counts"]
        LOGGER.info("Writing data to file for Dac val: %d", dac_val)
        for channel in channels:
            peds = np.array(dac_data[channel]).flatten()
            for sample_num, adc_count in enumerate(peds):
                capture_idx = (int(round(adc_count)), channel, sample_num)
                self.write_captured_data(capture_idx, mV, mode)
        # Write captured data to lookup table dataset
        self.update_lookup_table()

    def write_captured_data(self, capture_idx, mV: int, mode: str = "overwrite"):
        """Writes captured data to captured_data table and
        updates the times captured table accordingly

        Args:
            capture_idx (tuple or arr): Index of captured data table to write
            to in form: (adc_count (int), channel (int), sample_num (int))
            mV (float): mV value to set
            mode (str): 'overwrite' or 'average' how to handle repeat data
        """
        if not mode in ["overwrite", "average"]:
            raise ValueError("Invalid mode, 'overwrite' or 'average' is accepted")
        try:
            if mode == "overwrite" or self.times_captured[capture_idx] == 0:
                self._overwrite_data(capture_idx, mV)
            if mode == "average":
                self._average_data(capture_idx, mV)
        except KeyboardInterrupt:
            self.lookup_table.captured_data[capture_idx] = np.nan
            self.lookup_table.times_captured[capture_idx] = 0
            raise

    def write_event(self, event: dict, mV: float, mode: str = "overwrite"):
        """Writes captured events to captured_data table and
        updates the times captured table accordingly

        Args:
            event (dict): Parsed event
            mV (float): mV value to set
            mode (str): 'overwrite' or 'average' how to handle repeat data
        """
        if not all(x in event.keys() for x in ["data", "window_labels"]):
            raise KeyError("Event does not contain 'data' and/or 'window_labels'")
        samples = self.board.params["samples"]
        for channel in range(self.board.channels):
            for idx, window in enumerate(event["window_labels"][channel]):
                for relative_sample in range(samples):
                    sample_num = window * samples + relative_sample
                    adc_count = int(
                        event["data"][channel][relative_sample + idx * samples]
                    )
                    capture_idx = (adc_count, channel, sample_num)
                    self.write_captured_data(capture_idx, mV, mode)
        self.update_lookup_table()

    def interpolate_data(self):
        """Fills empty entries in from captured data, and stores it
        to lookup table with linear interpolation. Left and right
        ends of the dac entries are filled with min and max values
        for dac value respectively.
        """
        channels = self.board.params["channels"]
        for channel in range(channels):
            self.interpolate_channel_data(channel)

    def interpolate_channel_data(self, channel: int):
        """Fills empty entries in from captured data for a channel, and stores
        it to lookup table with linear interpolation.
        Left and right ends of the dac entries are filled with
        min and max values for dac value respectively.
        """
        validate_table_params_or_raise(self.lookup_table, self.board)
        num_samples = self.board.params["samples"] * self.board.params["windows"]
        data_width = self.board.params["ext_dac"]["max_counts"] + 1
        channel_slice = np.array(self.lookup_table[:, channel, :])
        channel_slice[0][:] = np.nanmin(channel_slice, axis=0)
        channel_slice[-1][:] = np.nanmax(channel_slice, axis=0)
        # Get indices of captured data
        data_indices = np.argwhere(
            ~np.isnan(channel_slice.transpose().ravel())
        ).flatten()
        interpolated_data = np.interp(
            range(data_width * num_samples),
            data_indices,
            channel_slice.transpose().ravel()[data_indices],
        )
        self.lookup_table[:, channel, :] = interpolated_data.reshape(
            (num_samples, data_width)
        ).transpose()

    def update_lookup_table(self):
        """Updates lookup table to captured_data table"""
        self.lookup_table.lookup[...] = self.lookup_table.captured_data

    def _overwrite_data(self, capture_idx, mV: int):
        """Overwrites captured data to captured_data table and
        updates the times captured table accordingly to 1
        """
        self.lookup_table.captured_data[capture_idx] = mV
        self.lookup_table.times_captured[capture_idx] = 1

    def _average_data(self, capture_idx, mV: int):
        """Averages captured data to captured_data table and
        updates the times captured table accordingly
        """
        cur_val = self.lookup_table.captured_data[capture_idx]
        times_cap = self.lookup_table.times_captured[capture_idx]
        avg = (cur_val * times_cap + mV) / (times_cap + 1)
        self.lookup_table.captured_data[capture_idx] = avg
        self.lookup_table.times_captured[capture_idx] += 1

    def _process_avg_ped(self, dac_data, channel):
        """Converts dac_data to pedestals"""
        ped_list = list()
        for capture_num in range(len(dac_data[0])):
            ped_data = [event[capture_num]["data"][channel] for event in dac_data]
            ped_data = np.array(ped_data).flatten().astype(int)
            ped_list.append(ped_data)
        avg_ped = np.array(ped_list).mean(axis=0).astype(int)

        return avg_ped

    @staticmethod
    def initialize_lookuptable(file_name: str, board: Board):
        """Writes a file and initializes three tables for the given board with two formats:
        lookup_table, captured_data:
        [ADC Counts][Channels][Windows * Samples] = mV (float)
        times_captured:
        [ADC Counts][Channels][Windows * Samples] = Num Captures (int)
        Initialized to np.nan for mV Values, and 0 for Num Captures
        Does not create a new lookup table if the file already exists

            Args:
                filename (str): Filename/path to store lookup table
                board (Board): Board object to initialize the lookup table's dimensions
        """

        if exists(file_name):
            raise FileExistsError("Lookup table already exists")

        with h5py.File(file_name, "w") as h5_file:
            h5_file.attrs.create("created_at", time.time())
            settings = {
                "model": board.params["model"],
                "params": board.params,
                "registers": board.registers,
            }
            settings_group = h5_file.create_group("settings")
            _store_attributes(settings_group, settings)
            windows = board.params["windows"]
            samples = board.params["samples"]
            data_width = 2 ** board.params["resolution"]
            channels = board.params["channels"]

            h5_file.create_dataset(
                name="captured_data",
                shape=(data_width, channels, windows * samples),
                compression="gzip",
                compression_opts=3,
                fillvalue=np.nan,
                dtype=np.float32,
            )
            h5_file.create_dataset(
                name="lookup_table",
                shape=(data_width, channels, windows * samples),
                compression="gzip",
                compression_opts=3,
                fillvalue=np.nan,
                dtype=np.float32,
            )
            h5_file.create_dataset(
                name="times_captured",
                shape=(data_width, channels, windows * samples),
                compression="gzip",
                compression_opts=3,
                fillvalue=0,
                dtype=np.int16,
            )

    def _validate_sweep_params(self, step_size, max_counts, min_counts, mode):
        if step_size < 1:
            raise ValueError("step_size must be > 0")
        if max_counts < min_counts:
            raise ValueError("Invalid min/max counts")
        if step_size > max_counts:
            raise ValueError("Invalid step_size value")
        if not mode in ["overwrite", "average"]:
            raise ValueError("Invalid mode, 'overwrite' or 'average' is accepted")


def validate_table_params_or_raise(lookup_table: LookupTable, board: Board) -> bool:
    """Validates the table dimensions from the lookuptable file, to the
    given board object
    """

    keys = [
        "lookup_table",
        "times_captured",
        "captured_data",
        "settings",
    ]
    h5_file = lookup_table._h5_file
    for key in keys:
        if key not in h5_file.keys():
            raise KeyError(f"Dataset doesn't exist {key}")

    if h5_file["settings"].attrs["model"] != board.model:
        raise ValueError("Given board model does not match the lookup table")
    h5_lookup_table = h5_file["lookup_table"]
    if (board.params["ext_dac"]["max_counts"] + 1) != h5_lookup_table.shape[0]:
        raise ValueError("Data width parameter does not match the given board")
    if board.params["channels"] != h5_lookup_table.shape[1]:
        raise ValueError("Channels parameter does not match the given board")
    if board.params["samples"] * board.params["windows"] != h5_lookup_table.shape[2]:
        raise ValueError(
            "Sample and/or windows parameters do not match the given board"
        )


def _store_attributes(group: h5py.Group, attrs: dict, recursive=True):
    """Adds a dict of attributes to a h5py group.

    Skips entries that have 'None' value.

    Args:
        group (Group): the group to add the attributes to
        attrs (dict): a dict of {name: value}
        recursive (bool): whether to create new groups for
            entries with dict values
    """
    for k, v in attrs.items():
        k = str(k)
        if v is None:  # HDF5 doesn't support None value
            continue
        elif isinstance(v, dict) and recursive:
            new_group = group.create_group(k, track_order=True)
            _store_attributes(new_group, v)
        else:
            v = _convert_for_storing(v)
            group.attrs.create(k, v)


def _convert_for_storing(value):
    """Converts a value to a type that can be stored in HDF5

    datetime -> float
    PathLike -> str
    bytes -> np.void

    Args:
        value: the value to convert

    Returns:
        A tuple of (converted value, new dtype)
    """
    if isinstance(value, datetime):
        value = value.timestamp()
    elif isinstance(value, PathLike):
        value = str(value)
    elif isinstance(value, bytes):
        value = np.void(value)
    return value
