import glob
import numpy as np
import matplotlib.pyplot as plt
from thunderhopper.modeltools import load_data
from thunderhopper.filters import sosfilter
from IPython import embed

# GENERAL SETTINGS:
target = 'Omocestus_rufipes'
data_path = glob.glob(f'../data/processed/{target}*.npz')[0]
save_path = '../data/inv/noise_env/'

# ANALYSIS SETTINGS:
scales = np.geomspace(0.1, 10000, 200)
sd_inputs = np.arange(10.9, 11.1, 0.01)
n_trials = 10
tol_to_one = 0.1

# EXECUTION:

# Load signal data:
data, config = load_data(data_path, files='filt')
signal, rate = data['filt'], 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()

# Get rescaled signals (time, scale):
signal = signal[:, None] * scales[None, :]

# Prepare storage:
current_match = 0
storage = dict(
    scales=scales,
    n_trials=n_trials,
    sd_factor=np.array([0.]),
    trials=np.zeros((scales.size, n_trials), dtype=float),
    mean=np.zeros(scales.size, dtype=float),
    spread=np.zeros(scales.size, dtype=float),
)

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

    # 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 = mix.std(axis=0)
    mean_sd = sd.mean(axis=-1)
    n_match = (np.abs(1 - mean_sd) <= tol_to_one).sum()
    if n_match > current_match:
        print(f'Found better SD: {sigma:.3f} with {n_match} matches (previous: {current_match})')
        storage['sd_factor'][0] = sigma
        storage['trials'][:, :] = sd
        storage['mean'][:] = mean_sd
        storage['spread'][:] = sd.std(axis=-1)
        current_match = n_match
    del mix
del signal
if save_path is not None:
    np.savez(save_path + 'sd_conversion.npz', **storage)

plt.plot(scales, storage['mean'], 'k')
plt.show()
embed()

print('Done.')
embed()