import plotstyle_plt
import glob
import numpy as np
import matplotlib.pyplot as plt
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:
scales = np.arange(0, 10.1, 0.1)
save_scales = np.array([0, 0.5, 1, 2, 5, 10])
floor_percents = dict(song=100, noise=99)
single_db_ref = True
plot_redundant = False

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

    # Load song envelope:
    data, config = load_data(data_path, files='env')
    song, rate = data['env'], config['env_rate']
    t_full = np.arange(song.shape[0]) / rate

    # Generate noise envelope:
    rng = np.random.default_rng()
    noise = rng.normal(size=song.shape)
    noise = extract_env(noise, rate, config=config)

    # Normalize each:
    song /= song.std()
    noise /= noise.std()

    # Mix song component and noise component:
    scaled = song[:, None] * scales[None, :]
    mix = scaled + noise[:, None]

    # Find noise floor (experimental):
    song_percent = np.percentile(scaled, floor_percents['song'], axis=0)
    noise_percent = np.percentile(noise, floor_percents['noise'])
    floor_scale = scales[np.nonzero(song_percent <= noise_percent)[0][-1]]

    # 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 variances per stage:
    var_env = mix.var(axis=0)
    var_log = mix_log.var(axis=0)
    var_inv = mix_inv.var(axis=0)

    # Get SNRs against pure noise per stage:
    base_ind = np.nonzero(scales == 0)[0][0]
    snr_env = var_env / var_env[base_ind]
    snr_log = var_log / var_log[base_ind]
    snr_inv = var_inv / var_inv[base_ind]

    if plot_redundant:
        # Normalize SNRs:
        norm_snr_env = snr_env / snr_env.max()
        norm_snr_log = snr_log / snr_log.max()
        norm_snr_inv = snr_inv / snr_inv.max()

        # Get SNR gain against env:
        gain_log = snr_log / snr_env
        gain_inv = snr_inv / snr_env

        # Plot results:
        fig, axes = plt.subplots(6, 1, sharex=True, layout='constrained')
        
        axes[0].set_title('variance')
        axes[0].plot(scales, var_env)
        axes[0].plot(scales, var_log)
        axes[0].plot(scales, var_inv)

        axes[1].set_title('normalized variance')
        axes[1].plot(scales, var_env / var_env.max())
        axes[1].plot(scales, var_log / var_log.max())
        axes[1].plot(scales, var_inv / var_inv.max())

        axes[2].set_title('SNR')
        axes[2].plot(scales, snr_env)
        axes[2].plot(scales, snr_log)
        axes[2].plot(scales, snr_inv)

        axes[3].set_title('normalized SNR')
        axes[3].plot(scales, norm_snr_env)
        axes[3].plot(scales, norm_snr_log)
        axes[3].plot(scales, norm_snr_inv)

        axes[4].set_title('gain (absolute SNR)')
        axes[4].plot(scales, gain_log)
        axes[4].plot(scales, gain_inv)

        axes[5].set_title('gain (normalized SNR)')
        axes[5].plot(scales, norm_snr_log / norm_snr_env)
        axes[5].plot(scales, norm_snr_inv / norm_snr_env)
        plt.show()

    # Save analysis results:
    save_inds = np.isin(scales, save_scales)
    if save_path is not None:
        data = dict(
            scales=scales,
            plot_scales=scales[save_inds],
            floor_scale=floor_scale,
            env=mix[:, save_inds],
            log=mix_log[:, save_inds],
            inv=mix_inv[:, save_inds],
            snr_env=snr_env,
            snr_log=snr_log,
            snr_inv=snr_inv,
            )
        save_data(save_path + name, data, config, overwrite=True)
print('Done.')
embed()