from dataclasses import dataclass
from typing import Any
from typing import List
from typing import Tuple
from typing import Union

from pyaaware import FeatureGenerator

from sonusai.mixture.types import AudioF
from sonusai.mixture.types import AudioT
from sonusai.mixture.types import AudiosF
from sonusai.mixture.types import AudiosT
from sonusai.mixture.types import Augmentation
from sonusai.mixture.types import Augmentations
from sonusai.mixture.types import ClassCount
from sonusai.mixture.types import DataClassSonusAIMixin
from sonusai.mixture.types import Feature
from sonusai.mixture.types import Location
from sonusai.mixture.types import MixtureIDs
from sonusai.mixture.types import NoiseFiles
from sonusai.mixture.types import Optional
from sonusai.mixture.types import Segsnr
from sonusai.mixture.types import TargetFiles
from sonusai.mixture.types import Truth
from sonusai.mixture.types import TruthSettings


@dataclass
class MRecord(DataClassSonusAIMixin):
    name: str = None
    noise_augmentation_index: int = None
    noise_file_index: int = None
    noise_offset: int = None
    noise_snr_gain: float = None
    random_snr: Optional[bool] = None
    samples: int = None
    snr: float = None
    target_augmentation_index: List[int] = None
    target_file_index: List[int] = None
    target_gain: List[int] = None
    target_snr_gain: float = None


MRecords = List[MRecord]


@dataclass
class MixtureDatabaseConfig(DataClassSonusAIMixin):
    class_balancing: Optional[bool] = False
    class_balancing_augmentation: Optional[Augmentation] = None
    class_labels: List[str] = None
    class_weights_threshold: List[float] = None
    exhaustive_noise: Optional[bool] = True
    feature: str = None
    first_cba_index: Optional[int] = None
    mixtures: MRecords = None
    noise_augmentations: Augmentations = None
    noises: NoiseFiles = None
    num_classes: int = None
    random_snrs: Optional[List[str]] = None
    seed: Optional[int] = 0
    snrs: List[float] = None
    target_augmentations: Augmentations = None
    targets: TargetFiles = None
    truth_mutex: bool = None
    truth_reduction_function: str = None
    truth_settings: TruthSettings = None


@dataclass
class TransformConfig:
    N: int
    R: int
    ttype: str


class MixtureDatabase:
    from sonusai.mixture.types import AudiosT
    from sonusai.mixture.types import ListAudiosT

    def __init__(self, config: Union[Location, MixtureDatabaseConfig]):
        if isinstance(config, MixtureDatabaseConfig):
            self.config = config
            self.location = None
        else:
            self.load(config)
            self.location = config

        self.fg = FeatureGenerator(feature_mode=self.feature,
                                   num_classes=self.num_classes,
                                   truth_mutex=self.truth_mutex)

        self.ft_config = TransformConfig(N=self.fg.ftransform_N,
                                         R=self.fg.ftransform_R,
                                         ttype=self.fg.ftransform_ttype)

        self.eft_config = TransformConfig(N=self.fg.eftransform_N,
                                          R=self.fg.eftransform_R,
                                          ttype=self.fg.eftransform_ttype)

        self.it_config = TransformConfig(N=self.fg.itransform_N,
                                         R=self.fg.itransform_R,
                                         ttype=self.fg.itransform_ttype)

        self._raw_target_audios = None
        self._augmented_noise_audios = None

    def load(self, location: Location) -> None:
        import json
        from os.path import exists
        from os.path import isdir
        from os.path import join

        from sonusai import SonusAIError

        if not isdir(location):
            raise SonusAIError(f'{location} is not a directory')

        filename = join(location, 'mixdb.json')
        if not exists(filename):
            raise SonusAIError(f'could not find mixture database in {location}')

        with open(file=filename, mode='r', encoding='utf-8') as f:
            self.config = MixtureDatabaseConfig.from_dict(json.loads(f.read()))

    def save(self):
        from os.path import join

        json_name = join(self.location, 'mixdb.json')
        with open(file=json_name, mode='w') as file:
            file.write(self.config.to_json(indent=2))

    def forward_transform(self, audio: AudioT) -> AudioF:
        from pyaaware import ForwardTransform

        from sonusai.utils import int16_to_float

        ftf = ForwardTransform(N=self.ft_config.N, R=self.ft_config.R, ttype=self.ft_config.ttype)
        return ftf.execute_all(int16_to_float(audio)).transpose()

    def inverse_transform(self, transform: AudioF, trim: bool = True) -> AudioT:
        from pyaaware import InverseTransform

        itf = InverseTransform(N=self.it_config.N, R=self.it_config.R, ttype=self.it_config.ttype)
        audio = itf.execute_all(transform.transpose())
        if trim:
            audio = audio[itf.N - itf.R:]

        return audio

    @property
    def class_balancing(self):
        return self.config.class_balancing

    @class_balancing.setter
    def class_balancing(self, value):
        self.config.class_balancing = value

    @property
    def class_balancing_augmentation(self):
        return self.config.class_balancing_augmentation

    @class_balancing_augmentation.setter
    def class_balancing_augmentation(self, value):
        self.config.class_balancing_augmentation = value

    @property
    def class_labels(self):
        return self.config.class_labels

    @class_labels.setter
    def class_labels(self, value):
        from sonusai import SonusAIError

        if value is not None and (not isinstance(value, list) or len(value) != self.num_classes):
            raise SonusAIError(f'invalid class_labels length')
        self.config.class_labels = value

    @property
    def class_weights_threshold(self):
        return self.config.class_weights_threshold

    @class_weights_threshold.setter
    def class_weights_threshold(self, value):
        from sonusai import SonusAIError

        if len(value) not in [1, self.num_classes]:
            raise SonusAIError(f'invalid class_weights_threshold length: {len(class_weights_threshold)}')
        self.config.class_weights_threshold = value

    @property
    def exhaustive_noise(self):
        return self.config.exhaustive_noise

    @exhaustive_noise.setter
    def exhaustive_noise(self, value):
        self.config.exhaustive_noise = value

    @property
    def feature(self):
        return self.config.feature

    @feature.setter
    def feature(self, value):
        self.config.feature = value

    @property
    def transform_frame_ms(self):
        from sonusai.mixture import SAMPLE_RATE

        return float(self.ft_config.R) / float(SAMPLE_RATE / 1000)

    @property
    def feature_ms(self):
        return self.transform_frame_ms * self.fg.decimation * self.fg.stride

    @property
    def feature_samples(self):
        return self.ft_config.R * self.fg.decimation * self.fg.stride

    @property
    def feature_step_ms(self):
        return self.transform_frame_ms * self.fg.decimation * self.fg.step

    @property
    def feature_step_samples(self):
        return self.ft_config.R * self.fg.decimation * self.fg.step

    @property
    def first_cba_index(self):
        return self.config.first_cba_index

    @first_cba_index.setter
    def first_cba_index(self, value):
        self.config.first_cba_index = value

    @property
    def mixtures(self):
        return self.config.mixtures

    @mixtures.setter
    def mixtures(self, value):
        self.config.mixtures = value

    @property
    def noise_augmentations(self):
        return self.config.noise_augmentations

    @noise_augmentations.setter
    def noise_augmentations(self, value):
        self.config.noise_augmentations = value

    @property
    def noises(self):
        return self.config.noises

    @noises.setter
    def noises(self, value):
        self.config.noises = value

    @property
    def num_classes(self):
        return self.config.num_classes

    @num_classes.setter
    def num_classes(self, value):
        self.config.num_classes = value

    @property
    def random_snrs(self):
        return self.config.random_snrs

    @random_snrs.setter
    def random_snrs(self, value):
        self.config.random_snrs = value

    @property
    def seed(self):
        return self.config.seed

    @seed.setter
    def seed(self, value):
        self.config.seed = value

    @property
    def snrs(self):
        return self.config.snrs

    @snrs.setter
    def snrs(self, value):
        self.config.snrs = value

    @property
    def target_augmentations(self):
        return self.config.target_augmentations

    @target_augmentations.setter
    def target_augmentations(self, value):
        self.config.target_augmentations = value

    @property
    def targets(self):
        return self.config.targets

    @targets.setter
    def targets(self, value):
        self.config.targets = value

    @property
    def truth_mutex(self):
        return self.config.truth_mutex

    @truth_mutex.setter
    def truth_mutex(self, value):
        self.config.truth_mutex = value

    @property
    def truth_reduction_function(self):
        return self.config.truth_reduction_function

    @truth_reduction_function.setter
    def truth_reduction_function(self, value):
        self.config.truth_reduction_function = value

    @property
    def truth_settings(self):
        return self.config.truth_settings

    @truth_settings.setter
    def truth_settings(self, value):
        self.config.truth_settings = value

    @property
    def raw_target_audios(self) -> AudiosT:
        return self._raw_target_audios

    def load_raw_target_audios(self, show_progress: bool = False) -> None:
        from tqdm import tqdm

        from sonusai.mixture import read_audio
        from sonusai.utils import p_tqdm_map

        names = [target.name for target in self.targets]
        progress = tqdm(total=len(names), desc='Read target audio', disable=not show_progress)
        self._raw_target_audios = p_tqdm_map(read_audio, names, progress=progress)
        progress.close()

    @property
    def augmented_noise_audios(self) -> ListAudiosT:
        return self._augmented_noise_audios

    def load_augmented_noise_audios(self, show_progress: bool = False) -> None:
        from functools import partial

        from tqdm import tqdm

        from sonusai.mixture import read_audio
        from sonusai.utils import p_tqdm_map

        names = [noise.name for noise in self.noises]

        progress = tqdm(total=len(names), desc='Read noise audio', disable=not show_progress)
        raw_noise_audios = p_tqdm_map(read_audio, names, progress=progress)
        progress.close()

        progress = tqdm(total=len(names), desc='Augment noise audio', disable=not show_progress)
        self._augmented_noise_audios = p_tqdm_map(partial(_augment_audio, augmentations=self.noise_augmentations),
                                                  raw_noise_audios,
                                                  progress=progress)
        progress.close()

    def total_samples(self, mixids: MixtureIDs = '*') -> int:
        return sum([self.mixture_samples(mixid) for mixid in self.mixids_to_list(mixids)])

    def mixture_samples(self, mixid: int) -> int:
        return self.mixtures[mixid].samples

    def mixture_transform_frames(self, mixid: int) -> int:
        return self.mixtures[mixid].samples // self.ft_config.R

    def mixture_feature_frames(self, mixid: int) -> int:
        return self.mixtures[mixid].samples // self.feature_step_samples

    def mixids_to_list(self, mixids: MixtureIDs = None) -> List[int]:
        from sonusai import SonusAIError

        all_mixids = list(range(len(self.mixtures)))

        if mixids is None:
            return all_mixids

        if isinstance(mixids, str):
            if mixids == '*':
                return all_mixids

            try:
                return eval(f'{all_mixids}[{mixids}]')
            except NameError:
                raise SonusAIError(f'Empty mixids {mixids}')

        if isinstance(mixids, range):
            result = list(mixids)
        elif isinstance(mixids, int):
            result = [mixids]
        else:
            result = mixids

        if not all(isinstance(x, int) and 0 <= x < len(self.mixtures) for x in result):
            raise SonusAIError(f'Invalid entries in mixids of {mixids}')

        if not result:
            raise SonusAIError(f'Empty mixids {mixids}')

        return result

    def mixture_metadata(self, mixid: int) -> str:
        mrecord = self.mixtures[mixid]
        metadata = ''
        for ti in range(len(mrecord.target_file_index)):
            tfi = mrecord.target_file_index[ti]
            tai = mrecord.target_augmentation_index[ti]
            metadata += f'target {ti} name: {self.targets[tfi].name}\n'
            metadata += f'target {ti} augmentation: {self.target_augmentations[tai].to_dict()}\n'
            metadata += f'target {ti} target_gain: {mrecord.target_gain[ti]}\n'
            truth_settings = self.targets[tfi].truth_settings
            for tsi in range(len(truth_settings)):
                metadata += f'target {ti} truth index {tsi}: {truth_settings[tsi].index}\n'
                metadata += f'target {ti} truth function {tsi}: {truth_settings[tsi].function}\n'
                metadata += f'target {ti} truth config {tsi}: {truth_settings[tsi].config}\n'
        metadata += f'noise name: {self.noises[mrecord.noise_file_index].name}\n'
        metadata += f'noise augmentation: {self.noise_augmentations[mrecord.noise_augmentation_index].to_dict()}\n'
        metadata += f'snr: {mrecord.snr}\n'
        metadata += f'random_snr: {mrecord.random_snr}\n'
        metadata += f'samples: {mrecord.samples}\n'
        metadata += f'target_snr_gain: {mrecord.target_snr_gain}\n'
        metadata += f'noise_snr_gain: {mrecord.noise_snr_gain}\n'

        return metadata

    def write_mixture_metadata(self, mixid: int) -> None:
        from os.path import splitext

        with open(file=splitext(self.mixture_filename(mixid))[0] + '.txt', mode='w') as f:
            f.write(self.mixture_metadata(mixid))

    def mixture_filename(self, mixid: int) -> Location:
        from os.path import join

        return join(self.location, self.mixtures[mixid].name) if self.location is not None else None

    def check_audio_files_exist(self) -> None:
        """Walk through all the noise and target audio files in a mixture database ensuring that they exist."""
        from os.path import exists

        from sonusai import SonusAIError
        from sonusai.mixture import tokenized_expandvars

        for file_index in range(len(self.noises)):
            file_name, _ = tokenized_expandvars(self.noises[file_index].name)
            if not exists(file_name):
                raise SonusAIError(f'Could not find {file_name}')

        for file_index in range(len(self.targets)):
            file_name, _ = tokenized_expandvars(self.targets[file_index].name)
            if not exists(file_name):
                raise SonusAIError(f'Could not find {file_name}')

    def read_mixture_data(self, mixid: int, items: Union[List[str], str]) -> Any:
        from os.path import exists

        import h5py
        import numpy as np

        def _get_dataset(file: h5py.File, name: str) -> Any:
            if name in file:
                return np.array(f[name])
            return None

        if not isinstance(items, list):
            items = [items]

        if exists(self.mixture_filename(mixid)):
            with h5py.File(self.mixture_filename(mixid), 'r') as f:
                result = ([_get_dataset(f, item) for item in items])
        else:
            result = ([None for item in items])

        if len(items) == 1:
            result = result[0]

        return result

    def write_mixture_data(self, mixid: int, items: Union[List[Tuple[str, Any]], Tuple[str, Any]]) -> None:
        import h5py

        if not isinstance(items, list):
            items = [items]

        with h5py.File(self.mixture_filename(mixid), 'a') as f:
            for item in items:
                if item[0] in f:
                    del f[item[0]]
                f.create_dataset(name=item[0], data=item[1])

    def mixture_targets(self, mixid: int, force: bool = False) -> AudiosT:
        """ Get the list of augmented target audio (one per target in the mixup) for the given mixid

        mixid
            Mixture ID.
        force
            Force computing data from original sources regardless of whether or not cached data exists.
        """
        from sonusai.mixture import apply_augmentation
        from sonusai.mixture import apply_gain
        from sonusai.mixture import pad_audio_to_length
        from sonusai.mixture import read_audio

        if not force:
            targets = self.read_mixture_data(mixid, 'targets')
            if targets is not None:
                return list(targets)

        mrecord = self.mixtures[mixid]
        targets = []
        for idx in range(len(mrecord.target_file_index)):
            if self.raw_target_audios is not None:
                target = self.raw_target_audios[mrecord.target_file_index[idx]]
            else:
                target = read_audio(self.targets[mrecord.target_file_index[idx]].name)

            target = apply_augmentation(audio=target,
                                        augmentation=self.target_augmentations[mrecord.target_augmentation_index[idx]],
                                        length_common_denominator=self.feature_step_samples)

            target = apply_gain(audio=target, gain=mrecord.target_snr_gain)
            target = pad_audio_to_length(audio=target, length=mrecord.samples)
            targets.append(target)

        return targets

    def mixture_targets_f(self, mixid: int, targets: AudiosT = None) -> AudiosF:
        """ Get the list of augmented target transform (one per target in the mixup) for the given mixid

        mixid
            Mixture ID.
        targets
            List of augmented target audio (one per target in the mixup) for the given mixid. (optional)
        """
        if targets is None:
            targets = self.mixture_targets(mixid=mixid)

        return [self.forward_transform(target) for target in targets]

    def mixture_target(self, mixid: int, targets: AudiosT = None, force: bool = False) -> AudioT:
        """ Get the augmented target audio for the given mixid

        mixid
            Mixture ID.
        targets
            List of augmented target audio (one per target in the mixup) for the given mixid. (optional)
        force
            Force computing data from original sources regardless of whether or not cached data exists.
        """
        if not force and targets is not None:
            return sum(targets)

        return sum(self.mixture_targets(mixid=mixid, force=force))

    def mixture_target_f(self, mixid: int, targets: AudiosT = None, target: AudioT = None) -> AudioF:
        """ Get the augmented target transform for the given mixid

        mixid
            Mixture ID.
        targets
            List of augmented target audio (one per target in the mixup) for the given mixid. (optional)
        target
            Augmented target audio for the given mixid. (optional)
        """
        if target is None:
            target = self.mixture_target(mixid=mixid, targets=targets)

        return self.forward_transform(target)

    def mixture_noise(self, mixid: int, force: bool = False) -> AudioT:
        """ Get the augmented noise audio for the given mixid

        mixid
            Mixture ID.
        force
            Force computing data from original sources regardless of whether or not cached data exists.
        """
        from sonusai.mixture import apply_augmentation
        from sonusai.mixture import apply_gain
        from sonusai.mixture import get_next_noise
        from sonusai.mixture import read_audio

        if not force:
            noise = self.read_mixture_data(mixid, 'noise')
            if noise is not None:
                return noise

        mrecord = self.mixtures[mixid]
        if self.augmented_noise_audios is not None:
            noise = self.augmented_noise_audios[mrecord.noise_file_index][mrecord.noise_augmentation_index]
        else:
            noise = apply_augmentation(audio=read_audio(self.noises[mrecord.noise_file_index].name),
                                       augmentation=self.noise_augmentations[mrecord.noise_augmentation_index])

        noise = get_next_noise(audio=noise, offset=mrecord.noise_offset, length=mrecord.samples)
        noise = apply_gain(audio=noise, gain=mrecord.noise_snr_gain)

        return noise

    def mixture_noise_f(self, mixid: int, noise: AudioT = None) -> AudioF:
        """ Get the augmented noise transform for the given mixid

        mixid
            Mixture ID.
        noise
            Augmented noise audio for the given mixid. (optional)
        """
        if noise is None:
            noise = self.mixture_noise(mixid=mixid)

        return self.forward_transform(noise)

    def mixture_mixture(self,
                        mixid: int,
                        targets: AudiosT = None,
                        target: AudioT = None,
                        noise: AudioT = None,
                        force: bool = False) -> AudioT:
        """ Get the mixture audio for the given mixid

        mixid
            Mixture ID.
        targets
            List of augmented target audio (one per target in the mixup) for the given mixid. (optional)
        target
            Augmented target audio for the given mixid. (optional)
        noise
            Augmented noise audio for the given mixid. (optional)
        force
            Force computing data from original sources regardless of whether or not cached data exists.
        """
        if not force:
            mixture = self.read_mixture_data(mixid, 'mixture')
            if mixture is not None:
                return mixture

        if target is None:
            target = self.mixture_target(mixid=mixid, targets=targets)

        if noise is None:
            noise = self.mixture_noise(mixid=mixid)

        mixture = target + noise

        return mixture

    def mixture_mixture_f(self, mixid: int,
                          targets: AudiosT = None,
                          target: AudioT = None,
                          noise: AudioT = None,
                          mixture: AudioT = None) -> AudioF:
        """ Get the mixture transform for the given mixid

        mixid
            Mixture ID.
        targets
            List of augmented target audio (one per target in the mixup) for the given mixid. (optional)
        target
            Augmented target audio for the given mixid. (optional)
        noise
            Augmented noise audio for the given mixid. (optional)
        mixture
            Mixture audio for the given mixid. (optional)
        """
        if mixture is None:
            mixture = self.mixture_mixture(mixid=mixid, targets=targets, target=target, noise=noise)

        return self.forward_transform(mixture)

    def mixture_truth_t(self,
                        mixid: int,
                        targets: AudiosT = None,
                        noise: AudioT = None,
                        force: bool = False) -> Truth:
        """ Get the truth_t data for the given mixid

        mixid
            Mixture ID.
        targets
            List of augmented target audio (one per target in the mixup) for the given mixid. (optional)
        noise
            Augmented noise audio for the given mixid. (optional)
        force
            Force computing data from original sources regardless of whether or not cached data exists.
        """
        import numpy as np

        from sonusai import SonusAIError
        from sonusai.mixture import TruthFunctionConfig
        from sonusai.mixture import truth_function

        if not force:
            truth_t = self.read_mixture_data(mixid, 'truth_t')
            if truth_t is not None:
                return truth_t

        if targets is None:
            targets = self.mixture_targets(mixid=mixid)

        if noise is None:
            noise = self.mixture_noise(mixid=mixid)

        mrecord = self.mixtures[mixid]
        if len(targets) != len(mrecord.target_file_index):
            raise SonusAIError('Number of target audio entries does not match number of targets')

        if not all(len(target_audio) == len(noise) for target_audio in targets):
            raise SonusAIError('Lengths of target audio do not match length of noise audio')

        truth_t = np.zeros((self.mixture_samples(mixid), self.num_classes), dtype=np.float32)
        for idx in range(len(targets)):
            for truth_setting in self.targets[mrecord.target_file_index[idx]].truth_settings:
                config = TruthFunctionConfig(
                    feature=self.feature,
                    index=truth_setting.index,
                    function=truth_setting.function,
                    config=truth_setting.config,
                    num_classes=self.num_classes,
                    mutex=self.truth_mutex,
                    target_gain=mrecord.target_gain[idx] * mrecord.target_snr_gain
                )
                truth_t += truth_function(target_audio=targets[idx], noise_audio=noise, config=config)

        return truth_t

    def mixture_segsnr_t(self,
                         mixid: int,
                         targets: AudiosT = None,
                         target: AudioT = None,
                         noise: AudioT = None,
                         force: bool = False) -> Segsnr:
        """ Get the segsnr_t data for the given mixid

        mixid
            Mixture ID.
        targets
            List of augmented target audio (one per target in the mixup) for the given mixid. (optional)
        target
            Augmented target audio for the given mixid. (optional)
        noise
            Augmented noise audio for the given mixid. (optional)
        force
            Force computing data from original sources regardless of whether or not cached data exists.
        """
        import numpy as np

        from pyaaware import ForwardTransform

        from sonusai.utils import int16_to_float

        if not force:
            segsnr_t = self.read_mixture_data(mixid, 'segsnr_t')
            if segsnr_t is not None:
                return segsnr_t

        if target is None:
            target = self.mixture_target(mixid=mixid, targets=targets)

        if noise is None:
            noise = self.mixture_noise(mixid=mixid)

        fft = ForwardTransform(N=self.ft_config.N, R=self.ft_config.R, ttype=self.ft_config.ttype)

        segsnr_t = np.empty(self.mixture_samples(mixid), dtype=np.float32)

        frame = 0
        for offset in range(0, self.mixture_samples(mixid), self.ft_config.R):
            indices = slice(offset, offset + self.ft_config.R)

            target_energy = fft.energy_t(int16_to_float(target[indices]))
            noise_energy = fft.energy_t(int16_to_float(noise[indices]))

            if noise_energy == 0:
                snr = np.float32(np.inf)
            else:
                snr = np.float32(target_energy / noise_energy)

            segsnr_t[indices] = snr

        return segsnr_t

    def mixture_segsnr(self,
                       mixid: int,
                       segsnr_t: Segsnr = None,
                       targets: AudiosT = None,
                       target: AudioT = None,
                       noise: AudioT = None,
                       force: bool = False) -> Segsnr:
        """ Get the segsnr data for the given mixid

        mixid
            Mixture ID.
        segsnr_t
            segsnr_t data for the given mixid. (optional)
        targets
            List of augmented target audio (one per target in the mixup) for the given mixid. (optional)
        target
            Augmented target audio for the given mixid. (optional)
        noise
            Augmented noise audio for the given mixid. (optional)
        force
            Force computing data from original sources regardless of whether or not cached data exists.
        """
        if not force:
            segsnr = self.read_mixture_data(mixid, 'segsnr')
            if segsnr is not None:
                return segsnr

            segsnr_t = self.read_mixture_data(mixid, 'segsnr_t')
            if segsnr_t is not None:
                return segsnr_t[0::self.ft_config.R]

        if segsnr_t is None:
            segsnr_t = self.mixture_segsnr_t(mixid=mixid, targets=targets, target=target, noise=noise)

        segsnr = segsnr_t[0::self.ft_config.R]

        return segsnr

    def mixture_ft(self,
                   mixid: int,
                   targets: AudiosT = None,
                   target: AudioT = None,
                   noise: AudioT = None,
                   mixture_f: AudioF = None,
                   mixture: AudioT = None,
                   truth_t: Truth = None,
                   force: bool = False) -> (Feature, Truth):
        """ Get the feature, truth_f data for the given mixid

        mixid
            Mixture ID.
        targets
            List of augmented target audio (one per target in the mixup) for the given mixid. (optional)
        target
            Augmented target audio for the given mixid. (optional)
        noise
            Augmented noise audio for the given mixid. (optional)
        mixture
            Mixture audio for the given mixid. (optional)
        truth_t
            truth_t for the given mixid. (optional)
        force
            Force computing data from original sources regardless of whether or not cached data exists.
        """
        import numpy as np

        from sonusai.mixture import truth_reduction

        if not force:
            feature, truth_f = self.read_mixture_data(mixid, ['feature', 'truth_f'])
            if feature is not None and truth_f is not None:
                return feature, truth_f

        if mixture_f is None:
            mixture_f = self.mixture_mixture_f(mixid=mixid,
                                               targets=targets,
                                               target=target,
                                               noise=noise,
                                               mixture=mixture)

            if truth_t is None:
                truth_t = self.mixture_truth_t(mixid=mixid, targets=targets, noise=noise)

            mrecord = self.mixtures[mixid]

            transform_frames = self.mixture_transform_frames(mixid)
            feature_frames = self.mixture_feature_frames(mixid)

            if truth_t is None:
                truth_t = np.zeros((self.mixture_samples(mixid), self.num_classes), dtype=np.float32)

            feature = np.empty((feature_frames, self.fg.stride, self.fg.num_bands), dtype=np.float32)
            truth_f = np.empty((feature_frames, self.num_classes), dtype=np.complex64)

            feature_frame = 0
            for transform_frame in range(transform_frames):
                indices = slice(transform_frame * self.ft_config.R,
                                (transform_frame + 1) * self.ft_config.R)
                self.fg.execute(mixture_f[transform_frame],
                                truth_reduction(truth_t[indices], self.truth_reduction_function))

                if self.fg.eof():
                    feature[feature_frame] = self.fg.feature()
                    truth_f[feature_frame] = self.fg.truth()
                    feature_frame += 1

            if np.isreal(truth_f).all():
                truth_f = np.float32(np.real(truth_f))

            return feature, truth_f

        def mixture_feature(self,
                            mixid: int,
                            targets: AudiosT = None,
                            noise: AudioT = None,
                            mixture: AudioT = None,
                            truth_t: Truth = None,
                            force: bool = False) -> Feature:
            """ Get the feature data for the given mixid

            mixid
                Mixture ID.
            targets
                List of augmented target audio (one per target in the mixup) for the given mixid. (optional)
            noise
                Augmented noise audio for the given mixid. (optional)
            mixture
                Mixture audio for the given mixid. (optional)
            truth_t
                truth_t for the given mixid. (optional)
            force
                Force computing data from original sources regardless of whether or not cached data exists.
            """
            feature, _ = self.mixture_ft(mixid=mixid,
                                         targets=targets,
                                         noise=noise,
                                         mixture=mixture,
                                         truth_t=truth_t,
                                         force=force)
            return feature

        def mixture_truth_f(self,
                            mixid: int,
                            targets: AudiosT = None,
                            noise: AudioT = None,
                            mixture: AudioT = None,
                            truth_t: Truth = None,
                            force: bool = False) -> Feature:
            """ Get the truth_f data for the given mixid

            mixid
                Mixture ID.
            targets
                List of augmented target audio (one per target in the mixup) for the given mixid. (optional)
            noise
                Augmented noise audio for the given mixid. (optional)
            mixture
                Mixture audio for the given mixid. (optional)
            truth_t
                truth_t for the given mixid. (optional)
            force
                Force computing data from original sources regardless of whether or not cached data exists.
            """
            _, truth_f = self.mixture_ft(mixid=mixid,
                                         targets=targets,
                                         noise=noise,
                                         mixture=mixture,
                                         truth_t=truth_t,
                                         force=force)
            return truth_f

    def mixture_class_count(self,
                            mixid: int,
                            targets: AudiosT = None,
                            noise: AudioT = None,
                            truth_t: Truth = None) -> ClassCount:
        """ Compute the number of samples for which each truth index is active for the given mixid

        targets
            List of augmented target audio (one per target in the mixup) for the given mixid. (optional)
        noise
            Augmented noise audio for the given mixid. (optional)
        mixture
            Mixture audio for the given mixid. (optional)
        truth_t
            truth_t for the given mixid. (optional)
        """
        if truth_t is None:
            truth_t = self.mixture_truth_t(mixid=mixid, targets=targets, noise=noise)

        class_count = [0] * self.num_classes
        num_classes = self.num_classes
        if self.truth_mutex:
            num_classes -= 1
        for cl in range(num_classes):
            class_count[cl] = int(np.sum(truth_t[:, cl] >= self.class_weights_threshold[cl]))

        return class_count


def get_mrecords_from_mixids(mixdb: MixtureDatabase, mixids: MixtureIDs = None) -> MRecords:
    from copy import deepcopy

    return [deepcopy(mixdb.mixtures[i]) for i in mixdb.mixids_to_list(mixids)]


def _augment_audio(audio: AudioT, augmentations: Augmentations) -> AudiosT:
    from sonusai.mixture import apply_augmentation

    return [apply_augmentation(audio, augmentation) for augmentation in augmentations]
