import numpy as np
from thunderhopper.filetools import search_files
from thunderhopper.modeltools import load_data, save_data
from misc_functions import get_saturation
from IPython import embed

# GENERAL SETTINGS:
target_species = [
    'Chorthippus_biguttulus',
    'Chorthippus_mollis',
    'Chrysochraon_dispar',
    'Euchorthippus_declivus',
    'Gomphocerippus_rufus',
    'Omocestus_rufipes',
    'Pseudochorthippus_parallelus',
]
search_path = '../data/inv/log_hp/collected/'
save_path = '../data/inv/log_hp/saturation/'

# ANALYSIS SETTINGS:
plateau_settings = dict(
    low=0.05,
    high=0.95,
    first=True,
    last=True,
    condense=None,
)
compute_hist = True
bins = 50
pad = 0.05

# PREPARATION:
if compute_hist:
    species_scales = []
    min_scale, max_scale = [], []
    archives = [{} for _ in target_species]

# EXECUTION:
for i, species in enumerate(target_species):
    print(f'Processing {species}')

    # Load accumulated invariance data:
    path = search_files(species, dir=search_path)[0]
    data, config = load_data(path, ['scales', 'measure_inv'])

    # Find upper saturation point per song file:
    crit_inds = np.array(get_saturation(data['measure_inv'], **plateau_settings)[1])
    crit_scales = data['scales'][crit_inds]

    # Output options:
    if not compute_hist:
        # Save species data immediately:
        archive = dict(crit_inds=crit_inds, crit_scales=crit_scales, scales=data['scales'])
        save_data(save_path + species, archive, config, overwrite=True)
        continue

    # Log but don't save data yet:
    archives[i]['crit_inds'] = crit_inds
    archives[i]['crit_scales'] = crit_scales
    archives[i]['scales'] = data['scales']
    min_scale.append(crit_scales.min())
    max_scale.append(crit_scales.max())

# Optional histogram:
if compute_hist:
    # Generated shared histogram edges:
    min_scale, max_scale = min(min_scale), max(max_scale)
    pad *= (max_scale - min_scale)
    edges = np.linspace(max(0, min_scale - pad), max_scale + pad, bins + 1)
    centers = edges[:-1] + np.diff(edges) / 2

    # Compute histogram and save species data:
    for i, (species, archive) in enumerate(zip(target_species, archives)):
        hist = np.histogram(archive['crit_scales'], bins=edges, density=True)[0]
        archive['hist'] = hist
        archive['bins'] = centers
        save_data(save_path + species, archive, config, overwrite=True)

print('Done.')