#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import print_function, absolute_import
import sys
import os
import soundfile
import tempfile
import numpy
import librosa
from tqdm import tqdm
from six.moves.http_client import BadStatusLine

from dcase_util.containers import ContainerMixin, FileMixin
from dcase_util.ui.ui import FancyStringifier
from dcase_util.utils import FileFormat, is_int

class AudioContainer(ContainerMixin, FileMixin):
    """Audio container class."""
    valid_formats = [FileFormat.WAV, FileFormat.FLAC,
                     FileFormat.OGG,
                     FileFormat.M4A, FileFormat.WEBM,
                     FileFormat.MP3, FileFormat.MP4]  #: Valid file formats

    def __init__(self,
                 data=None, fs=44100,
                 focus_start_samples=None, focus_stop_samples=None, focus_channel=None,
                 **kwargs):
        """Constructor

        Parameters
        ----------
        filename : str, optional
            File path
        fs : int
            Target sampling frequency, if loaded audio does have different sampling frequency, audio will be re-sampled.
            Default value "44100"

        """
        # Run ContainerMixin init
        ContainerMixin.__init__(self, **kwargs)

        # Run FileMixin init
        FileMixin.__init__(self, **kwargs)

        # Run super init to call init of mixins too
        super(AudioContainer, self).__init__(**kwargs)

        self.channel_axis = 0
        self.time_axis = 1

        # Audio data
        if data is None:
            # Initialize with array
            data = numpy.ndarray((0, ))

        self._data = data
        self.data_synced_with_file = False

        self.fs = fs
        self.filetype_info = None

        # Filename
        if self.filename:
            self.detect_file_format()
            self.validate_format()

        # Initialize focus segment variables
        self._focus_start = None
        self._focus_stop = None
        self._focus_channel = None

        # Set focus segment through properties
        self.focus_start_samples = focus_start_samples
        self.focus_stop_samples = focus_stop_samples
        self.focus_channel = focus_channel

    def __getstate__(self):
        d = super(AudioContainer, self).__getstate__()
        d.update({
            'channel_axis': self.channel_axis,
            'time_axis': self.time_axis,
            '_data': self._data,
            'data_synced_with_file': self.data_synced_with_file,
            'fs': self.fs,
            'filetype_info': self.filetype_info,
            'filename': self.filename,
            '_focus_start': self._focus_start,
            '_focus_stop': self._focus_stop,
            '_focus_channel': self._focus_channel,
        })

        return d

    def __setstate__(self, d):
        super(AudioContainer, self).__setstate__(d)

        self.channel_axis = d['channel_axis']
        self.time_axis = d['time_axis']

        self._data = d['_data']
        self.data_synced_with_file = d['data_synced_with_file']
        self.fs = d['fs']
        self.filetype_info = d['filetype_info']
        self.filename = d['filename']
        self._focus_start = None
        self._focus_stop = None
        self._focus_channel = None
        self.focus_start = d['_focus_start']
        self.focus_stop = d['_focus_stop']
        self.focus_channel = d['_focus_channel']

    def __str__(self):
        output = ''
        ui = FancyStringifier()
        output += ui.class_name(self.__class__.__name__) + '\n'
        if self.filename:
            output += ui.data(
                field='Filename',
                value=self.filename
            ) + '\n'

            if self.filetype_info and self.filetype_info.values:
                output += ui.data(
                    field='Format',
                    value=self.format + ' (' + ', '.join(self.filetype_info.values()) + ')'
                ) + '\n'

            else:
                output += ui.data(field='Format', value=self.format) + '\n'

            output += ui.data(
                field='Synced',
                value='Yes' if self.data_synced_with_file else 'No'
            ) + '\n'

        output += ui.data(
            field='Sampling rate',
            value=str(self.fs), unit='hz'
        ) + '\n'

        output += ui.data(
            field='Channels',
            value=str(self.channels)
        ) + '\n'

        output += ui.line(field='Duration') + '\n'
        output += ui.data(
            indent=4,
            field='Seconds',
            value=self.duration_sec,
            unit='sec'
        ) + '\n'

        output += ui.data(
            indent=4,
            field='Milliseconds',
            value=self.duration_ms,
            unit='ms'
        ) + '\n'

        output += ui.data(
            indent=4,
            field='Samples',
            value=self.duration_samples,
            unit='samples'
        ) + '\n'

        if self._focus_channel is not None or self._focus_start is not None or self._focus_stop is not None:
            output += ui.line(field='Focus segment') + '\n'
            if self.focus_channel is not None:
                if self.channels == 2:
                    if self._focus_channel == 0:
                        output += ui.data(
                            indent=4,
                            field='Channel',
                            value='{channel} [{label}]'.format(
                                channel=self._focus_channel,
                                label='Left Channel'
                            )
                        ) + '\n'

                    elif self._focus_channel == 1:
                        output += ui.data(
                            indent=4,
                            field='Channel',
                            value='{channel} [{label}]'.format(
                                channel=self._focus_channel,
                                label='Right Channel'
                            )
                        ) + '\n'

                else:
                    output += ui.data(
                        indent=4,
                        field='Channel',
                        value=self._focus_channel
                    ) + '\n'

            output += ui.line(
                indent=4,
                field='Duration'
            ) + '\n'

            output += ui.data(
                indent=6,
                field='Seconds',
                value=self.focus_stop_seconds - self.focus_start_seconds,
                unit='sec'
            ) + '\n'

            output += ui.data(
                indent=6,
                field='Samples',
                value=self.focus_stop_samples - self.focus_start_samples,
                unit='sec'
            ) + '\n'

            output += ui.line(
                indent=4,
                field='Start point'
            ) + '\n'

            output += ui.data(
                indent=6,
                field='Seconds',
                value=self.focus_start_seconds,
                unit='sec') + '\n'

            output += ui.data(
                indent=6,
                field='Samples',
                value=self.focus_start_samples,
                unit='samples'
            ) + '\n'

            output += ui.line(
                indent=4,
                field='Stop point'
            ) + '\n'

            output += ui.data(
                indent=6,
                field='Seconds',
                value=self.focus_stop_seconds,
                unit='sec'
            ) + '\n'

            output += ui.data(
                indent=6,
                field='Samples',
                value=self.focus_stop_samples,
                unit='samples'
            ) + '\n'

        return output

    def __nonzero__(self):
        return self.loaded

    def __getitem__(self, i):
        """Get ith sample, in case of multiple channels array is across channels is returned"""

        if not isinstance(i, int):
            raise TypeError("Index should be integer")

        if i < 0 or i > self.length:
            raise KeyError(i)

        if len(self._data.shape) == 1:
            return self._data[i]

        elif len(self._data.shape) > 1:
            return self._data[:, i]

        else:
            return None

    def __setitem__(self, i, value):
        """Set ith sample"""

        if not isinstance(i, int):
            raise TypeError("Index should be integer")

        if i < 0 or i > self.length:
            raise KeyError(i)

        if len(self._data.shape) == 1:
            self._data[i] = value

        elif len(self._data.shape) > 1:
            self._data[:, i] = value

    def __iter__(self):
        return iter(self._data)

    def __len__(self):
        return self.length

    @property
    def data(self):
        """Audio data

        Returns
        -------
        numpy.ndarray
            Audio data

        """

        return self._data

    @data.setter
    def data(self, value):
        self._data = value
        self.data_synced_with_file = False

    @property
    def focus_start_samples(self):
        """Focus segment start in samples.

        Returns
        -------
        int
            Focus segment start in samples

        """

        return self._focus_start

    @focus_start_samples.setter
    def focus_start_samples(self, value):
        if value is not None and value > 0:
            value = int(value)
            self._focus_start = value

            if self._focus_stop is not None and self._focus_stop < self._focus_start:
                # focus points are reversed
                start = self._focus_start
                self._focus_start = self._focus_stop
                self._focus_stop = start

        else:
            self._focus_start = 0

    @property
    def focus_start_seconds(self):
        """Focus segment start in seconds.

        Returns
        -------
        int
            Focus segment start in seconds

        """

        return self._sample_to_time(sample=self.focus_start_samples)

    @focus_start_seconds.setter
    def focus_start_seconds(self, value):
        self.focus_start_samples = self._time_to_sample(time=value)

    @property
    def focus_stop_samples(self):
        """Focus segment stop in samples.

        Returns
        -------
        int
            Focus segment stop in samples

        """
        if self._focus_stop is None:
            return self.length

        else:
            return self._focus_stop

    @focus_stop_samples.setter
    def focus_stop_samples(self, value):
        if value is None:
            self._focus_stop = None

        else:
            if value <= self.duration_samples and value is not None:
                value = int(value)
                self._focus_stop = value

                if self._focus_start is not None and self._focus_stop < self._focus_start:
                    # focus points are reversed
                    start = self._focus_start
                    self._focus_start = self._focus_stop
                    self._focus_stop = start

            else:
                self._focus_stop = self.duration_samples

    @property
    def focus_stop_seconds(self):
        """Focus segment stop in seconds.

        Returns
        -------
        int
            Focus segment stop in seconds

        """

        return self._sample_to_time(sample=self.focus_stop_samples)

    @focus_stop_seconds.setter
    def focus_stop_seconds(self, value):
        self.focus_stop_samples = self._time_to_sample(time=value)

    @property
    def focus_channel(self):
        """Focus channel

        Returns
        -------
        int or str
            Focus channel

        """

        return self._focus_channel

    @focus_channel.setter
    def focus_channel(self, value):
        if value is not None and is_int(value):
            if 0 <= value < self.channels:
                self._focus_channel = value

            else:
                self._focus_channel = None

        elif value is not None and isinstance(value, str):
            if value.lower() == 'mixdown':
                self._focus_channel = 'mixdown'

            elif value.lower() == 'left' or value.lower() == 'l':
                self._focus_channel = 0

            elif value.lower() == 'right' or value.lower() == 'r':
                self._focus_channel = 1

            else:
                # Unknown channel label given
                message = '{name}: Unknown channel [{channel}]'.format(name=self.__class__.__name__, channel=value)
                self.logger.exception(message)
                raise ValueError(message)

        else:
            self._focus_channel = None

    @property
    def loaded(self):
        """Audio load status.

        Returns
        -------
        bool
            Audio status

        """

        if isinstance(self._data, numpy.ndarray) and len(self._data) > 0:
            return True
        else:
            return False

    @property
    def shape(self):
        """Audio data shape.

        Returns
        -------
        tuple
            shape of audio data

        """

        if self.loaded:
            return self._data.shape
        else:
            return None

    @property
    def length(self):
        """Length of audio data in samples.

        Returns
        -------
        int
            Audio length

        """

        if self.loaded:
            if len(self._data.shape) == 1:
                return self._data.shape[0]

            elif len(self._data.shape) > 1:
                return self._data.shape[-1]

            else:
                return 0

        else:
            return 0

    @property
    def duration_samples(self):
        """Duration of audio data in samples.

        Returns
        -------
        int
            Audio duration

        """

        return self.length

    @property
    def duration_ms(self):
        """Duration of audio data in milliseconds.

        Returns
        -------
        float
            Audio duration

        """

        return (self.length / float(self.fs)) * 1000

    @property
    def duration_sec(self):
        """Duration of audio data in seconds.

        Returns
        -------
        float
            Audio duration

        """

        return self.length / float(self.fs)

    @property
    def channels(self):
        """Number of audio channels.

        Returns
        -------
        int
            Number of audio channels

        """

        if self.loaded:
            if len(self.data.shape) == 2:
                return self._data.shape[self.channel_axis]
            elif len(self.data.shape) == 1:
                return 1
            else:
                return 0
        else:
            return 0

    @property
    def streams(self):
        """Rename channels for compatibility.

        Returns
        -------
        int
            Number of streams

        """
        return self.channels

    @property
    def empty(self):
        """Check if audio data is empty.

        In case audio is not yet loaded it is first loaded into container from disk.

        Returns
        -------
        bool

        """

        if self.loaded:
            if self.length == 0:
                return True

            else:
                return False

        else:
            if self.filename and self.exists():
                # Audio data is not yet loaded and filename set and file exists, load the data from a file
                self.load()

                if self.length == 0:
                    return True

                else:
                    return False

            else:
                return True

    def load(self, filename=None, fs='native', mono=False, res_type='kaiser_best', start=None, stop=None):
        """Load file

        Parameters
        ----------
        filename : str, optional
            File path, if None given filename parameter given to class constructor is used.

        fs : int or str
            Target sampling frequency, if loaded audio does have different sampling frequency, audio will
            be re-sampled. If None given, value given to class constructor is used. If 'native' is given then
            native sampling frequency defined by audio file is used.
            Default value 'native'

        mono : bool
            Monophonic target, multi-channel audio will be down-mixed.
            Default value False

        res_type : str
            Resample type, defined by Librosa.
            Default value 'kaiser_best'

        start : float, optional
            Segment start time in seconds.
            Default value None

        stop : float, optional
            Segment stop time in seconds.
            Default value None

        Raises
        ------
        IOError:
            File does not exists or has unknown file format

        Returns
        -------
        self

        """

        if filename is not None:
            self.filename = filename
            self.detect_file_format()
            self.validate_format()

        if self.exists():
            if fs is None:
                # Use sampling frequency defined in class construction.
                fs = self.fs

            if self.format == FileFormat.WAV:
                info = soundfile.info(file=self.filename)

                self.filetype_info = {
                    'subtype': info.subtype,
                    'subtype_info': info.subtype_info
                }

                # Handle segment start and stop
                if start is not None and stop is not None:
                    start_sample = int(start * info.samplerate)
                    stop_sample = int(stop * info.samplerate)

                    if stop_sample > info.frames:
                        stop_sample = info.frames

                else:
                    start_sample = None
                    stop_sample = None

                self._data, source_fs = soundfile.read(
                    file=self.filename,
                    start=start_sample,
                    stop=stop_sample
                )

                self._data = self._data.T

                # Down-mix audio
                if mono and len(self._data.shape) > 1:
                    self._data = numpy.mean(self._data, axis=self.channel_axis)

                if fs == 'native':
                    # Use native sampling frequency.
                    self.fs = source_fs

                else:
                    # Target sampling frequency defined, possibly re-sample signal.
                    if fs != source_fs:
                        self._data = librosa.core.resample(
                            self._data,
                            source_fs,
                            fs,
                            res_type=res_type
                        )

                    # Store sampling frequency
                    self.fs = fs

            elif self.format in [FileFormat.FLAC, FileFormat.OGG,
                                 FileFormat.MP3,
                                 FileFormat.M4A, FileFormat.MP4, FileFormat.WEBM]:
                # Handle segment start and stop
                if start is not None and stop is not None:
                    offset = start
                    duration = stop - start

                else:
                    offset = 0.0
                    duration = None

                if fs == 'native':
                    # Use native sampling frequency
                    sr = None

                else:
                    # Use target sampling frequency
                    sr = fs

                self._data, self.fs = librosa.load(
                    self.filename,
                    sr=sr,
                    mono=mono,
                    res_type=res_type,
                    offset=offset,
                    duration=duration
                )

            else:
                message = '{name}: Unknown format [{format}]'.format(
                    name=self.__class__.__name__,
                    format=self.filename
                )

                self.logger.exception(message)
                raise IOError(message)

        else:
            message = '{name}: File does not exists [{file}]'.format(
                name=self.__class__.__name__,
                file=self.filename
            )

            self.logger.exception(message)
            raise IOError(message)

        # Check if after load function is defined, call if found
        if hasattr(self, '_after_load'):
            self._after_load()

        # Internal data is synced with the file, until it is edited.
        self.data_synced_with_file = True

        return self

    def save(self, filename=None, bit_depth=16, bit_rate=None):
        """Save audio

        Parameters
        ----------
        filename : str, optional
            File path, if None given filename parameter given to class constructor is used.
            Default value None

        bit_depth : int, optional
            Bit depth for audio.
            Default value 16

        bit_rate : int, optional
            Bit rate for compressed audio formats.
            Default value None

        Raises
        ------
        ImportError:
            Error if file format specific module cannot be imported
        IOError:
            File has unknown file format

        Returns
        -------
        self

        """

        if filename:
            self.filename = filename
            self.detect_file_format()
            self.validate_format()

        if self.filename is None or self.filename == '':
            message = '{name}: Filename is empty [{filename}]'.format(
                name=self.__class__.__name__,
                filename=self.filename
            )

            self.logger.exception(message)
            raise IOError(message)

        # Check if before save function is defined, call if found
        if hasattr(self, '_before_save'):
            self._before_save()

        if self.format == FileFormat.WAV:
            if bit_depth == 16:
                subtype = 'PCM_16'

            elif bit_depth == 24:
                subtype = 'PCM_24'

            elif bit_depth == 32:
                subtype = 'PCM_32'

            else:
                message = '{name}: Unexpected bit depth [{bitdepth}]'.format(
                    name=self.__class__.__name__,
                    bitdepth=bit_depth
                )

                self.logger.exception(message)
                raise IOError(message)

            soundfile.write(
                file=self.filename,
                data=self._data.T,
                samplerate=self.fs,
                subtype=subtype
            )

        elif self.format == FileFormat.FLAC:
            if bit_depth == 16:
                subtype = 'PCM_16'

            elif bit_depth == 24:
                subtype = 'PCM_24'

            elif bit_depth == 32:
                subtype = 'PCM_32'

            else:
                message = '{name}: Unexpected bit depth [{bitdepth}]'.format(
                    name=self.__class__.__name__,
                    bitdepth=bit_depth
                )

                self.logger.exception(message)
                raise IOError(message)

            soundfile.write(
                file=self.filename,
                data=self._data.T,
                samplerate=self.fs,
                format='flac',
                subtype=subtype
            )

        elif self.format == FileFormat.OGG:
            soundfile.write(
                file=self.filename,
                data=self._data.T,
                samplerate=self.fs,
                format='OGG',
                subtype='VORBIS'
            )

        elif self.format == FileFormat.MP3:
            # Notice: Saving with MP3 format results in slightly longer signal than original.
            # Difference is due to padding in the compression algorithm, and is usually around 200 - 1000 samples.
            import subprocess
            import platform

            if platform.system() == 'Windows':
                ffmpeg_binary = "ffmpeg.exe"

            else:
                ffmpeg_binary = "ffmpeg"

            if bit_rate not in [8, 16, 24, 32, 40, 48, 64, 80, 96, 112, 128, 160, 192, 224, 256, 320]:
                message = '{name}: Unsupported bit rate [{bitrate}]'.format(
                    name=self.__class__.__name__,
                    bitrate=bit_rate
                )

                self.logger.exception(message)
                raise IOError(message)

            command = [
                ffmpeg_binary,
                '-y',                                               # enable overwrite file
                '-f', 's16le',                                      # input format
                '-acodec', 'pcm_s16le',                             # input bit depth
                '-r', str(self.fs),                                 # sampling rate
                '-ac', str(self.channels),                          # amount of channels
                '-i', '-',                                          # input from pipe
                '-vn',                                              # no video input
                '-acodec', 'libmp3lame',                            # output audio codec
                '-b:a', "{bitrate:d}k".format(bitrate=bit_rate),    # bit rate
                self.filename                                       # output filename
            ]

            popen_parameters = {
                'stdin': subprocess.PIPE,
                'stdout': subprocess.PIPE,
                'stderr': subprocess.PIPE
            }

            pipe = subprocess.Popen(
                command,
                **popen_parameters
            )

            # Convert signal data from float [-1,1] to signed 16-bit
            audio_signal = numpy.asarray(self.data).T
            signal_max_value = 2 ** (16 - 1)
            audio_signal = (audio_signal * signal_max_value).clip(
                -signal_max_value,
                signal_max_value - 1
            ).astype('int16')

            try:
                try:
                    pipe.stdin.write(
                        audio_signal.tobytes()
                    )

                except NameError:
                    pipe.stdin.write(
                        audio_signal.tostring()
                    )

            except IOError as error:
                pipe_error = pipe.stderr.read()
                error = str(error)
                error += "\n\nFFMPEG encountered the following error {filename}:".format(filename=self.filename)
                error += "\n\n" + str(pipe_error)

                raise IOError(error)

            pipe.stdin.close()
            if pipe.stderr is not None:
                pipe.stderr.close()

            pipe.wait()

        else:
            message = '{name}: Unknown format for saving [{format}]'.format(
                name=self.__class__.__name__,
                format=self.filename
            )

            self.logger.exception(message)
            raise IOError(message)

        # Check if after save function is defined, call if found
        if hasattr(self, '_after_save'):
            self._after_save()

        # Internal data is synced with the file, until it is edited.
        self.data_synced_with_file = True

        return self

    def load_from_youtube(self, query_id, start=None, stop=None, mono=False, silent=True):
        """Load audio data from youtube

        Parameters
        ----------
        query_id : str
            Youtube query id.

        start : float, optional
            Segment start time in seconds.
            Default value None

        stop : float, optional
            Segment stop time in seconds.
            Default value None

        mono : bool
            Monophonic target, multi-channel audio will be down-mixed.
            Default value False

        silent : bool
            Switch to show progress bar.
            Default value True

        Raises
        ------
        IOError:
            Youtube video does not exists or cannot be downloaded

        Returns
        -------
        self

        """

        def progress_hook(t):
            """Wraps tqdm instance. Don't forget to close() or __exit__()
            the tqdm instance once you're done with it (easiest using `with` syntax).
            """

            def inner(total, recvd, ratio, rate, eta):
                t.total = int(total / 1024.0)
                t.update(int(recvd / 1024.0))

            return inner

        try:
            import pafy

        except ImportError:
            message = '{name}: Unable to import pafy module. You can install it with `pip install pafy`.'.format(
                name=self.__class__.__name__
            )

            self.logger().exception(message)
            raise ImportError(message)

        try:
            from youtube_dl.utils import ExtractorError

        except ImportError:
            message = '{name}: Unable to import youtube_dl module. You can install it with `pip install youtube-dl`.'.format(
                name=self.__class__.__name__
            )

            self.logger().exception(message)
            raise ImportError(message)

        try:
            # Access youtube video and get best quality audio stream
            youtube_audio = pafy.new(
                url='https://www.youtube.com/watch?v={query_id}'.format(query_id=query_id),
                basic=False,
                gdata=False,
                size=False
            ).getbestaudio()

            # Get temp file
            tmp_file = tempfile.NamedTemporaryFile(suffix='.'+youtube_audio.extension)

            # Get temporary filename
            tmp_filename = tmp_file.name

            # Remove temporary file (avoid FileExistsError on Windows)
            tmp_file.close()

            download_progress_bar = None
            if not silent:
                # Create download progress bar
                download_progress_bar = tqdm(
                    desc="{0: <25s}".format('Download youtube item '),
                    file=sys.stdout,
                    unit='B',
                    unit_scale=True,
                    leave=False,
                    disable=self.disable_progress_bar,
                    ascii=self.use_ascii_progress_bar
                )
                callback = progress_hook(download_progress_bar)
            else:
                callback = None

            # Download audio
            youtube_audio.download(
                filepath=tmp_filename,
                quiet=True,
                callback=callback
            )

            if not silent:
                # Close progress bar
                download_progress_bar.close()

                # Create audio processing progress bar
                audio_processing_progress_bar = tqdm(
                    desc="{0: <25s}".format('Processing '),
                    initial=0,
                    total=4,
                    file=sys.stdout,
                    leave=False,
                    disable=self.disable_progress_bar,
                    ascii=self.use_ascii_progress_bar
                )

            # Store current filename
            filename = self.filename

            # Load audio segment
            self.load(
                filename=tmp_filename,
                mono=mono,
                fs=self.fs,
                res_type='kaiser_best',
                start=float(start) if start is not None else None,
                stop=float(stop) if stop is not None else None
            )

            # Restore filename
            if filename:
                self.filename = filename
                self.detect_file_format()

            if not silent:
                audio_processing_progress_bar.update(1)
                audio_processing_progress_bar.update(3)
                audio_processing_progress_bar.close()

        except (IOError, BadStatusLine, ExtractorError) as e:
            # Store files with errors
            raise IOError(e.message)

        except (KeyboardInterrupt, SystemExit):
            # Remove temporal file and current audio file.
            os.remove(self.filename)
            raise

        return self

    def normalize(self, headroom=0.005):
        """Normalize audio data.

        Data is normalized between -(1.0 - headroom) and +(1.0 - headroom)

        Parameters
        ----------
        headroom : float
            How much headroom there should be left under 1.0.
            Default value 0.005

        Returns
        -------
        self

        """

        if self.channels > 1:
            for channel_data in self._data:
                mean_value = numpy.mean(channel_data)
                channel_data -= mean_value

                max_value = max(abs(channel_data)) + headroom
                channel_data /= max_value
        else:
            mean_value = numpy.mean(self._data)
            self._data -= mean_value

            max_value = max(abs(self._data)) + headroom
            self._data /= max_value

        return self

    def resample(self, target_fs, scale=True, res_type='kaiser_best'):
        """Resample audio data.

        Parameters
        ----------
        target_fs : int
            Target sampling rate

        scale : bool
            Scale the resampled signal to have approximately equal total energy (see `librosa.core.resample`).
            Default value True

        res_type : str
            Resample type (see `librosa.core.resample`)
            Default value 'kaiser_best'

        Returns
        -------
        self

        """

        if target_fs != self.fs:
            self._data = librosa.resample(
                y=self._data,
                orig_sr=self.fs,
                target_sr=target_fs,
                scale=scale,
                res_type=res_type
            )
            self.fs = target_fs

        return self

    def mixdown(self):
        """Mix all audio channels into single channel.

        Returns
        -------
        self

        """

        self.reset_focus()
        self.set_focus(channel='mixdown')
        self.freeze()

        return self

    def reset_focus(self):
        """Reset focus segment.

        Returns
        -------
        self

        """

        self._focus_start = None
        self._focus_stop = None
        self._focus_channel = None

        return self

    def set_focus(self,
                  start=None, stop=None, duration=None,
                  start_seconds=None, stop_seconds=None, duration_seconds=None,
                  channel=None):
        """Set focus segment

        Parameters
        ----------
        start : int
            Sample index of focus segment start.
            Default value None

        stop : int
            Sample index of focus segment stop.
            Default value None

        duration : int
            Sample count of focus segment.
            Default value None

        start_seconds : float
            Time stamp (in seconds) of focus segment start.
            Default value None

        stop_seconds : float
            Time stamp (in seconds) of focus segment stop.
            Default value None

        duration_seconds : float
            Duration (in seconds) of focus segment.
            Default value None

        channel : int or str
            Audio channel id or name to focus. In case of stereo signal, valid channel labels to select
            single channel are 'L', 'R', 'left', and 'right' or 0, 1, and to get mixed down
            version of all channels 'mixdown'.
            Default value None

        Returns
        -------
        self

        """

        if start is not None or stop is not None or duration is not None:
            # Sample based setting
            if start is not None and stop is not None:
                self.reset_focus()
                self.focus_start_samples = start
                self.focus_stop_samples = stop

            elif start is not None and duration is not None:
                self.reset_focus()
                self.focus_start_samples = start
                self.focus_stop_samples = start + duration

        elif start_seconds is not None or stop_seconds is not None or duration_seconds is not None:
            # Time based setting
            if start_seconds is not None and stop_seconds is not None:
                self.reset_focus()
                self.focus_start_samples = self._time_to_sample(time=start_seconds)
                self.focus_stop_samples = self._time_to_sample(time=stop_seconds)

            elif start_seconds is not None and duration_seconds is not None:
                self.reset_focus()
                self.focus_start_samples = self._time_to_sample(time=start_seconds)
                self.focus_stop_samples = self._time_to_sample(time=start_seconds + duration_seconds)

        else:
            # Reset
            self._focus_start = None
            self._focus_stop = None

        self.focus_channel = channel

        return self

    def get_focused(self):
        """Get focus segment from audio data.

        Returns
        -------
        numpy.ndarray

        """

        focused_data = None

        if self.focus_start_samples is not None or self.focus_stop_samples is not None:
            if self.focus_start_samples is not None:
                focus_start_samples = self.focus_start_samples

            else:
                focus_start_samples = 0

            if self.focus_stop_samples is not None:
                focus_stop_samples = self.focus_stop_samples

            else:
                focus_stop_samples = self.length

            if self.channels == 1:
                # We have single channel
                focused_data = self._data[focus_start_samples:focus_stop_samples]

            elif self.channels > 1:
                # We have multichannel audio
                focused_data = []
                for channel_data in self._data:
                    focused_data.append(channel_data[focus_start_samples:focus_stop_samples])

                focused_data = numpy.vstack(focused_data)

        else:
            focused_data = self._data

        if self.focus_channel is not None and is_int(self.focus_channel) and 0 <= self.focus_channel < self.channels:
            return focused_data[self.focus_channel, :]

        elif self.focus_channel == 'mixdown' and self.channels > 1:
            return numpy.mean(focused_data, axis=self.channel_axis)

        else:
            return focused_data

    def freeze(self):
        """Freeze focus segment, copy segment to be container's data.

        Returns
        -------
        self

        """

        self._data = self.get_focused()
        self.reset_focus()

        return self

    def frames(self,
               frame_length=None, hop_length=None,
               frame_length_seconds=None, hop_length_seconds=None):
        """Slice audio into overlapping frames.

        Parameters
        ----------
        frame_length : int, optional
            Frame length in samples. Set either frame_length or frame_length_seconds.
            Default value None

        hop_length : int, optional
            Frame hop length in samples. Set either hop_length or hop_length_seconds.
            Default value None

        frame_length_seconds : float, optional
            Frame length in seconds, converted into samples based on sampling rate.
            Default value None

        hop_length_seconds: float, optional
            Frame hop length in seconds, converted into samples based on sampling rate.
            Default value None

        Raises
        ------
        ValueError:
            No frame_length and no frame_length_seconds given.
            No hop_length and no hop_length_seconds given.

        Returns
        -------
        numpy.ndarray

        """

        if not frame_length and frame_length_seconds:
            frame_length = int(self.fs * frame_length_seconds)

        if not hop_length and hop_length_seconds:
            hop_length = int(self.fs * hop_length_seconds)

        if not frame_length:
            message = '{name}: Specify frame_length parameter for frame splitting.'.format(
                name=self.__class__.__name__
            )
            self.logger.exception(message)
            raise ValueError(message)

        if not hop_length:
            message = '{name}: Specify hop_length parameter for frame splitting.'.format(
                name=self.__class__.__name__
            )
            self.logger.exception(message)
            raise ValueError(message)

        if self.channels == 1:
            return librosa.util.frame(
                y=self.get_focused(),
                frame_length=frame_length,
                hop_length=hop_length
            )

        else:
            data = []
            for channel_id, channel_data in enumerate(self.get_focused()):
                data.append(
                    librosa.util.frame(
                        y=channel_data,
                        frame_length=frame_length,
                        hop_length=hop_length
                    )
                )
            return numpy.array(data)

    def segments(self,
                 segment_length=None, segment_length_seconds=None,
                 segments=None,
                 skip_segments=None):
        """Slice audio into segments.

        Parameters
        ----------
        segment_length : int, optional
            Segment length in samples. Set either segment_length or segment_length_seconds. Used to produce
            consecutive non-overlapping segments.
            Default value None

        segment_length_seconds : float, optional
            Segment length in seconds, converted into samples based on sampling rate. Used to produce consecutive
            non-overlapping segments.
            Set either segment_length or segment_length_seconds.
            Default value None

        segments : list of dict or MetaDataContainer, optional
            List of time segments (onset and offset). If none given, segment length is used to produce consecutive
            non-overlapping segments.
            Default value None

        skip_segments : list of dict or MetaDataContainer, optional
            List of time segments (onset and offset) to be skipped when creating segments.
            Only used when segment_length or segment_length_seconds are given and segments are generated
            within this method.
            Default value None

        Raises
        ------
        ValueError:
            No segments and no segment_length given.

        Returns
        -------
        list, MetaDataContainer

        """
        from dcase_util.containers import MetaDataContainer

        if not segment_length and segment_length_seconds:
            # Get segment_length from segment_length_seconds
            segment_length = int(self.fs * segment_length_seconds)

        if segments is None and segment_length is not None:
            if skip_segments is not None:
                # Make sure skip segments is MetaDataContainer
                skip_segments = MetaDataContainer(skip_segments)

            # No segments given, get segments based on segment_length
            segment_start = 0
            segments = MetaDataContainer()
            while True:
                # Segment stop
                segment_stop = segment_start + segment_length
                if skip_segments is not None:
                    # Go through skip segments and adjust segment start and stop to avoid segments
                    for item in skip_segments:
                        if item.active_within_segment(
                                start=segment_start/float(self.fs),
                                stop=segment_stop/float(self.fs)
                        ):
                            # Adjust segment start to avoid current skip segment
                            segment_start = int(self.fs * item.offset)
                            # Adjust segment stop accordingly
                            segment_stop = segment_start + segment_length

                if segment_stop < self.length:
                    # Valid segment found, store it
                    segments.append(
                        {
                            'onset': segment_start/float(self.fs),
                            'offset': segment_stop/float(self.fs),
                        }
                    )

                # Set next segment start
                segment_start = segment_stop

                # Stop loop if segment_start is out of signal
                if segment_start > self.length:
                    break

        elif segments is not None:
            # Make sure segments is MetadataContainer
            segments = MetaDataContainer(segments)

        else:
            message = '{name}: Specify segments parameter or segment_length for segment creation.'.format(
                name=self.__class__.__name__
            )
            self.logger.exception(message)
            raise ValueError(message)

        # Get audio segments
        data = []
        for segment in segments:
            segment_start_samples = int(self.fs * segment.onset)
            segment_stop_samples = int(self.fs * segment.offset)
            if self.channels > 1:
                data.append(
                    self._data[:, segment_start_samples:segment_stop_samples]
                )

            else:
                data.append(
                    self._data[segment_start_samples:segment_stop_samples]
                )

        return data, segments

    def pad(self, type='silence', length=None, length_seconds=None):
        """Generate signal

        Parameters
        ----------
        type : str
            Default value 'silence'

        length : int, optional
            Default value None

        length_seconds : float, optional
            Default value None

        Returns
        -------
        list, MetaDataContainer

        """

        if not length and length_seconds is not None:
            # Get length from length_seconds
            length = int(self.fs * length_seconds)

        if self.length < length:
            if type=='silence':
                if len(self.data.shape) == 1:
                    self._data = numpy.pad(
                        array=self._data,
                        pad_width=(0,length-self.length),
                        mode='constant'
                    )
                else:
                    self._data = numpy.pad(
                        array=self._data,
                        pad_width=((0,0),(0,length-self.length)),
                        mode='constant'
                    )

        return self

    def plot(self, plot_type='wave', **kwargs):
        """Visualize audio data

        Parameters
        ----------

        plot_type : str
            Visualization type, 'wave' for waveform plot, 'spec' for spectrogram.
            Default value 'wave'

        Returns
        -------
        self

        """

        if plot_type == 'wave':
            self.plot_wave(**kwargs)

        elif plot_type == 'spec':
            self.plot_spec(**kwargs)

    def plot_wave(self, x_axis='time', max_points=50000.0, offset=0.0, color='#333333', alpha=1.0,
                  show_filename=True, plot=True):
        """Visualize audio data as waveform.

        Parameters
        ----------

        x_axis : str
            X-axis type.
            Default value 'time'

        max_points : float
            Maximum number of time-points to plot (see `librosa.display.waveplot`).
            Default value 50000

        offset : float
            Horizontal offset (in time) to start the waveform plot (see `librosa.display.waveplot`).
            Default value 0.0

        color : str
            Waveform fill color in hex-code.
            Default value '#333333'

        alpha : float
            Alpha of the waveform fill color.
            Default value 1.0

        show_filename : bool
            Show filename as figure title.
            Default value True

        plot : bool
            If true, figure is shown automatically. Set to False if collecting multiple plots into same figure
            outside this method.
            Default value True

        Returns
        -------
        self

        """

        import matplotlib.pyplot as plt
        from librosa.display import waveplot
        if plot:
            plt.figure()

        if self.channels > 1:
            # Plotting for multi-channel audio
            for channel_id, channel_data in enumerate(self.get_focused()):
                plt.subplot(self.channels, 1, channel_id + 1)
                if channel_id + 1 != self.channels:
                    current_x_axis = None
                else:
                    current_x_axis = x_axis

                waveplot(
                    y=channel_data.ravel(),
                    sr=self.fs,
                    x_axis=current_x_axis,
                    max_points=max_points,
                    offset=offset,
                    color=color,
                    alpha=alpha
                )

                plt.ylabel('Channel {channel:d}'.format(channel=channel_id))
                if channel_id == 0 and show_filename:
                    if self.filename:
                        plt.title(self.filename)

                if channel_id+1 != self.channels:
                    plt.xlabel('')

        else:
            # Plotting for single channel audio
            waveplot(
                y=self.get_focused().ravel(),
                sr=self.fs,
                x_axis=x_axis,
                max_points=max_points,
                offset=offset,
                color=color,
                alpha=alpha
            )

            plt.ylabel('Channel {channel:d}'.format(channel=0))

            if self.filename and show_filename:
                plt.title(self.filename)

        if plot:
            plt.show()

        return self

    def plot_spec(self, spec_type='log', hop_length=512, cmap='magma',
                  show_filename=True, show_colorbar=True, plot=True):
        """Visualize audio data as spectrogram.

        Parameters
        ----------

        spec_type : str
            Spectrogram type, use 'linear', 'log', 'cqt', 'cqt_hz', and 'cqt_note'.
            Default value 'log'

        hop_length : float
            Hop length, also used to determine time scale in x-axis (see `librosa.display.specshow`).
            Default value 512

        cmap : float
            Color map (see `librosa.display.specshow`).
            Default value 'magma'

        show_filename : bool
            Show filename as figure title.
            Default value True

        show_colorbar : bool
            Show color bar next to plot.
            Default value True

        plot : bool
            If true, figure is shown automatically. Set to False if collecting multiple plots into same
            figure outside this method.
            Default value True

        Returns
        -------
        self

        """

        from librosa.display import specshow
        import matplotlib.pyplot as plt

        if plot:
            plt.figure()

        if self.channels > 1:
            for channel_id, channel_data in enumerate(self.get_focused()):
                plt.subplot(self.channels, 1, channel_id+1)

                if spec_type in ['linear', 'log']:
                    D = librosa.core.amplitude_to_db(numpy.abs(librosa.stft(channel_data.ravel())) ** 2, ref=numpy.max)

                elif spec_type.startswith('cqt'):
                    D = librosa.core.amplitude_to_db(librosa.cqt(channel_data.ravel(), sr=self.fs), ref=numpy.max)

                if spec_type == 'linear':
                    specshow(
                        data=D,
                        sr=self.fs,
                        y_axis='linear',
                        x_axis='time',
                        hop_length=hop_length,
                        cmap=cmap
                    )

                elif spec_type == 'log':
                    specshow(
                        data=D,
                        sr=self.fs,
                        y_axis='log',
                        x_axis='time',
                        hop_length=hop_length,
                        cmap=cmap
                    )

                elif spec_type == 'cqt_hz' or 'cqt':
                    specshow(
                        data=D,
                        sr=self.fs,
                        y_axis='cqt_hz',
                        x_axis='time',
                        hop_length=hop_length,
                        cmap=cmap
                    )

                elif spec_type == 'cqt_note':
                    specshow(
                        data=D,
                        sr=self.fs,
                        y_axis='cqt_note',
                        x_axis='time',
                        hop_length=hop_length,
                        cmap=cmap
                    )

                if show_colorbar:
                    plt.colorbar(format='%+2.0f dB')

                plt.ylabel('Channel {channel:d}'.format(channel=channel_id))
                if channel_id == 0 and self.filename:
                    plt.title(self.filename)

        else:
            channel_id = 0

            if spec_type in ['linear', 'log']:
                D = librosa.core.amplitude_to_db(
                    numpy.abs(librosa.stft(self.get_focused().ravel())) ** 2,
                    ref=numpy.max
                )

            elif spec_type.startswith('cqt'):
                D = librosa.core.amplitude_to_db(
                    librosa.cqt(self.get_focused().ravel(), sr=self.fs),
                    ref=numpy.max
                )

            if spec_type == 'linear':
                specshow(
                    data=D,
                    sr=self.fs,
                    y_axis='linear',
                    x_axis='time',
                    hop_length=hop_length,
                    cmap=cmap
                )

            elif spec_type == 'log':
                specshow(
                    data=D,
                    sr=self.fs,
                    y_axis='log',
                    x_axis='time',
                    hop_length=hop_length,
                    cmap=cmap
                )

            elif spec_type == 'cqt_hz' or 'cqt':
                specshow(
                    data=D,
                    sr=self.fs,
                    y_axis='cqt_hz',
                    x_axis='time',
                    hop_length=hop_length,
                    cmap=cmap
                )

            elif spec_type == 'cqt_note':
                specshow(
                    data=D,
                    sr=self.fs,
                    y_axis='cqt_note',
                    x_axis='time',
                    hop_length=hop_length,
                    cmap=cmap
                )

            if show_colorbar:
                plt.colorbar(format='%+2.0f dB')

            plt.ylabel('Channel {channel:d}'.format(channel=channel_id))
            if show_filename and channel_id == 0:
                plt.title(self.filename)

        if plot:
            plt.show()

    def _time_to_sample(self, time):
        """Time to sample index.

        Parameters
        ----------
        time : float
            Time stamp in seconds.

        Returns
        -------
        int

        """

        return int(time * self.fs)

    def _sample_to_time(self, sample):
        """Sample index to time.

        Parameters
        ----------
        sample : int
            Sample index.

        Returns
        -------
        float

        """

        return sample / float(self.fs)


"""
Photon emission and absoprtion models.
"""
import numpy as np
import h5py

from soxs.spectra import ApecGenerator, \
    get_wabs_absorb, get_tbabs_absorb
from soxs.utils import parse_prng
from pyxsim.utils import mylog
from yt.funcs import get_pbar
from yt.units.yt_array import YTArray, YTQuantity
from yt.utilities.physical_constants import hcgs, clight

hc = (hcgs*clight).in_units("keV*angstrom").v
# NOTE: XSPEC has hc = 12.39854 keV*A, so there may be slight differences in
# placement of spectral lines due to the above
cl = clight.v

class TableApecModel(ApecGenerator):
    r"""
    Initialize a thermal gas emission model from the AtomDB APEC tables
    available at http://www.atomdb.org. This code borrows heavily from Python
    routines used to read the APEC tables developed by Adam Foster at the
    CfA (afoster@cfa.harvard.edu).

    Parameters
    ----------
    emin : float
        The minimum energy for the spectral model.
    emax : float
        The maximum energy for the spectral model.
    nchan : integer
        The number of channels in the spectral model. If one
        is thermally broadening lines, it is recommended that
        this value result in an energy resolution per channel
        of roughly 1 eV.
    var_elem : list of strings, optional
        The names of elements to allow to vary freely
        from the single abundance parameter. Default:
        None
    model_root : string, optional
        The directory root where the model files are stored. If
        not provided, the default SOXS-provided files are used.
    model_vers : string, optional
        The version identifier string for the APEC files, e.g.
        "3.0.3". Default: 3.0.8
    thermal_broad : boolean, optional
        Whether or not the spectral lines should be thermally
        broadened. Default: True
    nolines : boolean, optional
        Turn off lines entirely for generating spectra.
        Default: False
    abund_table : string or array_like, optional
        The abundance table to be used for solar abundances.
        Either a string corresponding to a built-in table or an array
        of 30 floats corresponding to the abundances of each element
        relative to the abundance of H. Default is "angr".
        Built-in options are:
        "angr" : from Anders E. & Grevesse N. (1989, Geochimica et
        Cosmochimica Acta 53, 197)
        "aspl" : from Asplund M., Grevesse N., Sauval A.J. & Scott
        P. (2009, ARAA, 47, 481)
        "wilm" : from Wilms, Allen & McCray (2000, ApJ 542, 914
        except for elements not listed which are given zero abundance)
        "lodd" : from Lodders, K (2003, ApJ 591, 1220)
    nei : boolean, optional
        If True, use the non-equilibrium ionization tables. These are
        not supplied with pyXSIM/SOXS but must be downloaded separately, in
        which case the *apec_root* parameter must also be set to their
        location. Default: False

    Examples
    --------
    >>> apec_model = TableApecModel(0.05, 50.0, 1000, apec_vers="3.0.3",
    ...                             thermal_broad=False)
    """
    def __init__(self, emin, emax, nchan, var_elem=None,
                 model_root=None, model_vers=None,
                 thermal_broad=True, nolines=False,
                 abund_table="angr", nei=False):
        super(TableApecModel, self).__init__(emin, emax, nchan, var_elem=var_elem,
                                             apec_root=model_root, apec_vers=model_vers,
                                             broadening=thermal_broad, nolines=nolines,
                                             abund_table=abund_table, nei=nei)
        self.nchan = self.nbins

    def prepare_spectrum(self, zobs):
        """
        Prepare the thermal model for execution given a redshift *zobs* for the spectrum.
        """
        cosmic_spec, metal_spec, var_spec = \
            self._get_table(list(range(self.nT)), zobs, 0.0)
        self.cosmic_spec = YTArray(cosmic_spec, "cm**3/s")
        self.metal_spec = YTArray(metal_spec, "cm**3/s")
        if var_spec is None:
            self.var_spec = var_spec
        else:
            self.var_spec = YTArray(var_spec, "cm**3/s")

    def get_spectrum(self, kT):
        """
        Get the thermal emission spectrum given a temperature *kT* in keV.
        """
        tindex = np.searchsorted(self.Tvals, kT)-1
        var_spec = None
        if tindex >= self.Tvals.shape[0]-1 or tindex < 0:
            cosmic_spec = YTArray(np.zeros(self.nchan), "cm**3/s")
            metal_spec = cosmic_spec
            if self.var_spec is not None:
                var_spec = YTArray(np.zeros((self.num_var_elem, self.nchan)), "cm**3/s")
        else:
            dT = (kT-self.Tvals[tindex])/self.dTvals[tindex]
            cspec_l = self.cosmic_spec[tindex, :]
            mspec_l = self.metal_spec[tindex, :]
            cspec_r = self.cosmic_spec[tindex+1, :]
            mspec_r = self.metal_spec[tindex+1, :]
            cosmic_spec = cspec_l*(1.-dT)+cspec_r*dT
            metal_spec = mspec_l*(1.-dT)+mspec_r*dT
            if self.var_spec is not None:
                vspec_l = self.var_spec[:, tindex, :]
                vspec_r = self.var_spec[:, tindex+1, :]
                var_spec = vspec_l*(1.-dT) + vspec_r*dT
        return cosmic_spec, metal_spec, var_spec

    def return_spectrum(self, temperature, metallicity, redshift, norm,
                        velocity=0.0, elem_abund=None):
        """
        Given the properties of a thermal plasma, return a spectrum.

        Parameters
        ----------
        temperature : float
            The temperature of the plasma in keV.
        metallicity : float
            The metallicity of the plasma in solar units.
        redshift : float
            The redshift of the plasma.
        norm : float
            The normalization of the model, in the standard Xspec units of
            1.0e-14*EM/(4*pi*(1+z)**2*D_A**2).
        velocity : float, optional
            Velocity broadening parameter in km/s. Default: 0.0
        elem_abund : dict of element name, float pairs
            A dictionary of elemental abundances to vary
            freely of the abund parameter. Default: None
        """
        spec = super(TableApecModel, self).get_spectrum(temperature, metallicity,
                                                        redshift, norm, velocity=velocity,
                                                        elem_abund=elem_abund)
        return YTArray(spec.flux*spec.de, "photons/s/cm**2")

thermal_models = {"apec": TableApecModel}

class AbsorptionModel(object):
    _name = ""
    def __init__(self, nH, energy, cross_section):
        self.nH = YTQuantity(nH*1.0e22, "cm**-2")
        self.emid = YTArray(energy, "keV")
        self.sigma = YTArray(cross_section, "cm**2")

    def get_absorb(self, e):
        """
        Get the absorption spectrum.
        """
        sigma = np.interp(e, self.emid, self.sigma, left=0.0, right=0.0)
        return np.exp(-sigma*self.nH)

    def absorb_photons(self, eobs, prng=None):
        r"""
        Determine which photons will be absorbed by foreground
        galactic absorption.

        Parameters
        ----------
        eobs : array_like
            The energies of the photons in keV.
        prng : integer, :class:`~numpy.random.RandomState` object or :mod:`~numpy.random`, optional
            A pseudo-random number generator. Typically will only be specified
            if you have a reason to generate the same set of random numbers, such as for a
            test. Default is the :mod:`numpy.random` module.
        """
        prng = parse_prng(prng)
        n_events = eobs.size
        if n_events == 0:
            return np.array([], dtype='bool')
        detected = np.zeros(n_events, dtype='bool')
        nchunk = n_events // 100
        if nchunk == 0:
            nchunk = n_events
        k = 0
        pbar = get_pbar("Absorbing photons", n_events)
        while k < n_events:
            absorb = self.get_absorb(eobs[k:k+nchunk])
            nabs = absorb.size
            randvec = prng.uniform(size=nabs)
            detected[k:k+nabs] = randvec < absorb
            k += nabs
            pbar.update(k)
        pbar.finish()
        return detected

class TBabsModel(AbsorptionModel):
    r"""
    Initialize a Tuebingen-Boulder (Wilms, J., Allen, A., &
    McCray, R. 2000, ApJ, 542, 914) ISM absorption model.

    Parameters
    ----------
    nH : float
        The foreground column density in units of 10^22 cm^{-2}.

    Examples
    --------
    >>> tbabs_model = TBabsModel(0.1)
    """
    _name = "tbabs"
    def __init__(self, nH):
        self.nH = YTQuantity(nH, "1.0e22*cm**-2")

    def get_absorb(self, e):
        e = np.array(e)
        return get_tbabs_absorb(e, self.nH.v)

class WabsModel(AbsorptionModel):
    r"""
    Initialize a Wisconsin (Morrison and McCammon; ApJ 270, 119)
    absorption model.

    Parameters
    ----------
    nH : float
        The foreground column density in units of 10^22 cm^{-2}.

    Examples
    --------
    >>> wabs_model = WabsModel(0.1)
    """
    _name = "wabs"
    def __init__(self, nH):
        self.nH = YTQuantity(nH, "1.0e22*cm**-2")

    def get_absorb(self, e):
        e = np.array(e)
        return get_wabs_absorb(e, self.nH.v)

absorb_models = {"wabs": WabsModel,
                 "tbabs": TBabsModel}


import cv2 as cv
import numpy as np

haar_cascade =cv.CascadeClassifier('haar_face.xml')

people =['Indian','Non-Indian']
features =np.load('features.npy',allow_pickle=True)
labels =np.load('labels.npy')

face_recognizer=cv.face.LBPHFaceRecognizer_create()
face_recognizer.read('face_trained.yml')

img =cv.imread(r'F:\Machine learning verzeo\Major Project\dataset\test\1-SS687216.jpg')

gray =cv.cvtColor(img,cv.COLOR_BGR2GRAY)
cv.imshow('Person',gray)

m = -1.0
i = None
#Detect the face in the image
faces_rect =haar_cascade.detectMultiScale(gray,1.1,4)
for (x,y,w,h) in faces_rect:
    faces_roi=gray[y:y+h,x:x+w]

    label,confidence =face_recognizer.predict(faces_roi)
    if i is None:
        i=label
    if confidence>m:
        i=label;
        m=confidence

print(f'label ={people[i]} with a confidence of {m}')
cv.putText(img,str(people[i]),(20,20),cv.FONT_HERSHEY_COMPLEX,1.0,(0,255,0),2)
cv.rectangle(img,(x,y),(x+w,y+h),(0,255,0),1)

cv.imshow('Detected face',img)
cv.waitKey(0)


import traceback
import queue
from pathlib import Path
import cv2

import sys
import argparse
import os
from datetime import datetime, timedelta

import depthai as dai
import numpy as np
from collections import namedtuple
from math import sqrt, ceil
# from FPS import FPS

import time

# OAK-D palm detection
# https://github.com/geaxgx/oakd_palm_detection

# Install steps:
# pip install --extra-index-url https://artifacts.luxonis.com/artifactory/luxonis-python-snapshot-local/ depthai==0.0.2.1+6ec3f3181b4e46fa6a9f9b20a5b4a3dac5e876b4
# cd ..
# git clone https://github.com/geaxgx/oakd_palm_detection

# Status: working

class Anchor:
    def __init__(self, x_center=0, y_center=0, w=0, h=0):
        self.x_center = x_center
        self.y_center = y_center
        self.w = w
        self.h = h

class HandRegion:
    def __init__(self, pd_score, pd_box, pd_kps=0):
        self.pd_score = pd_score # Palm detection score
        self.pd_box = pd_box # Palm detection box [x, y, w, h] normalized
        self.pd_kps = pd_kps # Palm detection keypoints

    def print(self):
        attrs = vars(self)
        print('\n'.join("%s: %s" % item for item in attrs.items()))

SSDAnchorOptions = namedtuple('SSDAnchorOptions',[
        'num_layers',
        'min_scale',
        'max_scale',
        'input_size_height',
        'input_size_width',
        'anchor_offset_x',
        'anchor_offset_y',
        'strides',
        'aspect_ratios',
        'reduce_boxes_in_lowest_layer',
        'interpolated_scale_aspect_ratio',
        'fixed_anchor_size'])

def calculate_scale(min_scale, max_scale, stride_index, num_strides):
    if num_strides == 1:
        return (min_scale + max_scale) / 2
    else:
        return min_scale + (max_scale - min_scale) * stride_index / (num_strides - 1)

def generate_anchors(options):
    """
    option : SSDAnchorOptions
    # https://github.com/google/mediapipe/blob/master/mediapipe/calculators/tflite/ssd_anchors_calculator.cc
    """
    anchors = []
    layer_id = 0
    n_strides = len(options.strides)
    while layer_id < n_strides:
        anchor_height = []
        anchor_width = []
        aspect_ratios = []
        scales = []
        # For same strides, we merge the anchors in the same order.
        last_same_stride_layer = layer_id
        while last_same_stride_layer < n_strides and \
                options.strides[last_same_stride_layer] == options.strides[layer_id]:
            scale = calculate_scale(options.min_scale, options.max_scale, last_same_stride_layer, n_strides)
            if last_same_stride_layer == 0 and options.reduce_boxes_in_lowest_layer:
                # For first layer, it can be specified to use predefined anchors.
                aspect_ratios += [1.0, 2.0, 0.5]
                scales += [0.1, scale, scale]
            else:
                aspect_ratios += options.aspect_ratios
                scales += [scale] * len(options.aspect_ratios)
                if options.interpolated_scale_aspect_ratio > 0:
                    if last_same_stride_layer == n_strides -1:
                        scale_next = 1.0
                    else:
                        scale_next = calculate_scale(options.min_scale, options.max_scale, last_same_stride_layer+1, n_strides)
                    scales.append(sqrt(scale * scale_next))
                    aspect_ratios.append(options.interpolated_scale_aspect_ratio)
            last_same_stride_layer += 1

        for i,r in enumerate(aspect_ratios):
            ratio_sqrts = sqrt(r)
            anchor_height.append(scales[i] / ratio_sqrts)
            anchor_width.append(scales[i] * ratio_sqrts)

        stride = options.strides[layer_id]
        feature_map_height = ceil(options.input_size_height / stride)
        feature_map_width = ceil(options.input_size_width / stride)

        for y in range(feature_map_height):
            for x in range(feature_map_width):
                for anchor_id in range(len(anchor_height)):
                    x_center = (x + options.anchor_offset_x) / feature_map_width
                    y_center = (y + options.anchor_offset_y) / feature_map_height
                    new_anchor = Anchor(x_center=x_center, y_center=y_center)
                    if options.fixed_anchor_size:
                        new_anchor.w = 1.0
                        new_anchor.h = 1.0
                    else:
                        new_anchor.w = anchor_width[anchor_id]
                        new_anchor.h = anchor_height[anchor_id]
                    anchors.append(new_anchor)

        layer_id = last_same_stride_layer
    return anchors

# Create anchors
# https://github.com/google/mediapipe/blob/master/mediapipe/modules/palm_detection/palm_detection_cpu.pbtxt

anchor_options = SSDAnchorOptions(num_layers=4,
                                    min_scale=0.1484375,
                                    max_scale=0.75,
                                    input_size_height=128,
                                    input_size_width=128,
                                    anchor_offset_x=0.5,
                                    anchor_offset_y=0.5,
                                    strides=[8, 16, 16, 16],
                                    aspect_ratios= [1.0],
                                    reduce_boxes_in_lowest_layer=False,
                                    interpolated_scale_aspect_ratio=1.0,
                                    fixed_anchor_size=True)
anchors = generate_anchors(anchor_options)
print(f"{len(anchors)} anchors have been created")

def decode_bboxes(score_thresh, wi, hi, scores, bboxes, anchors):
    """
    wi, hi : NN input shape
    mediapipe/calculators/tflite/tflite_tensors_to_detections_calculator.cc
    # Decodes the detection tensors generated by the model, based on
    # the SSD anchors and the specification in the options, into a vector of
    # detections. Each detection describes a detected object.

    https://github.com/google/mediapipe/blob/master/mediapipe/modules/palm_detection/palm_detection_cpu.pbtxt :
    node {
        calculator: "TensorsToDetectionsCalculator"
        input_stream: "TENSORS:detection_tensors"
        input_side_packet: "ANCHORS:anchors"
        output_stream: "DETECTIONS:unfiltered_detections"
        options: {
            [mediapipe.TensorsToDetectionsCalculatorOptions.ext] {
            num_classes: 1
            num_boxes: 896
            num_coords: 18
            box_coord_offset: 0
            keypoint_coord_offset: 4
            num_keypoints: 7
            num_values_per_keypoint: 2
            sigmoid_score: true
            score_clipping_thresh: 100.0
            reverse_output_order: true

            x_scale: 128.0
            y_scale: 128.0
            h_scale: 128.0
            w_scale: 128.0
            min_score_thresh: 0.5
            }
        }
    }
    """
    sigmoid_scores = 1 / (1 + np.exp(-scores))
    regions = []
    for i,anchor in enumerate(anchors):
        score = sigmoid_scores[i]

        if score > score_thresh:
            # If reverse_output_order is true, sx, sy, w, h = bboxes[i,:4]
            # Here reverse_output_order is true

            sx, sy, w, h = bboxes[i,:4]
            cx = sx * anchor.w / wi + anchor.x_center
            cy = sy * anchor.h / hi + anchor.y_center
            w = w * anchor.w / wi
            h = h * anchor.h / hi
            box = [cx - w*0.5, cy - h*0.5, w, h]

            kps = {}
            # 0 : wrist
            # 1 : index finger joint
            # 2 : middle finger joint
            # 3 : ring finger joint
            # 4 : little finger joint
            # 5 :
            # 6 : thumb joint
            for j, name in enumerate(["0", "1", "2", "3", "4", "5", "6"]):
                # Here reverse_output_order is true
                lx, ly = bboxes[i,4+j*2:6+j*2]
                lx = lx * anchor.w / wi + anchor.x_center
                ly = ly * anchor.h / hi + anchor.y_center
                kps[name] = [lx, ly]
            regions.append(HandRegion(float(score), box, kps))
    return regions

def non_max_suppression(regions, nms_thresh):

    # cv2.dnn.NMSBoxes(boxes, scores, 0, nms_thresh) needs:
    # boxes = [ [x, y, w, h], ...] with x, y, w, h of type int
    # Currently, x, y, w, h are float between 0 and 1, so we arbitrarily multiply by 1000 and cast to int
    # boxes = [r.box for r in regions]
    boxes = [ [int(x*1000) for x in r.pd_box] for r in regions]
    scores = [r.pd_score for r in regions]
    indices = cv2.dnn.NMSBoxes(boxes, scores, 0, nms_thresh)
    return [regions[i[0]] for i in indices]

frame = None
bboxes = []

# fps = FPS()

def to_planar(arr: np.ndarray, shape: tuple) -> list:
    return [val for channel in cv2.resize(arr, shape).transpose(2, 0, 1) for y_col in channel for val in y_col]

# nn data, being the bounding box locations, are in <0..1> range - they need to be normalized with frame width/height

def frame_norm3(frame, bbox):
    # frame.shape[:2],
    fh, fw, d = frame.shape
    sw = fw // 128
    sh = fh // 128
    x = bbox[0] * sw
    y = bbox[1] * sh
    w = bbox[2] * sw
    h = bbox[3] * sh
    return [x,y,w,h]

pd_score_thresh = 0.5
pd_nms_thresh = 0.3

device = None
q_rgb = None
q_nn = None
q_in = None

debug=True

useOAKDCam=False

def init_model(transform):

    if transform == 'oakd_palm':
        # # Start defining a pipeline
        # pipeline = dai.Pipeline()

        # if useOAKDCam:
        #     # Define a source - color camera
        #     cam_rgb = pipeline.createColorCamera()
        #     cam_rgb.setPreviewSize(128, 128)
        #     cam_rgb.setFps(90.0)
        #     cam_rgb.setInterleaved(False)

        # # Define a neural network that will make predictions based on the source frames
        # detection_nn = pipeline.createNeuralNetwork()
        # detection_nn.setBlobPath(str(Path("../oakd_palm_detection/models/palm_detection.blob").resolve().absolute()))

        # cam_rgb.preview.link(detection_nn.input)

        # # Create outputs
        # xout_rgb = pipeline.createXLinkOut()
        # xout_rgb.setStreamName("rgb")
        # cam_rgb.preview.link(xout_rgb.input)

        # xout_nn = pipeline.createXLinkOut()
        # xout_nn.setStreamName("nn")
        # detection_nn.out.link(xout_nn.input)

        # # Pipeline defined, now the device is assigned and pipeline is started
        # device = dai.Device(pipeline)
        # device.startPipeline()

        # if useOAKDCam:
        #     # Output queues will be used to get the rgb frames and nn data from the outputs defined above
        #     q_rgb = device.getOutputQueue(name="rgb", maxSize=4, blocking=False)
        #     q_nn = device.getOutputQueue(name="nn", maxSize=4, blocking=False)

    # return (q_rgb,q_nn), None
        return None, None

def process_image(transform,processing_model,img):
    global useOAKDCam, bboxes, results, pd_score_thresh, pd_nms_thresh, bboxes, anchors, device, q_rgb, q_nn, fps, q_in
    tracks = []
    # (q_rgb,q_nn) = processing_model
    try:
        # if useOAKDCam:
        # #     ret, frame = cap.read()
        #     frame = np.array(cam_out.get().getData()).reshape((3, 300, 300)).transpose(1, 2, 0).astype(np.uint8)
        #     shape = (3, frame.getHeight(), frame.getWidth())
        #     frame = in_rgb.getData().reshape(shape).transpose(1, 2, 0).astype(np.uint8)
        #     frame = np.ascontiguousarray(frame)
        # else:
        frame = img

        #palm detection https://github.com/geaxgx/oakd_palm_detection
        if transform == 'oakd_palm':
            if device is None:
                    # Start defining a pipeline
                    pipeline = dai.Pipeline()

                    if useOAKDCam:
                        # Define a source - color camera
                        cam_rgb = pipeline.createColorCamera()
                        cam_rgb.setPreviewSize(128, 128)
                        cam_rgb.setFps(90.0)
                        cam_rgb.setInterleaved(False)

                    # Define a neural network that will make predictions based on the source frames
                    detection_nn = pipeline.createNeuralNetwork()
                    detection_nn.setBlobPath(str(Path("../oakd_palm_detection/models/palm_detection.blob").resolve().absolute()))

                    if useOAKDCam:
                        cam_rgb.preview.link(detection_nn.input)
                    else:
                        detection_in = pipeline.createXLinkIn()
                        detection_in.setStreamName("detection_in")
                        detection_in.out.link(detection_nn.input)

                    # Create outputs
                    if useOAKDCam:
                        xout_rgb = pipeline.createXLinkOut()
                        xout_rgb.setStreamName("rgb")
                        cam_rgb.preview.link(xout_rgb.input)

                    xout_nn = pipeline.createXLinkOut()
                    xout_nn.setStreamName("nn")
                    detection_nn.out.link(xout_nn.input)

                    # Pipeline defined, now the device is assigned and pipeline is started
                    device = dai.Device(pipeline)
                    device.startPipeline()

                    if useOAKDCam:
                        # Output queues will be used to get the rgb frames and nn data from the outputs defined above
                        q_rgb = device.getOutputQueue(name="rgb", maxSize=4, blocking=False)
                    else:
                        q_in = device.getInputQueue("detection_in")

                    q_nn = device.getOutputQueue(name="nn", maxSize=4, blocking=False)

            # fps.update()
            # if frame is not None:

            if not useOAKDCam:
                nn_data = dai.NNData()
                nn_data.setLayer("input", to_planar(frame, (128, 128)))
                q_in.send(nn_data)

                # in_nn = q_nn.get()
                in_nn = q_nn.tryGet()
                # 2 output layers:
                # - classificators:
                # - regressors :
                # From: print(in_nn.getAllLayerNames())

                if in_nn is not None:
                    scores = np.array(in_nn.getLayerFp16("classificators"))
                    bboxes = np.array(in_nn.getLayerFp16("regressors")).reshape((896,18))

                    # Decode bboxes
                    regions = decode_bboxes(pd_score_thresh, 128, 128, scores, bboxes, anchors)
                    # Non maximum suppression
                    regions = non_max_suppression(regions, pd_nms_thresh)
                    tracks = regions
                    for r in regions:
                        raw_bbox = (np.array(r.pd_box) * 128).astype(int)
                        # box = raw_bbox
                        # print("raw_bbox",raw_bbox)
                        # print("frame.shape",frame.shape)
                        box = frame_norm3(frame, raw_bbox)
                        # print("box3",box)
                        # cv2.rectangle(frame, (box[0], box[1]), (box[2], box[3]), (255, 255, 255), 2)
                        cv2.rectangle(frame, (box[0], box[1]), (box[0]+box[2], box[1]+box[3]), (255,255,0), 2)

                    # if frame is not None:
                    #     img = frame
                    if frame is not None:
                        # cv2.putText(frame, "FPS: {:.2f}".format(fps.get()), (10,10), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0,0,255), 1)
                        # cv2.imshow("rgb", frame)
                        img = frame

            else:
                # in_rgb = q_rgb.tryGet()

                in_rgb = q_rgb.get()

                if in_rgb is not None:
                    # if the data from the rgb camera is available, transform the 1D data into a HxWxC frame
                    shape = (3, in_rgb.getHeight(), in_rgb.getWidth())

                    frame = in_rgb.getData().reshape(shape).transpose(1, 2, 0).astype(np.uint8)
                    frame = np.ascontiguousarray(frame)
                    in_nn = q_nn.get()
                    # 2 output layers:
                    # - classificators:
                    # - regressors :
                    # From: print(in_nn.getAllLayerNames())

                    if in_nn is not None:
                        scores = np.array(in_nn.getLayerFp16("classificators"))
                        bboxes = np.array(in_nn.getLayerFp16("regressors")).reshape((896,18))

                        # Decode bboxes
                        regions = decode_bboxes(pd_score_thresh, 128, 128, scores, bboxes, anchors)
                        # Non maximum suppression
                        regions = non_max_suppression(regions, pd_nms_thresh)
                        tracks = regions
                        for r in regions:
                            box = (np.array(r.pd_box) * 128).astype(int)
                            cv2.rectangle(frame, (box[0], box[1]), (box[0]+box[2], box[1]+box[3]), (255,255,0), 2)

                    # if frame is not None:
                    #     img = frame
                    if frame is not None:
                        # cv2.putText(frame, "FPS: {:.2f}".format(fps.get()), (10,10), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0,0,255), 1)
                        # cv2.imshow("rgb", frame)
                        img = frame

                    # if cv2.waitKey(1) == ord('q'):
                    #     pass

    except Exception as e:
        track = traceback.format_exc()
        print(track)
        print("OAK-D Exception",e)
        pass

    return tracks,img

def create_pipeline_palm():
    global useOAKDCam
    print("Creating pipeline...")
    pipeline = depthai.Pipeline()

    if useOAKDCam:
        # ColorCamera
        print("Creating Color Camera...")
        cam = pipeline.createColorCamera()
        cam.setPreviewSize(300, 300)
        cam.setResolution(depthai.ColorCameraProperties.SensorResolution.THE_1080_P)
        cam.setInterleaved(False)
        cam.setBoardSocket(depthai.CameraBoardSocket.RGB)
        cam_xout = pipeline.createXLinkOut()
        cam_xout.setStreamName("cam_out")
        cam.preview.link(cam_xout.input)

    # Define a neural network that will make predictions based on the source frames
    detection_nn = pipeline.createNeuralNetwork()
    detection_nn.setBlobPath(str(Path("../oakd_palm_detection/models/palm_detection.blob").resolve().absolute()))
    detection_nn_xout = pipeline.createXLinkOut()
    detection_nn_xout.setStreamName("detection_nn")
    detection_nn.out.link(detection_nn_xout.input)

    if useOAKDCam:
        cam.preview.link(detection_nn.input)
    else:
        detection_in = pipeline.createXLinkIn()
        detection_in.setStreamName("detection_in")
        detection_in.out.link(detection_nn.input)

    print("Pipeline created.")
    return pipeline