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

# GENERAL SETTINGS:
example_file = 'Omocestus_rufipes_DJN_32-40s724ms-48s779ms'
search_target = ['*', example_file][0]
data_paths = search_files(search_target, excl='noise', dir='../data/processed/')
noise_path = '../data/processed/white_noise_sd-1.npz'
save_path = '../data/inv/rect_lp/'

# ANALYSIS SETTINGS:
mode = ['pure', 'noise'][1]
example_scales = np.array([0.1, 0.3, 1, 3, 10])
scales = np.geomspace(0.01, 100, 1000)
scales = np.unique(np.concatenate(([0], scales, example_scales)))
cutoffs = np.array([np.nan, 2500, 250, 25])

# PREPARATION:
if mode == 'noise':
    pure_noise = np.load(noise_path)['raw']

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

    # Get filtered song (prior to envelope extraction):
    data, config = load_data(data_path, files='raw')
    song, rate = data['raw'], 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 mode == 'noise':
        # Get normalized noise component:
        noise = draw_noise_segment(pure_noise, song.shape[0])
        noise /= noise[segment].std()

    # Prepare storage:
    measure_filt = np.zeros_like(scales)
    measure_env = np.zeros((scales.size, len(cutoffs)), dtype=float)
    if save_detailed:
        # Prepare optional storage:
        shape = (song.shape[0], example_scales.size)
        snip_raw = np.zeros(shape)
        snip_filt = np.zeros(shape)
        snip_env = np.zeros(shape + (len(cutoffs),))

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

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

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

        # Store non-envelope results:
        measure_filt[i] = mix[segment].std()
        if save_detailed and scale in example_scales:
            scale_ind = np.nonzero(example_scales == scale)[0][0]
            snip_raw[:, scale_ind] = mix
            snip_filt[:, scale_ind] = mix

        # Process piecewise again:
        for j, cutoff in enumerate(cutoffs):
            if np.isnan(cutoff):
                mix_env = mix_rect
            else:
                mix_env = sosfilter(mix_rect, rate, cutoff, 'lp',
                                    padtype='even', padlen=config['padlen'])

            # Store envelope results:
            measure_env[i, j] = mix_env[segment].std()
            if save_detailed and scale in example_scales:
                snip_env[:, scale_ind, j] = mix_env

    # Save analysis results:
    if save_path is not None:
        archive = dict(
            scales=scales,
            example_scales=example_scales,
            cutoffs=cutoffs,
            measure_filt=measure_filt,
            measure_env=measure_env,
        )
        if save_detailed:
            archive.update(
                snip_raw=snip_raw,
                snip_filt=snip_filt,
                snip_env=snip_env,
            )
        save_name = save_path + name + '_' + mode
        save_data(save_name, archive, config, overwrite=True)

print('Done.')
embed()