import numpy as np
from thunderhopper.modeltools import load_data, save_data
from thunderhopper.filetools import search_files, crop_paths
from thunderhopper.filters import decibel, sosfilter
from IPython import embed

# GENERAL SETTINGS:
target = ['Omocestus_rufipes', '*'][0]
data_paths = search_files(target, excl='noise', dir='../data/processed/')
save_path = '../data/inv/log_hp/'

# ANALYSIS SETTINGS:
add_noise = False
save_snippets = target == 'Omocestus_rufipes'
example_scales = np.array([0.1, 1, 10, 30, 100, 300])
scales = np.geomspace(0.1, 10000, 500)
scales = np.unique(np.concatenate((scales, example_scales)))

# EXECUTION:
for data_path, name in zip(data_paths, crop_paths(data_paths)):
    print(f'Processing {name}')

    # Get filtered song (prior to envelope extraction):
    data, config = load_data(data_path, files='filt')
    song, rate = data['filt'], config['rate']

    # Get song segment to be analyzed:
    time = np.arange(song.shape[0]) / rate
    start, end = data['songs_0'].ravel()
    segment = (time >= start) & (time <= end)

    # Normalize song component:
    song /= song[segment].std()

    # Rescale song component:
    mix = song[:, None] * scales[None, :]

    if add_noise:
        # Add normalized envelopenoise:
        rng = np.random.default_rng()
        noise = rng.normal(scale=1, size=song.shape)
        noise /= noise[segment].std()
        mix += noise[:, None]

    # Process mixture:
    mix = sosfilter(np.abs(mix), rate, config['env_fcut'], 'lp',
                    padtype='even', padlen=config['padlen'])
    mix_log = decibel(mix, ref=1)
    mix_inv = sosfilter(mix_log, rate, config['inv_fcut'], 'hp',
                        padtype='constant', padlen=config['padlen'])

    # Get intensity measure per stage:
    measure_env = mix[segment, :].std(axis=0)
    measure_log = mix_log[segment, :].std(axis=0)
    measure_inv = mix_inv[segment, :].std(axis=0)

    # Save analysis results:
    save_inds = np.nonzero(np.isin(scales, example_scales))[0]
    if save_path is not None:
        data = dict(
            scales=scales,
            example_scales=example_scales,
            measure_env=measure_env,
            measure_log=measure_log,
            measure_inv=measure_inv,
            )
        if save_snippets:
            data.update(
                snip_env=mix[:, save_inds],
                snip_log=mix_log[:, save_inds],
                snip_inv=mix_inv[:, save_inds],
            )
        file_name = save_path + name
        if add_noise:
            file_name += '_noise'
        save_data(file_name, data, config, overwrite=True)
print('Done.')
embed()