import numpy as np
from thunderhopper.filetools import search_files
from thunderhopper.modeltools import load_data
from thunderhopper.filters import sosfilter
from IPython import embed

# GENERAL SETTINGS:
target = 'Omocestus_rufipes'
data_path = search_files(target, dir='../data/processed/')[0]
save_path = '../data/inv/noise_env/'

# ANALYSIS SETTINGS:
scales = np.geomspace(0.01, 1000, 100)
sd_inputs = np.array([1.0])
n_trials = 100

# EXECUTION:

# Load signal data:
data, config = load_data(data_path, files='raw')
signal, rate = data['raw'], config['rate']

# Reduce to song segment and normalize:
time = np.arange(signal.shape[0]) / rate
start, end = data['songs_0'].ravel()
segment = (time >= start) & (time <= end)
signal /= signal[segment].std()

# Rescale signal (time, scale):
signal = signal[:, None] * scales[None, :]

# Prepare storage:
sd_noise = np.zeros((scales.size, n_trials, sd_inputs.size), dtype=float)

# Analyze piece-wise:
rng = np.random.default_rng()
for i, sigma in enumerate(sd_inputs):
    print(f'Testing SD: {sigma:.3f} ...')

    # Prepare trial storage:
    sd_trials = np.zeros((segment.sum(), scales.size, n_trials), dtype=float)

    # Run trials:
    for j in range(n_trials):
        # Mix signals with white noise of target SD:
        mix = signal + rng.normal(0, sigma, signal.shape)

        # Process mixture:
        mix = sosfilter(mix, rate, config['bp_fcut'], 'bp',
                        padtype='fixed', padlen=config['padlen'])
        mix = sosfilter(np.abs(mix), rate, config['env_fcut'], 'lp',
                        padtype='even', padlen=config['padlen'])

        # Log current trial:
        sd_trials[..., j] = mix[segment, :]

    # Get noise remainders of mean over trials:
    sd_trials -= sd_trials.mean(axis=-1, keepdims=True)

    # Estimate noise SD:
    sd_noise[:, :, i] = sd_trials.std(axis=0)

    # # Add Gaussian noise of given SD to rescaled signals (time, scale, trial):
    # mix = signal[..., None] + rng.normal(0, sigma, (*signal.shape, n_trials))

    # # Get mixture envelopes (time, scale, trial):
    # mix = sosfilter(np.abs(mix), rate, config['env_fcut'], 'lp',
    #                 padtype='even', padlen=config['padlen'])[segment, ...]

    # # Get noise remainders of mean over trials:
    # mix -= mix.mean(axis=-1, keepdims=True)

    # # Estimate noise SD:
    # sd_noise[:, :, i] = mix.std(axis=0)

if save_path is not None:
    archive = dict(
        scales=scales,
        sd_input=sd_inputs,
        sd_noise=sd_noise,
    )
    np.savez(save_path + 'sd_conversion.npz', **archive)

print('Done.')
embed()