from pathlib import Path
import struct
import numpy as np
from scipy.stats import mode
from typing import BinaryIO

from .load import _is_likely_eeg_column
from .recording import Recording
from .RecordingMetadata import RecordingMetadata
from .errors import BreesyError, BreesyInternalError
from .log import logger
from .type_hints import enforce_type_hints


@enforce_type_hints
def load_edf(filename: str) -> Recording:
    """Load an EDF (European Data Format) or EDF+ file.

    :param filename: Path to the EDF file

    :return: Recording object containing the loaded data
    """
    with open(filename, 'rb') as f:
        # assert it is EDF
        first_byte = f.read(1)
        f.seek(0)
        if first_byte != b"0":
            raise BreesyError(f'The file is not correctly formatted for an EDF file.',
                              f'Ensure that "{filename}" is a proper EDF(+) file. Consider contacting authors of the file.')

        data, channel_names, sample_rate = _read_edf_or_bdf(f, bytes_per_sample=2)


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

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


@enforce_type_hints
def load_bdf(filename: str) -> Recording:
    """Load a BDF (BioSemi Data Format) or BDF+ file.

    :param filename: Path to the BDF file

    :return: Recording object containing the loaded data
    """        
    with open(filename, 'rb') as f:
        # assert it is BDF
        first_byte = f.read(1)
        f.seek(0)
        if first_byte[0] != 255:
            raise BreesyError(f'The file is not correctly formatted for an BDF file.',
                              f'Ensure that "{filename}" is a proper BDF(+) file. Consider contacting authors of the file.')

        data, channel_names, sample_rate = _read_edf_or_bdf(f, bytes_per_sample=3)

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

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


def _read_header(f: BinaryIO) -> dict:
    # Read fixed header (256 bytes)
    header = {}
    header['version'] = f.read(8).decode('ascii', errors='ignore').strip()
    header['patient_id'] = f.read(80).decode('ascii').strip()
    header['recording_id'] = f.read(80).decode('ascii').strip()
    header['start_date'] = f.read(8).decode('ascii').strip()
    header['start_time'] = f.read(8).decode('ascii').strip()
    header['header_bytes'] = int(f.read(8).decode('ascii').strip())
    header['reserved'] = f.read(44).decode('ascii').strip()
    header['num_records'] = int(f.read(8).decode('ascii').strip())
    header['record_duration'] = float(f.read(8).decode('ascii').strip())
    header['num_signals'] = int(f.read(4).decode('ascii').strip())
    return header


def _read_signal_header(f: BinaryIO, num_signals: int) -> dict:
    # Read signal headers (256 bytes per signal)
    signal_headers = {}
    signal_headers['label'] = [f.read(16).decode('ascii').strip() for _ in range(num_signals)]
    signal_headers['transducer'] = [f.read(80).decode('ascii').strip() for _ in range(num_signals)]
    signal_headers['physical_dim'] = [f.read(8).decode('ascii').strip() for _ in range(num_signals)]
    signal_headers['physical_min'] = [float(f.read(8).decode('ascii').strip()) for _ in range(num_signals)]
    signal_headers['physical_max'] = [float(f.read(8).decode('ascii').strip()) for _ in range(num_signals)]
    signal_headers['digital_min'] = [int(f.read(8).decode('ascii').strip()) for _ in range(num_signals)]
    signal_headers['digital_max'] = [int(f.read(8).decode('ascii').strip()) for _ in range(num_signals)]
    signal_headers['prefilter'] = [f.read(80).decode('ascii').strip() for _ in range(num_signals)]
    signal_headers['samples_per_record'] = [int(f.read(8).decode('ascii').strip()) for _ in range(num_signals)]
    signal_headers['reserved'] = [f.read(32) for _ in range(num_signals)]
    return signal_headers


def _get_scaling_factors_and_offsets(signal_headers: dict, num_signals: int) -> tuple[list, list]:
    scaling_factors, offsets = [], []
    for i in range(num_signals):
        digital_range = signal_headers['digital_max'][i] - signal_headers['digital_min'][i]
        physical_range = signal_headers['physical_max'][i] - signal_headers['physical_min'][i]
        if digital_range != 0:
            scaling_factors.append(physical_range / digital_range)
            offsets.append(signal_headers['physical_min'][i] - signal_headers['digital_min'][i] * scaling_factors[i])
        else:
            scaling_factors.append(1.0)
            offsets.append(0.0)
    return scaling_factors, offsets


def _filter_plus_annotations(data: list[np.ndarray], channel_names: list[str]) -> tuple[list[np.ndarray], list[str], list[int]]:
    annot_labels = ['EDF Annotations', 'BDF Annotations']
    other_labels = ['Status']
    likely_eeg = [_is_likely_eeg_column(ch) for ch in data]
    eeg_indices = [i for i, label in enumerate(channel_names)
                   if all(x not in label for x in annot_labels+other_labels)
                   and likely_eeg[i]]
    removed = [name for i, name in enumerate(channel_names) if i not in eeg_indices]
    if removed:
        logger.warning(f"Ignored {len(removed)} channel(s) that don't look like EEG signal: {removed}")
        data = [data[i] for i in eeg_indices]
        channel_names = [channel_names[i] for i in eeg_indices]
        # TODO: create events from annotations
    return data, channel_names, eeg_indices


def _determine_sample_rate(samples_per_record: list[int], header: dict, eeg_indices: list[int], channel_names: list[str]) -> int:
    sample_rates = [samples_per_record[i] / header['record_duration'] for i in eeg_indices]
    if len(set(sample_rates)) == 1:
        sample_rate = float(sample_rates[0])
    else:
        sample_rate = mode(sample_rates)
        diff_channels_str = ', '.join([ch for i, ch in enumerate(channel_names) if sample_rates[i] != sample_rate])
        logger.warning(f'Detected variable sample rates. Will use {sample_rate} as the most common one. Channel names with different sample rates: {diff_channels_str}')
    return sample_rate


def _convert_to_microvolts(data: list[np.ndarray], units: list[str]) -> list[np.ndarray]:
    converted_data = []
    for i, (channel_data, unit) in enumerate(zip(data, units)):
        unit_normalized = unit.lower().strip()
        
        if unit_normalized in ['μv', 'uv', 'µv']:
            conversion_factor = 1.0
        elif unit_normalized == 'mv':
            conversion_factor = 1000.0
        elif unit_normalized == 'v':
            conversion_factor = 1_000_000.0
        else:
            raise BreesyError(f'Unknown unit specified for channel nr. {i+1}: "{unit}".',
                              f'Contact the authors of the file and ask to fix it. Breesy only supports "μV", "mV", and "V" (not case-sensitive). '
                              f'If you think that "{unit}" unit should also be supported, contact Breesy authors.')
        
        converted_data.append(channel_data * conversion_factor)
    
    return converted_data


def _read_edf_or_bdf(f: BinaryIO, bytes_per_sample: int) -> tuple[np.ndarray, list[str], int]:
    if bytes_per_sample not in [2, 3]:
        raise BreesyInternalError(f'bytes_per_sample can only be 2 or 3, provided {bytes_per_sample}.')
    
    header = _read_header(f)
    num_signals = header['num_signals']
    signal_headers = _read_signal_header(f, num_signals=num_signals)
    scaling_factors, offsets = _get_scaling_factors_and_offsets(signal_headers=signal_headers, num_signals=num_signals)

    num_records = header['num_records']
    samples_per_record = signal_headers['samples_per_record']
    total_samples = [samples_per_record[i] * num_records for i in range(num_signals)]

    data = [np.zeros(total_samples[i]) for i in range(num_signals)]

    # Read all data records
    for record_idx in range(num_records):
        for signal_idx in range(num_signals):
            n_samples = samples_per_record[signal_idx]

            if bytes_per_sample == 2:
                raw_samples = struct.unpack(f'<{n_samples}h', f.read(n_samples * 2))
            else:
                raw_bytes = f.read(n_samples * 3)
                raw_samples = []
                for i in range(0, len(raw_bytes), 3):
                    value = int.from_bytes(raw_bytes[i:i+3], byteorder='little', signed=True)
                    raw_samples.append(value)
            
            # Convert to physical values
            start_idx = record_idx * n_samples
            end_idx = start_idx + n_samples
            data[signal_idx][start_idx:end_idx] = np.array(raw_samples) * scaling_factors[signal_idx] + offsets[signal_idx]

    channel_names = [x.replace('.', '') for x in signal_headers['label']]
    data, channel_names, eeg_indices = _filter_plus_annotations(data=data, channel_names=channel_names)
    units = [signal_headers['physical_dim'][i] for i in eeg_indices]
    data = _convert_to_microvolts(data, units)
    sample_rate = _determine_sample_rate(samples_per_record=samples_per_record, header=header,
                                         eeg_indices=eeg_indices, channel_names=channel_names)
    return np.array(data), channel_names, sample_rate
