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:
example_file = 'Omocestus_rufipes_DJN_32-40s724ms-48s779ms'
search_target = ['*', example_file][1]
data_paths = search_files(search_target, excl='noise', dir='../data/processed/')
noise_path = '../data/processed/white_noise_sd-1.npz'
save_path = '../data/inv/log_hp/'

# ANALYSIS SETTINGS:
add_noise = search_target == '*' or False
save_detailed = search_target == example_file
example_scales = np.array([0.1, 1, 10, 30, 100, 300])
scales = np.geomspace(0.01, 10000, 1000)
scales = np.unique(np.concatenate((scales, example_scales)))

# PREPARATION:
pure_noise = np.load(noise_path)['filt']

# 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()
    if add_noise:
        # Get normalized noise component:
        noise = pure_noise[:song.shape[0]]
        noise /= noise[segment].std()

    # Prepare storage:
    measure_inv = np.zeros_like(scales)
    if save_detailed:
        # Prepare optional storage:
        measure_env = np.zeros_like(scales)
        measure_log = np.zeros_like(scales)
        snip_env = np.zeros((song.shape[0], example_scales.size))
        snip_log = np.zeros((song.shape[0], example_scales.size))
        snip_inv = np.zeros((song.shape[0], example_scales.size))

    # Execute piecewise:
    for i, scale in enumerate(scales):

        # Get scaled mixture:
        mix = song * scale
        if add_noise:
            mix += noise

        # 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'])

        # Log intensity measures:
        measure_inv[i] = mix_inv[segment].std()
        if save_detailed:
            measure_env[i] = mix[segment].std()
            measure_log[i] = mix_log[segment].std()
            if scale in example_scales:
                # Log snippet data:
                save_ind = np.nonzero(example_scales == scale)[0][0]
                snip_env[:, save_ind] = mix
                snip_log[:, save_ind] = mix_log
                snip_inv[:, save_ind] = mix_inv

    # Save analysis results:
    if save_path is not None:
        archive = dict(
            scales=scales,
            example_scales=example_scales,
            measure_inv=measure_inv,
        )
        if save_detailed:
            archive.update(
                measure_env=measure_env,
                measure_log=measure_log,
                snip_env=snip_env,
                snip_log=snip_log,
                snip_inv=snip_inv,
            )
        file_name = save_path + name
        if add_noise:
            file_name += '_noise'
        save_data(file_name, archive, config, overwrite=True)

print('Done.')
embed()