import numpy as np
from math import gcd
import itertools
from tqdm.auto import tqdm
from scipy.stats import pearsonr

from . import breesy_scipy, constants, Event
from .recording import Recording, update_data, cut_by_sample_range
from .filtering import notch_filter
from .breesy_scipy import _si_welch, _si_decimate, _si_resample_poly
from .breesy_sklearn import _sklearn_fastica
from .errors import BreesyError, protect_from_lib_error
from .type_hints import enforce_type_hints
from .log import logger


FREQ_SPECTRUM_PARAMS = """
:param scaling: Scaling to use, either "density" (default) for power spectral density, or "spectrum" for squared magnitude spectrum
:param window: Window type to use (see https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.windows.get_window.html)
"""


_DOC_PARAMS = {
    name: value for name, value in globals().items()
    if name.endswith('_PARAMS') and isinstance(value, str)
}


def _format_doc(func):
    """Decorator that expands shared parameter docs in the docstring."""
    if func.__doc__:
        func.__doc__ = func.__doc__.format(**_DOC_PARAMS)
    return func


@_format_doc
@enforce_type_hints
def get_frequency_spectrum_channel(recording: Recording, channel_name: str, scaling: str = 'density',
                                   window: str | None = None) -> tuple[np.ndarray, np.ndarray]:
    """Compute the frequency spectrum of a single EEG channel.

    :param recording: Input EEG recording
    :param channel_name: Name of the channel to process
    {FREQ_SPECTRUM_PARAMS}
    :return: Tuple of (frequencies, power_spectrum)
    """
    if channel_name not in recording.channel_names:
        raise BreesyError(
            f"Channel '{channel_name}' not found in recording.",
            "Use one of the channels in the recording; you can view them by doing: `recording.channel_names`"
        )

    channel_index = recording.channel_names.index(channel_name)
    channel_data = recording.data[channel_index]

    f, spectrum = _si_welch(data=channel_data[np.newaxis, :], sample_rate=recording.sample_rate,
                            scaling=scaling, window=window, nperseg=4*recording.sample_rate)
    return f, spectrum[0]


@_format_doc
@enforce_type_hints
def get_frequency_spectrum(recording: Recording, scaling: str = 'density', window: str | None = None) -> tuple[np.ndarray, np.ndarray]:
    """Compute the frequency spectra of each EEG channel.

    :param recording: Input EEG recording
    {FREQ_SPECTRUM_PARAMS}
    :return: Tuple containing frequencies array and periodogram array of shape (n_channels, n_frequencies)
    """

    f, spectrum = _si_welch(data=recording.data, sample_rate=recording.sample_rate, scaling=scaling, window=window,
                            nperseg=4*recording.sample_rate)
    return f, spectrum


@_format_doc
@enforce_type_hints
def get_mean_frequency_spectrum(recording: Recording, scaling: str = 'density',
                                window: str | None = None) -> tuple[np.ndarray, np.ndarray]:
    """Compute the mean frequency spectrum of all EEG channels.

    :param recording: Input EEG recording
    {FREQ_SPECTRUM_PARAMS}
    :return: Tuple containing frequencies array and periodogram array of shape (n_channels, n_frequencies)
    """

    f, spectrum = get_frequency_spectrum(recording=recording, scaling=scaling, window=window)
    return f, spectrum.mean(axis=0)


# Detection functions --------------------------


@enforce_type_hints
def check_for_data_problems(recording: Recording) -> dict[str, list]:
    """Detect common problems in data:
    - Missing values (NaNs) - if any found, it is highly recommended to repair the data file first before doing any analysis
    - Data scale - estimated unit of data
    - Dead channels - channels which have very low amplitude and likely disconnected during recording
    - Noisy channels - channels which have very high variance and likely overwhelmed by noise
    - Powerline frequency - if there are 50 Hz or 60Hz peaks in some channels, it is recommended to filter out that frequency in all channels

    :param recording: Input recording

    :return: Dictionary with found parameters and problematic values
    """
    nan_channels = detect_nans(recording=recording)
    data_scale = detect_data_scale(recording=recording)
    dead_channels = detect_dead_channels(recording=recording, std_threshold=constants.DEAD_CHANNEL_SD_THRESHOLD, skip_channels=[])
    noisy_channels = detect_noisy_channels(recording=recording, variance_threshold=constants.NOISY_CHANNEL_VARIANCE_THRESHOLD, skip_channels=dead_channels)
    found_powerline_peaks = detect_powerline_noise(recording=recording,
                                                   power_threshold=constants.POWERLINE_POWER_THRESHOLD,
                                                   baseline_l_freq=constants.BASELINE_FOR_PEAKS_LOW,
                                                   baseline_h_freq=constants.BASELINE_FOR_PEAKS_HIGH,
                                                   skip_channels=dead_channels)
    return {
        'missing_values': nan_channels,
        'scale': data_scale,
        'dead': dead_channels,
        'noisy': noisy_channels,
        'powerline': found_powerline_peaks
    }


@enforce_type_hints
@protect_from_lib_error("numpy")
def detect_dead_channels(recording: Recording, std_threshold: float = constants.DEAD_CHANNEL_SD_THRESHOLD, skip_channels: list[str] = None) -> list[str]:
    """Detect channels with abnormally low variability (possible dead channels).

    :param recording: Input recording
    :param std_threshold: Standard deviation threshold for flagging the channel as dead
    :param skip_channels: List of channel names to skip

    :return: List of channel names suspected to be dead
    """
    sid, eid = int(recording.number_of_samples * 0.05), int(recording.number_of_samples * 0.95)
    cut_recording = cut_by_sample_range(recording=recording, start_index=sid, end_index=eid)
    sds = np.std(cut_recording.data, axis=1)
    median_std = np.median(sds)
    dead_channels = [ch for i, ch in enumerate(recording.channel_names) if sds[i] < median_std * std_threshold]
    dead_channels = _skip_channel_by_names(ch_names=dead_channels, skip_names=skip_channels)
    _output_problematic_channels(problem_channels=dead_channels, all_channels=recording.channel_names, reason="dead channels")
    return dead_channels


@enforce_type_hints
@protect_from_lib_error("numpy")
def detect_noisy_channels(recording: Recording, variance_threshold: int | float = constants.NOISY_CHANNEL_VARIANCE_THRESHOLD, skip_channels: list[str] = None) -> list[str]:
    """Detect channels with abnormally high variance (likely noisy channels).

    :param recording: Input recording
    :param variance_threshold: Variance threshold multiplier for flagging the channel as noisy
    :param skip_channels: List of channel names to skip

    :return: List of channel names suspected to be noisy
    """
    variances = np.var(recording.data, axis=1)
    mean_var = np.mean(variances)
    noisy_channels = [ch for i, ch in enumerate(recording.channel_names) if variances[i] > mean_var * variance_threshold]
    noisy_channels = _skip_channel_by_names(ch_names=noisy_channels, skip_names=skip_channels)
    _output_problematic_channels(problem_channels=noisy_channels, all_channels=recording.channel_names, reason="highly noisy channels")
    return noisy_channels


@enforce_type_hints
@protect_from_lib_error("numpy")
def detect_nans(recording: Recording) -> list[str]:
    n_nans_per_channel = np.isnan(recording.data).sum(axis=1)
    nan_channels = [ch for i, ch in enumerate(recording.channel_names) if n_nans_per_channel[i] > 0]
    _output_problematic_channels(problem_channels=nan_channels, all_channels=recording.channel_names, reason="channels with missing values")
    return nan_channels


@enforce_type_hints
@protect_from_lib_error("numpy")
def detect_data_scale(recording: Recording) -> float:
    """Detect the likely scale of EEG data and return a multiplier to convert to microvolts.
    
    This function analyzes the amplitude range of EEG data to determine if it's in
    volts (V), millivolts (mV), or microvolts (µV), and returns the appropriate
    multiplier to convert the data to microvolts.
    
    Typical EEG amplitudes:
    - In microvolts (µV): 10-100 µV (multiplier: 1.0)
    - In millivolts (mV): 0.01-0.1 mV (multiplier: 1000.0)
    - In volts (V): 0.00001-0.0001 V (multiplier: 1000000.0)
    
    :param recording: Input recording

    :return: Multiplier to convert data to microvolts
    """
    channel_stds = np.std(recording.data, axis=-1)
    typical_amplitude = np.median(channel_stds)
    
    if typical_amplitude > 1.0:
        scale = "microvolts (µV)"
        multiplier = 1.0
    elif typical_amplitude > 0.001:
        scale = "millivolts (mV)"
        multiplier = 1000.0
    else:
        scale = "volts (V)"
        multiplier = 1_000_000.0
    
    logger.info(f"EEG data is likely in {scale} (typical amplitude: {typical_amplitude:.6f})")
    
    return multiplier



def _skip_channel_by_names(ch_names: list[str], skip_names: list[str] | None) -> list[str]:
    if skip_names:
        return [ch for ch in ch_names if ch not in skip_names]
    return ch_names


def _output_problematic_channels(problem_channels: list[str], all_channels: list[str], reason: str) -> None:
    if not problem_channels:
        logger.info(f'Checked for {reason}: no problematic channels detected.')
    elif len(problem_channels) == len(all_channels):
        logger.info(f'Checked for {reason}: detected in all channels.')
    else:
        ratio = round(len(problem_channels) / len(all_channels) * 100)
        ch_names_str = ' - ' + ', '.join(problem_channels) if len(problem_channels) < 10 else ''
        logger.info(f"Checked for {reason}: likely detected in {len(problem_channels)} ({ratio}%) channel(s){ch_names_str}.")


@enforce_type_hints
def detect_powerline_noise(recording: Recording, power_threshold: float = constants.POWERLINE_POWER_THRESHOLD,
                           baseline_l_freq: int | float = constants.BASELINE_FOR_PEAKS_LOW,
                           baseline_h_freq: int | float = constants.BASELINE_FOR_PEAKS_HIGH,
                           skip_channels: list[str] = None) -> list[str]:
    if recording.sample_rate / 2 < 50:
        logger.info(f'Powerline noise cannot be detected because sample rate is too low.')
        return []

    f, spectra = get_frequency_spectrum(recording=recording)
    baseline = _get_baseline_for_peak_detection(f=f, x=spectra, start_freq=baseline_l_freq, end_freq=baseline_h_freq)
    if skip_channels:
        keep_mask = np.array([False if ch in skip_channels else True for ch in recording.channel_names])
    else:
        keep_mask = np.full(recording.number_of_channels, True)

    powerline_peaks, power_freq = None, None

    powerline_freqs = [50, 60]
    for freq in powerline_freqs:
        if recording.sample_rate / 2 < freq:
            break
        if _notch_filter_detected(f=f, x=spectra, freq=freq):
            continue

        peaks = _measure_rel_peaks_at(f=f, x=spectra, freq=freq, baseline=baseline)
        if peaks[keep_mask].max() > power_threshold:
            powerline_peaks = peaks
            power_freq = freq
            break

    if power_freq is not None:
        powerline_peaks[~keep_mask] = 1  # set to baseline
        noisy_channels = [ch for i, ch in enumerate(recording.channel_names)
                          if powerline_peaks[i] > power_threshold]
        _output_problematic_channels(problem_channels=noisy_channels, all_channels=recording.channel_names,
                                     reason=f'prominent {power_freq} Hz powerline noise')
        # harmonics
        noisiest_i = np.argmax(powerline_peaks).ravel()
        for harmonic in range(power_freq*2, int(recording.sample_rate // 2), power_freq):
            harmonic_peak = _measure_rel_peaks_at(f=f, x=spectra[noisiest_i], freq=harmonic, baseline=baseline[noisiest_i])
            if harmonic_peak > power_threshold:
                logger.info(f'\t+ harmonic of {harmonic} Hz')

        return noisy_channels
    else:
        logger.info(f'No powerline noise detected or was already filtered.')
        return []


def _notch_filter_detected(f: np.ndarray, x: np.ndarray, freq: int) -> bool:
    freq_i = np.searchsorted(f, freq)
    power_at_freq = x[:, freq_i]

    window_size = 5
    attenuation_threshold = 0.3
    ch_ratio_with_notch = 0.5
    lower_i = np.searchsorted(f, freq - window_size)
    upper_i = np.searchsorted(f, freq + window_size)
    surrounding_power = np.concatenate([
        x[:, lower_i:freq_i], x[:, freq_i+1:upper_i]
    ], axis=1)
    avg_surrounding = np.mean(surrounding_power, axis=1)
    ratio = power_at_freq / avg_surrounding
    return np.sum(ratio < attenuation_threshold) > x.shape[0] * ch_ratio_with_notch


@protect_from_lib_error("numpy")
def _get_baseline_for_peak_detection(f: np.ndarray, x: np.ndarray, start_freq: int | float, end_freq: int | float) -> np.ndarray:
   start_i = np.searchsorted(f, start_freq)
   end_i = np.searchsorted(f, end_freq)
   baseline = np.percentile(x[:, start_i:end_i], q=95, axis=1)
   return baseline


@protect_from_lib_error("numpy")
def _measure_rel_peaks_at(f: np.ndarray, x: np.ndarray, freq: int | float, baseline: float) -> np.ndarray:
    freq_i = np.searchsorted(f, freq)
    return x[:, freq_i] / baseline


# ----------------

# TODO: move to a new file called decomposition? or to transform?
@enforce_type_hints
def get_ica_components(recording: Recording, n_components: int, random_state: int | None = None) -> np.ndarray:
    """Compute Independent Component Analysis (ICA) components.

    :param recording: Input recording with data shaped (n_channels, n_samples)
    :param n_components: Number of ICA components to compute
    :param random_state: Optional random seed for reproducibility

    :return: Components array of shape (n_components, n_samples)
    """
    ica = _sklearn_fastica(n_components=n_components, random_state=random_state)
    components = ica.fit_transform(recording.data.T)
    return components.T


@enforce_type_hints
def get_channel_correlation(recording: Recording, absolute_values: bool = True) -> np.ndarray:
    """Compute matrix of Pearson correlation between all channels in data.

    :param recording: Input EEG recording
    :absolute_values: Whether to ues absolute (in [0, 1] range) correlation values. Setting this False will keep correlation in [-1, 1] range. 

    :return: Correlation array of shape (n_channels, n_channels)
    """
    n_channels = recording.number_of_channels
    matrix = np.ones((n_channels, n_channels))
    matrix_iter = tqdm(itertools.combinations(range(n_channels), r=2), desc="Calculating correlations between channels",
                       total=int(n_channels * (n_channels-1) / 2), disable=n_channels < 33)
    for i, j in matrix_iter:
        corr = pearsonr(recording.data[i], recording.data[j]).statistic
        if absolute_values:
            corr = abs(corr)
        matrix[i, j] = corr
        matrix[j, i] = corr
    return matrix

# ----------------


@enforce_type_hints
def remove_powerline_noise(recording: Recording, smoothness_factor: int | float = constants.NOTCH_DEFAULT_SMOOTHNESS,
                           power_threshold: float = constants.POWERLINE_POWER_THRESHOLD) -> Recording:
    """Remove powerline noise from a recording.

    :param recording: Input recording
    :param quality_factor: Quality factor for the notch filter

    :return: Filtered recording
    """
    frequency = detect_powerline_noise(recording=recording, power_threshold=power_threshold)
    if frequency:
        for freq in frequency:
            recording = notch_filter(recording=recording, frequency=freq, smoothness_factor=smoothness_factor)
        return recording
    return recording


@enforce_type_hints
@protect_from_lib_error("numpy")
def mean_centering(recording: Recording) -> Recording:
    """Center the recording data by subtracting the mean.

    :param recording: Input recording

    :return: Mean-centered recording
    """
    new_eeg_data = recording.data - np.expand_dims(recording.data.mean(axis=-1), 1)
    return update_data(recording, new_eeg_data)


@enforce_type_hints
def remove_slow_drift(recording: Recording, highpass_freq: int | float = constants.SLOWDRIFT_REMOVAL_FREQ,
                      filter_order: int = constants.BUTTER_FILTER_ORDER) -> Recording:
    """Remove slow drift using a high-pass Butterworth filter.

    :param recording: Input recording
    :param highpass_freq: High-pass cutoff frequency in Hz
    :param filter_order: Order of the Butterworth high-pass filter

    :return: Filtered recording
    """
    highpass = breesy_scipy._si_butter_highpass_sos(freq=highpass_freq, sample_rate=recording.sample_rate, filter_order=filter_order)
    new_eeg_data = breesy_scipy._si_sosfiltfilt(sos=highpass, data=recording.data)
    return update_data(recording, new_eeg_data)


@enforce_type_hints
@protect_from_lib_error("numpy")
def rereference(recording: Recording, channel_name: str) -> Recording:
    """Re-reference EEG by subtracting a reference signal from each channel.

    :param recording: Input recording
    :param channel_name: Name of the reference signal to subtract

    :return: Re-referenced recording
    """
    ch_i = recording.channel_names.find(channel_name)
    if ch_i == -1:
        available_channels = ", ".join([f'"{x}"' for x in recording.channel_names])
        raise BreesyError(f'Channel "{channel_name}" was not found in recording.',
                          f'Use one of available channels: {available_channels}.')
    new_eeg_data = recording.data - recording.data[ch_i]
    return update_data(recording, new_eeg_data)


@enforce_type_hints
def downsample(recording: Recording, new_sample_rate: int | float) -> Recording:
    """Downsample the recording to a new sample rate.
    
    :param recording: Input recording
    :param new_sample_rate: Target sample rate (must be lower than current sample rate)
    
    :return: Downsampled recording
    """
    if new_sample_rate >= recording.sample_rate:
        raise BreesyError(f'New sample rate ({new_sample_rate} Hz) must be lower than the current sample rate ({recording.sample_rate} Hz).',
                          "Provide a lower sample rate to downsample the data.")

    downsample_factor = recording.sample_rate / new_sample_rate
    if downsample_factor.is_integer():
        new_data = _si_decimate(recording.data, int(downsample_factor))
    else:
        g = gcd(int(new_sample_rate), int(recording.sample_rate))
        up = int(new_sample_rate / g)
        down = int(recording.sample_rate / g)
        new_data = _si_resample_poly(
            recording.data,
            up=up,
            down=down,
            axis=-1)

    new_events = [
        Event(event.name, index=int(event.index / downsample_factor))
        for event in recording.events
    ]
    return Recording(new_data, recording.channel_names, new_sample_rate, new_events)
