import numpy as np
from thunderhopper.filters import decibel, sosfilter
from thunderhopper.model import convolve_kernels, process_signal
from thunderhopper.modeltools import load_data
from IPython import embed

## SETTINGS:

# General:
mode = ['log_hp', 'thresh_lp', 'full', 'short'][3]
noise_path = '../data/processed/white_noise_sd-1.npz'
save_path = '../data/inv/'
pad = np.array([0.1, 0.9])

stages = dict(
    log_hp=['filt', 'env', 'log', 'inv'],
    thresh_lp=['inv', 'conv', 'feat'],
    full=['raw', 'filt', 'env', 'log', 'inv', 'conv', 'feat'],
    short=['raw', 'filt', 'env', 'conv', 'feat']
)[mode]

# PROCESSING:

print(f'Fetching references for {mode} invariance...')

# Load pure-noise starter representation:
noise_data, config = load_data(noise_path, stages[0])
starter = noise_data[stages[0]]

# Prepare buffered measurement segment:
pad = (pad * starter.shape[0]).astype(int)
segment = np.arange(starter.shape[0])[pad[0]:pad[1]]

# Normalize starter:
starter /= starter[segment].std()

# Run pipeline:
if mode == 'log_hp':
    data = {'filt': starter}
    data['env'] = sosfilter(np.abs(data['filt']), config['rate'], config['env_fcut'],
                            'lp', padtype='even', padlen=config['padlen'])
    data['log'] = decibel(data['env'], ref=1)
    data['inv'] = sosfilter(data['log'], config['env_rate'], config['inv_fcut'],
                            'hp', padtype='constant', padlen=config['padlen'])
elif mode == 'thresh_lp':
    data = {'inv': starter}
    data['conv'] = convolve_kernels(data['inv'], config['kernels'], config['k_specs'])
    data['feat'] = sosfilter((data['conv'] > config['feat_thresh']).astype(float),
                             config['env_rate'], config['feat_fcut'], 'lp',
                             padtype='fixed', padlen=config['padlen'])
elif mode == 'full':
    data = process_signal(config, stages, signal=starter, rate=config['rate'])[0]
elif mode == 'short':
    data = process_signal(config, ['raw', 'filt', 'env'], signal=starter, rate=config['rate'])[0]
    data['conv'] = convolve_kernels(data['env'], config['kernels'], config['k_specs'])
    data['feat'] = sosfilter((data['conv'] > config['feat_thresh']).astype(float),
                             config['env_rate'], config['feat_fcut'], 'lp',
                             padtype='fixed', padlen=config['padlen'])

# Get measures:
measures = {}
for stage in stages:
    if stage == 'feat':
        measures[stage] = data[stage][segment, :].mean(axis=0)
    else:
        measures[stage] = data[stage][segment, ...].std(axis=0)

# Save results:
np.savez(save_path + f'{mode}/ref_measures.npz', **measures)

print('Done.')
embed()