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 = False
single_db_ref = True
# find_saturation = add_noise and False
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)))
# if find_saturation:
#     scales = np.append(scales, 10e10)

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

    # Get normalized song envelope:
    data, config = load_data(data_path, files='env')
    song, rate = data['env'], config['env_rate']
    song /= song.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.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.std(axis=0)
    measure_log = mix_log.std(axis=0)
    measure_inv = mix_inv.std(axis=0)

    # # Find saturation level:
    # if find_saturation:
    #     limit = measure_inv[-1]
    #     scales = scales[:-1]
    #     measure_env = measure_env[:-1]
    #     measure_log = measure_log[:-1]
    #     measure_inv = measure_inv[:-1]

    # 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,
            )
        # if find_saturation:
        #     data['limit'] = limit
        file_name = save_path + name
        if add_noise:
            file_name += '_noise'
        save_data(file_name, data, config, overwrite=True)
print('Done.')
embed()