from pathlib import Path
import numpy as np
from .RecordingMetadata import RecordingMetadata
from .errors import BreesyError
from .load import _calculate_sample_rate_from_time_column, _detect_time_column, _find_keys_in_data, _leave_only_eeg_channels
from .breesy_scipy import _loadmat
from .recording import Recording
from .log import logger
from .type_hints import enforce_type_hints


@enforce_type_hints
def load_mat(filename: str, sample_rate: int | float | None = None) -> Recording:
    """Load a MATLAB .mat file.

    :param filename: Path to the .mat file
    :param sample_rate: Sampling rate in Hz

    :return: Recording object containing the loaded data
    """
    mat_data = _loadmat(filename)

    meaningful_keys = _mat_extract_meaningful_keys(mat_data)

    data_key = None

    if len(meaningful_keys) == 1:
        data_key = meaningful_keys[0]
    else:
        # Find the largest array that looks like EEG data
        largest_size = 0
        for key in meaningful_keys:
            value = mat_data[key]
            if isinstance(value, np.ndarray) and value.size > largest_size and value.ndim >= 2:
                largest_size = value.size
                data_key = key

    data = mat_data[data_key]  # TODO: there can be 'O' types, as AI previously suggested?
    if not isinstance(data, np.ndarray):
        raise BreesyError(
            f"Problem reading data which is of type {type(data)}. It may be that Breesy failed to read the file correctly."  # TODO: test different files
        )
    if data.ndim != 2:
        raise BreesyError(
            "Found data has an incorrect number of dimensions. "
            "Expected data.ndim to be 2, but got "
            f"{data.ndim} dimensions instead.",
            "The data may be not EEG data, "
            "there are several records joined into one file, "
            "or Breesy did not found the correct data array in the file."
        )

    if data.shape[0] > data.shape[1]:
        logger.warning("The first data dimension seems to be the samples, because data contains more rows than columns. Dimensions will be switched places.")
        data = data.T
    n_channels, n_samples = data.shape
    default_channel_names = [f'Channel_{i + 1}' for i in range(data.shape[0])]

    # find or create channel names
    channel_name_keys = _find_keys_in_data(meaningful_keys, ['ch', 'ch name', 'ch names', 'channel names', 'channels', 'names', '10-20', '10-10', 'pos', 'positions'])
    if len(channel_name_keys) == 0:
        channel_names = default_channel_names
        channel_name_key = None
    elif len(channel_name_keys) > 1:
        logger.warning(
            "found several keys in data which may contain channel names: "
            f"{channel_name_keys}. "
            "Using the first one.")
        channel_name_key = channel_name_keys[0]
    else:
        channel_name_key = channel_name_keys[0]

    if channel_name_key:
        found_names = mat_data[channel_name_key]
        if not isinstance(found_names, (list, tuple, np.ndarray)):
            logger.warning(
                "found channel name information does not appear to be a list: "
                f"{found_names}. "
                "Default channel names will be used instead.")
            channel_names = default_channel_names
        elif (n_channels != len(found_names)) and (n_samples != len(found_names)):
            logger.warning(
                f"the number of channel names ({len(found_names)}) provided does not match the "
                f"number of channels ({n_channels}) in the EEG data. "
                "Default channel names will be used instead."
            )
            channel_names = default_channel_names
        elif (n_channels != len(channel_names)) and (n_samples == len(channel_names)):
            logger.warning('data rows and columns will be switched places according to the number of found channel names.')
            data = data.T
            n_channels, n_samples = data.shape
            channel_names = found_names
        else:
            channel_names = found_names

    # try to find it in keys
    extracted_sample_rate = _detect_sample_rate_in_mat_keys(mat_data, meaningful_keys)
    if extracted_sample_rate is None:
        # else in time column
        time_column_idx = _detect_time_column(data)
        if time_column_idx is not None:
            extracted_sample_rate = _calculate_sample_rate_from_time_column(data[time_column_idx])
            logger.info(f"Detected time column at index {time_column_idx}. Extracted sample rate: {extracted_sample_rate} Hz")

    # Determine final sample rate
    if sample_rate is not None:
        if extracted_sample_rate is not None and abs(sample_rate - extracted_sample_rate) > 0.01:
            logger.warning(f"Warning: Provided sample_rate ({sample_rate} Hz) differs from extracted sample rate ({extracted_sample_rate} Hz). Using the provided sample_rate.")
        sample_rate = float(sample_rate)
    elif extracted_sample_rate is not None:
        sample_rate = extracted_sample_rate
    else:
        raise BreesyError(
            "No sample rate provided and could not detect sample rate from data.",
            "Please provide the sample rate manually by passing the 'sample_rate' parameter. "
            "Example: load_mat(filename, sample_rate=250)"
        )

    data, channel_names = _leave_only_eeg_channels(data, channel_names)
    # channel names may not start from 1 after this, but maybe we should leave it like this

    metadata = RecordingMetadata(
        name=Path(filename).stem,
        file_path=Path(filename),
        description=f"MATLAB data loaded from {filename}"
    )

    return Recording(
        data=data,
        channel_names=channel_names,
        sample_rate=sample_rate,
        events=None,  # No events for now - could be added as a separate function
        metadata=metadata
    )


def _mat_extract_meaningful_keys(mat_data: dict) -> list[str]:
    keys = [k for k in mat_data.keys() if not k.startswith('__')]
    if not keys:
        raise BreesyError(
            "No valid data found. The MATLAB file appears to be empty or contains only metadata."
        )
    elif len(keys) == 1:
        logger.warning(
            "data file likely contains only the recording and "
            "no useful metadata. Please find the sample rate and channel names "
            "in your data description."
        )  # TODO: fast instructions or link to documentation
    return keys


def _detect_sample_rate_in_mat_keys(mat_data: dict, keys: list[str]) -> float | None:
    sample_rate_keys = _find_keys_in_data(keys, ['sample rate', 'sp', 'sr', 'hz', 'sampling rate', 'sfreq', 'fs', 'srate'])
    if len(sample_rate_keys) == 0:
        return None

    sample_rate_key = sample_rate_keys[0]
    if len(sample_rate_keys) > 1:
        logger.warning(
            "found several keys in data which may contain sample rate:"
            f"\n{sample_rate_keys}\n"
            "Using the first one.")

    sr_value = mat_data[sample_rate_key]
    try:
        if isinstance(sr_value, (int, float, str)):
            return float(sr_value)
        elif isinstance(sr_value, np.ndarray) and sr_value.size == 1:
            return float(sr_value.item())
        elif isinstance(sr_value, (list, tuple)) and len(sr_value) == 1:
            return float(sr_value[0])
        else:
            raise ValueError
    except ValueError:
        logger.warning(
            f"Information found in column '{sample_rate_key}' "
            f"could not be read automatically; column contents are:\n{sr_value}"
        )
        return None