from typing import List
from typing import Union

import numpy as np


def strictly_decreasing(list_to_check: list) -> bool:
    return all(x > y for x, y in zip(list_to_check, list_to_check[1:]))


def generate_truth(mixdb: dict,
                   record: dict,
                   target: np.ndarray,
                   noise: np.ndarray,
                   compute: bool = True) -> np.ndarray:
    from copy import deepcopy

    from sonusai.mixture.class_count import get_class_count
    from sonusai.mixture.get_class_weights_threshold import get_class_weights_threshold

    if not compute:
        return np.empty(0, dtype=np.single)

    truth = np.zeros((len(target), mixdb['num_classes']), dtype=np.single)
    for truth_setting in mixdb['targets'][record['target_file_index']]['truth_settings']:
        config = deepcopy(truth_setting)
        config['frame_size'] = mixdb['frame_size']
        config['num_classes'] = mixdb['num_classes']
        config['mutex'] = mixdb['truth_mutex']
        config['target_gain'] = record['target_gain']
        new_truth = truth_function(target=target,
                                   noise=noise,
                                   config=config)
        truth = truth + new_truth

    indices = [sub['index'] for sub in mixdb['targets'][record['target_file_index']]['truth_settings']]
    record['class_count'] = get_class_count(
        truth_index=indices,
        truth=truth,
        class_weights_threshold=get_class_weights_threshold(mixdb))

    return truth


def truth_function(target: np.ndarray,
                   noise: np.ndarray,
                   config: dict) -> np.ndarray:
    import h5py
    from pyaaware import ForwardTransform
    from pyaaware import SED

    from sonusai import SonusAIError
    from sonusai.utils import int16_to_float

    if config['function'] == 'sed':
        if config['target_gain'] == 0:
            return np.zeros((len(target), config['num_classes']), dtype=np.single)

        if len(target) % config['frame_size'] != 0:
            raise SonusAIError(f'Number of samples in audio is not a multiple of {config["frame_size"]}')

        if 'config' not in config:
            raise SonusAIError('Truth function SED missing config')

        parameters = ['thresholds']
        for parameter in parameters:
            if parameter not in config['config']:
                raise SonusAIError(f'Truth function SED config missing required parameter: {parameter}')

        thresholds = config['config']['thresholds']
        if not strictly_decreasing(thresholds):
            raise SonusAIError(f'Truth function SED thresholds are not strictly decreasing: {thresholds}')

        fft = ForwardTransform(N=config['frame_size'] * 4, R=config['frame_size'])
        sed = SED(thresholds=thresholds,
                  index=config['index'],
                  frame_size=config['frame_size'],
                  num_classes=config['num_classes'],
                  mutex=config['mutex'])

        audio = np.int16(np.single(target) / config['target_gain'])
        truth = np.empty((0, config['num_classes']), dtype=np.single)
        for offset in range(0, len(audio), config['frame_size']):
            new_truth = sed.execute(fft.energy(int16_to_float(audio[offset:offset + config['frame_size']])))
            truth = np.vstack((truth, np.reshape(new_truth, (1, len(new_truth)))))

        truth = truth.repeat(config['frame_size'], axis=0)
        return truth

    elif config['function'] == 'file':
        if 'config' not in config:
            raise SonusAIError('Truth function file missing config')

        parameters = ['file']
        for parameter in parameters:
            if parameter not in config['config']:
                raise SonusAIError(f'Truth function file config missing required parameter: {parameter}')

        with h5py.File(name=config['config']['file'], mode='r') as f:
            truth_in = f['/truth_t'][:]

        if truth_in.ndim != 2:
            raise SonusAIError('Truth file data is not 2 dimensions')

        if truth_in.shape[0] != len(target):
            raise SonusAIError('Truth file does not contain the right amount of samples')

        truth = np.zeros((len(target), config['num_classes']), dtype=np.single)
        if config['target_gain'] == 0:
            return truth

        if isinstance(config['index'], list):
            if len(config['index']) != truth_in.shape[1]:
                raise SonusAIError('Truth file does not contain the right amount of classes')

            truth[:, config['index']] = truth_in
        else:
            if config['index'] + truth_in.shape[1] > config['num_classes']:
                raise SonusAIError('Truth file contains too many classes')

            truth[:, config['index']:config['index'] + truth_in.shape[1]] = truth_in

        return truth

    elif config['function'] == 'snr':
        truth = np.zeros((len(target), config['num_classes']), dtype=np.single)
        fft = ForwardTransform(N=config['frame_size'] * 4, R=config['frame_size'])
        for offset in range(0, len(target), config['frame_size']):
            target_energy = fft.energy(int16_to_float(target[offset:offset + config['frame_size']]))
            noise_energy = fft.energy(int16_to_float(noise[offset:offset + config['frame_size']]))
            truth[offset:offset + config['frame_size'], config['index']] = np.single(target_energy / noise_energy)

        return truth

    elif config['function'] == 'phoneme':
        # Read in .txt transcript and run a Python function to generate text grid data
        # (indicating which phonemes are active)
        # Then generate truth based on this data and put in the correct classes based on config['index']
        raise SonusAIError('Truth function phoneme is not supported yet')

    raise SonusAIError(f'Unsupported truth function: {config["function"]}')


def get_truth_indices(mixdb: dict, mixid: Union[str, List[int]]) -> list:
    from sonusai.mixture.get_mixtures_from_mixid import convert_mixid_to_list

    indices = list()

    for m in convert_mixid_to_list(mixdb, mixid):
        target_file_index = mixdb['mixtures'][m]['target_file_index']

        # extract indices from truth_settings
        index = [sub['index'] for sub in mixdb['targets'][target_file_index]['truth_settings']]

        # flatten, uniquify, and sort
        index = sorted(list(set([item for sublist in index for item in sublist])))

        indices.append(index)

    return indices
