from pathlib import Path, PosixPath
import numpy as np

from .log import logger
from .recording import Recording
from .Event import Event
from breesy.load_events import find_and_load_event_for_file
from .RecordingMetadata import RecordingMetadata
from .errors import BreesyError, BreesyInternalError, protect_from_lib_error
from .type_hints import enforce_type_hints

# Note: if changing this, then update load_dataset if/else statement as well
SUPPORTED_FILE_EXTENSIONS = ['.csv', '.mat', '.tsv', '.npy', '.eeg', '.vhdr', '.avg', '.seg', '.fdt', '.set', '.edf', '.xdf', '.bdf']


@enforce_type_hints
def load_recording(filename: str | Path | PosixPath, sample_rate: int | float | None = None) -> Recording:
    """Load a recording from a file with a supported extension.

    :param filename: Path to the file to load
    :param sample_rate: Optional sampling rate in Hz

    :return: Recording object containing the loaded data
    """
    filename = str(filename)
    file_extension = Path(filename).suffix
    if not file_extension:
        raise BreesyError("File name should contain an extension",
                          "Provide a file with a supported extension, example: 'my_data.gdf'")

    is_supported_extension = file_extension in SUPPORTED_FILE_EXTENSIONS
    if not is_supported_extension:
        raise BreesyError(f"Unsupported file extension: {file_extension}",
                          "Provide a file with a supported extension, example: 'my_data.gdf'. "
                          f"Currently supported extensions: {SUPPORTED_FILE_EXTENSIONS}")

    path_filename = Path(filename)
    is_valid_filename = Path.is_file(path_filename)
    if not is_valid_filename:
        raise BreesyError(
            f"File not found: {filename}", "a) Make sure that the file is in your working directory; "
                                           f"the current working directory is: {Path.cwd()}. "
                                           f"b) Double check that your file explorer is showing the full file name "
                                           f"- Windows does not show file extensions by default.")

    if file_extension == '.xdf':
        from .load_xdf import load_xdf
        return load_xdf(filename)
    elif file_extension == '.edf':
        from .load_edf import load_edf
        return load_edf(filename)
    elif file_extension == '.bdf':
        from .load_edf import load_bdf
        return load_bdf(filename)
    elif file_extension in ['.eeg', '.avg', '.seg', '.vhdr']:
        from .load_brainvision import load_brainvision_eeg
        if sample_rate is not None:
            logger.warning(f"Ignoring provided sample_rate of {sample_rate} because BrainVision files contain their own sample rate information.")
        return load_brainvision_eeg(filename)
    elif file_extension == '.csv':
        return load_csv(filename, sample_rate=sample_rate)
    elif file_extension == '.mat':
        from .load_matlab import load_mat
        return load_mat(filename, sample_rate=sample_rate)
    elif file_extension == '.tsv':
        return load_tsv(filename, sample_rate=sample_rate)
    elif file_extension == '.npy':
        return load_numpy(filename, sample_rate=sample_rate)
    elif file_extension in ['.set', '.fdt']:
        from .load_eeglab import load_eeglab
        if sample_rate is not None:
            logger.warning(f"Ignoring provided sample_rate of {sample_rate} because EEGLab files contain their own sample rate information.")
        return load_eeglab(filename)
    else:
        raise BreesyInternalError("File extension valid but not handled")


@enforce_type_hints
def load_delimited_text(filename: str, delimiter: str = ',', sample_rate: int | float | None = None) -> Recording:
    """Load a delimited text file (CSV, TSV, etc.).

    :param filename: Path to the text file
    :param delimiter: Character used to separate values in the file
    :param sample_rate: Sampling rate in Hz. If not provided, will attempt to guess from time column. If sample rate cannot be guessed, an error will be raised.

    :return: Recording object containing the loaded data
    """
    import csv

    # First pass: read headers and determine structure
    with open(filename, 'r', newline='') as file:
        reader = csv.reader(file, delimiter=delimiter)
        headers = next(reader)  # Get column names

        # Check if first column might be time
        time_column_idx = None
        if headers[0].lower() in ['time', 't', 'timestamp']:
            time_column_idx = 0

        # Exclude time column from channel names if present
        if time_column_idx is not None:
            channel_names = headers[1:]
        else:
            channel_names = headers

        # Count rows for pre-allocation
        row_count = sum(1 for _ in reader)

    # Second pass: read all data
    with open(filename, 'r', newline='') as file:
        reader = csv.reader(file, delimiter=delimiter)
        next(reader)  # Skip header

        # Pre-allocate data array
        data = np.zeros((len(channel_names), row_count))
        time_values = None

        if time_column_idx is not None:
            time_values = np.zeros(row_count)

        # Fill data
        for i, row in enumerate(reader):
            if i >= row_count:
                break

            # Convert values to float, handling potential errors
            try:
                row_values = [float(val) for val in row]
            except ValueError:
                # Skip rows that can't be converted to float
                continue

            if time_column_idx is not None:
                time_values[i] = row_values[time_column_idx]
                data_values = row_values[:time_column_idx] + row_values[time_column_idx + 1:]
            else:
                data_values = row_values

            for j, val in enumerate(data_values):
                if j < len(channel_names):
                    data[j, i] = val

    if sample_rate is not None:
        sample_rate = float(sample_rate)
    else:
        if time_values is not None and len(time_values) > 1:
            # Estimate sample rate
            time_diffs = np.diff(time_values)
            avg_diff = float(np.mean(time_diffs))
            if avg_diff > 0:
                sample_rate = float(round(1.0 / avg_diff))
        
        if sample_rate is None:
            raise BreesyError(
                "Sample rate was not provided and it could not be guessed from the data. "
                "Please provide the sample rate manually by passing the 'sample_rate' parameter.",
                "Provide the sample rate as a named argument to the function, e.g.: load_delimited_text('file.csv', sample_rate=500)"
            )


    # Look for event columns (columns with mostly 0s and few 1s)
    events = []
    for i, col_name in enumerate(channel_names):
        col_data = data[i, :]
        unique_vals = set(np.unique(col_data))
        # Check if column is binary and sparse (potential event marker)
        if unique_vals <= {0, 1} and np.mean(col_data) < 0.1:
            event_indices = np.where(col_data == 1)[0]
            for idx in event_indices:
                events.append(Event(name=col_name, index=idx))

    data, channel_names = _leave_only_eeg_channels(data, channel_names)

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

    return Recording(
        data=data,
        channel_names=channel_names,
        sample_rate=sample_rate,
        events=events,
        metadata=metadata
    )

# wrapper function for backward compatibility
@enforce_type_hints
def load_csv(filename: str) -> Recording:
    """Load a CSV file.

    :param filename: Path to the CSV file

    :return: Recording object containing the loaded data
    """
    return load_delimited_text(filename, delimiter=',')


# wrapper function for backward compatibility
@enforce_type_hints
def load_tsv(filename: str) -> Recording:
    """Load a TSV file.

    :param filename: Path to the TSV file

    :return: Recording object containing the loaded data
    """
    return load_delimited_text(filename, delimiter='\t')


@enforce_type_hints
def load_numpy(filename: str, sample_rate: int | float | None = None, channel_names: list[str] | None = None) -> Recording:
    """Load a numpy .npy file.

    :param filename: Path to the .npy file
    :param sample_rate: Sampling rate in Hz. If not provided, will attempt to detect from time column.
                        If no time column is found and sample_rate is None, an error will be raised.
    :param channel_names: List of channel names. If not provided, default names will be generated.

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

    # TODO: #12 we need to handle this case properly, currently just taking first trial
    if data.ndim == 3:
        # likely many recordings instead of one
        logger.warning("The loaded NumPy array has 3 dimensions, which likely means it contains multiple recordings; loading only the first one")
        data = data[0]

    if data.ndim != 2:
        raise BreesyError(
            f"The NumPy array has {data.ndim} dimensions, but 2 dimensions are required.",
            "Make sure the NumPy file contains a 2D array with shape (channels, samples)."
        )

    # if the first dimension is larger, assume it's (samples, channels) and transpose
    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

    extracted_sample_rate = None
    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")
        # remove time column from data
        data = np.delete(data, time_column_idx, axis=0)

    # 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_numpy(filename, sample_rate=250)"
        )

    # Create channel names
    if channel_names is None:
        channel_names = [f'Channel_{i+1}' for i in range(data.shape[0])]
    elif len(channel_names) != data.shape[0]:
        raise BreesyError(
            f"The number of channel names ({len(channel_names)}) does not match the number of channels ({data.shape[0]}) in the data.",
            "Make sure that the number of channel names matches the number of channels in the data. "
            f"Expected {data.shape[0]} channel names."
        )

    events = find_and_load_event_for_file(filename, sample_rate)

    data, channel_names = _leave_only_eeg_channels(data, channel_names)

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

    return Recording(
        data=data,
        channel_names=channel_names,
        sample_rate=sample_rate,
        events=events,
        metadata=metadata
    )


@enforce_type_hints
def from_array(data: np.ndarray, sample_rate: int | float | None = None, channel_names: list[str] | None = None) -> Recording:
    """Convert a numpy array into a Recording object.

    :param data: NumPy array containing EEG data. Should be 2D with shape (channels, samples) or (samples, channels).
                 If a time column is detected, it will be used to extract the sample rate and removed from the data.
    :param sample_rate: Sampling rate in Hz. If not provided, will attempt to detect from time column.
                        If no time column is found and sample_rate is None, an error will be raised.
    :param channel_names: List of channel names. If not provided, default names will be generated.

    :return: Recording object containing the loaded data
    """
    # Ensure it's a 2D array
    if data.ndim != 2:
        raise BreesyError(
            f"The NumPy array has {data.ndim} dimensions, but 2 dimensions are required.",
            "Make sure the NumPy array is 2D with shape (channels, samples) or (samples, channels)."
        )

    # Auto-transpose if the first dimension is larger (likely samples, channels)
    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

    extracted_sample_rate = None
    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")
        # remove time column from data
        data = np.delete(data, time_column_idx, axis=0)

    # 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: from_array(data, sample_rate=250)"
        )

    # Create channel names
    if channel_names is None:
        channel_names = [f'Channel_{i+1}' for i in range(data.shape[0])]
    elif len(channel_names) != data.shape[0]:
        raise BreesyError(
            f"The number of channel names ({len(channel_names)}) does not match the number of channels ({data.shape[0]}) in the data.",
            "Make sure that the number of channel names matches the number of channels in the data. "
            f"Expected {data.shape[0]} channel names."
        )
    
    data, channel_names = _leave_only_eeg_channels(data, channel_names)

    # Create metadata
    metadata = RecordingMetadata(
        name="NumPy Recording",
        file_path=None,
        description="Recording created from NumPy array"
    )

    return Recording(
        data=data,
        channel_names=channel_names,
        sample_rate=sample_rate,
        events=None,
        metadata=metadata
    )


# @enforce_type_hints
# def load_all_files_in_folder_by_extension(file_extension: str, folder_path: str = ""):
#     """Load all files with a specific extension from a folder.

#     :param file_extension: File extension to filter by (e.g., '.mat')
#     :param folder_path: Path to the folder containing the files

#     :return: List of Recording objects containing the loaded data
#     """
#     if not isinstance(folder_path, str):
#         raise BreesyError("Invalid folder_path argument type",
#                           "Provide a string as the folder path, example: 'data_folder', or '' (empty string or period) for current working directory.")

#     if not isinstance(file_extension, str):
#         raise BreesyError("Invalid file_extension argument type",
#                           "Provide a string as the file extension, example: '.gdf'")

#     folder = Path(folder_path)
#     if not folder.is_dir():
#         raise BreesyError(f"The folder path was not found: {folder.resolve()}",
#                           f"Check the folder path manually to make sure that the folder actually exists.")

#     full_path = folder.resolve()

#     files = list(full_path.glob(f"*{file_extension}"))
#     if not files:
#         raise BreesyError(
#             f"No {file_extension} files found in folder: '{folder_path}' (resolved absolute path: {folder.resolve()})",
#             f"Check the folder to make sure that it actually contains {file_extension} files.")

#     data_list = []
#     for file in files:
#         data = load_dataset(str(file))
#         data_list.append(data)

#     return data_list


# --------- Helper functions ---------

@protect_from_lib_error("numpy")
def _np_load(filename):
    return np.load(filename)


def _detect_time_column(data: np.ndarray) -> int | None:
    for i in range(data.shape[0]):
        if _is_likely_time_column(data[i]):
            return i
    return None


def _calculate_sample_rate_from_time_column(data: np.ndarray) -> float:
    if (data.size < 2) or data[0] == data[1]:
        raise BreesyInternalError(
            "Cannot calculate sample rate from time column with less than 2 unique values. Make sure the time column contains valid time values."
        )
    time_step = float(data[1] - data[0])
    sample_rate = float(1 / time_step)
    return sample_rate


def _is_likely_eeg_column(column: np.ndarray) -> bool:
    return (not _column_is_monotonic(column)) and (not _column_is_sparse(column))


def _is_likely_time_column(column: np.ndarray) -> bool:
    u_diffs = np.unique(np.diff(column).round(6))
    if len(u_diffs) == 1 and u_diffs[0] > 0:
        return True
    return False


def _leave_only_eeg_channels(data: np.ndarray, ch_names: list[str]) -> tuple[np.ndarray, list[str]]:
    eeg_ix = [i for i in range(data.shape[0]) if _is_likely_eeg_column(data[i])]
    removed = [name for i, name in enumerate(ch_names) if i not in eeg_ix]
    if removed:
        logger.warning(f"Ignored {len(removed)} channel(s) that don't look like EEG signal: {removed}")
    
    new_channel_names = [x for i, x in enumerate(ch_names) if i in eeg_ix]
    return data[eeg_ix], new_channel_names


def _column_is_monotonic(column: np.ndarray) -> bool:
    diffs = np.diff(column)
    return bool(np.all(diffs >= 0))


def _column_is_sparse(column: np.ndarray, sparsity_threshold: float = 0.2) -> bool:
    u_values, u_counts = np.unique(column, return_counts=True)
    most_common_value_rate = np.max(u_counts) / len(column)
    return bool(most_common_value_rate > sparsity_threshold)


def _find_keys_in_data(keys: list[str], searching_for: list[str]) -> list:
    found = []
    translation = str.maketrans('', '', '_- .')
    to_check = [k.lower().translate(translation) for k in searching_for]
    for k in keys:
        if k.lower().translate(translation) in to_check:
            found.append(k)
    return found
