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

# GENERAL SETTINGS:
target = 'Omocestus_rufipes'
data_paths = glob.glob(f'../data/processed/{target}*.npz')
save_path = '../data/inv/log_hp/'

# ANALYSIS SETTINGS:
add_noise = True
single_db_ref = True
example_scales = np.array([0, 0.1, 1, 10, 100, 200])
scales = np.geomspace(0.1, 1000, 100)
if not add_noise:
    example_scales = example_scales[example_scales > 0]
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 song envelope:
    data, config = load_data(data_path, files='env')
    song, rate = data['env'], config['env_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 noise envelope:
        rng = np.random.default_rng()
        noise = rng.normal(size=song.shape)
        noise = extract_env(noise, rate, config=config)
        noise /= noise[segment].std()
        mix += noise[:, None]

    # Process mixture:
    mix_log = decibel(mix, axis=None if single_db_ref else 0)
    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,
            env=mix[:, save_inds],
            log=mix_log[:, save_inds],
            inv=mix_inv[:, save_inds],
            measure_env=measure_env,
            measure_log=measure_log,
            measure_inv=measure_inv,
            )
        file_name = save_path + name
        if add_noise:
            file_name += '_noise'
        save_data(file_name, data, config, overwrite=True)
print('Done.')
embed()