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

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

# ANALYSIS SETTINGS:
add_noise = False
thresh_percent = 90
example_scales = np.array([0, 1, 10, 50])
scales = np.geomspace(0.01, 100, 100)
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}')
    save_name = save_path + name

    # Get pure-song kernel responses:
    data, config = load_data(data_path, files='conv')
    song, rate = data['conv'], data['conv_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(axis=0)

    if add_noise:
        # Get normalized noise:
        rng = np.random.default_rng()
        noise = rng.normal(size=(song.shape[0], 1))
        noise /= noise[segment].std()

        # Prepare noise-bound threshold:
        threshold = np.percentile(noise, thresh_percent, axis=0)
    else:
        # Reuse threshold from previous noise run:
        threshold = np.load(save_name + '_noise.npz')['thresh']

    # Prepare measure storage:
    shape = (scales.size, song.shape[1])
    # measure_conv = np.zeros(shape, dtype=float)
    measure_feat = np.zeros(shape, dtype=float)

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

        # Rescale song component:
        scaled_conv = song * scale
        if add_noise:
            # Add noise:
            scaled_conv += noise

        # Process mixture:
        scaled_bi = (scaled_conv > threshold).astype(float)
        scaled_feat = sosfilter(scaled_bi, rate, config['feat_fcut'], 'lp',
                                padtype='fixed', padlen=config['padlen'])

        # Get intensity measure per stage:
        # measure_conv[i] = scaled_conv[segment, :].std(axis=0)
        measure_feat[i] = scaled_feat[segment, :].mean(axis=0)

    # Save analysis results:
    if save_path is not None:
        data = dict(
            scales=scales,
            # measure_conv=measure_conv,
            measure_feat=measure_feat,
            thresh=threshold,
            thresh_perc=thresh_percent,
            )
        if add_noise:
            save_name += '_noise'
        save_data(save_name, data, config, overwrite=True)
print('Done.')
embed()