from __future__ import annotations

import io
import os
import subprocess
import tempfile
import time
import warnings
from subprocess import Popen, PIPE
from typing import NamedTuple, Callable

import matplotlib.pyplot as plt
import numpy as np
import scipy.io.wavfile
from PIL import Image
from inaSpeechSegmenter import *
from matplotlib.axes import Axes
from matplotlib.figure import Figure
import tensorflow as tf


gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

seg = Segmenter()


class ResultFrame(NamedTuple):
    gender: str
    start: float
    end: float
    prob: float


class Result(NamedTuple):
    frames: list[ResultFrame]
    file: str


class BatchResults(NamedTuple):
    results: list[Result]
    time_full: float
    time_avg: float
    successes: int
    messages: list[tuple[str, int]]


def _process_helper(args: tuple[Segmenter, str]) -> Result:
    seg, f = args
    lseg = seg(f)
    return Result([ResultFrame(*s) for s in lseg], f)


def process(self: Segmenter, inp: list[str], tmpdir=None, verbose=False, skip_if_exist=False,
            nbtry=1, try_delay=2.) -> list[Result]:
    from tqdm.contrib.concurrent import process_map
    import tqdm
    results = []
    for i in tqdm.tqdm(inp):
        results.append(_process_helper((self, i)))
    return results

    # return process_map(_process_helper, [(self, i) for i in inp], max_workers=2)

    #
    # t_batch_start = time.time()
    #
    # results: list[Result] = []
    # lmsg = []
    # fg = featGenerator(inp.copy(), inp.copy(), tmpdir, self.ffmpeg, skip_if_exist, nbtry, try_delay)
    # i = 0
    # for feats, msg in fg:
    #     lmsg += msg
    #     i += len(msg)
    #     if verbose:
    #         print('%d/%d' % (i, len(inp)), msg)
    #     if feats is None:
    #         break
    #     mspec, loge, diff_len = feats
    #     lseg = self.segment_feats(mspec, loge, diff_len, 0)
    #     results.append()
    #
    # t_batch_dur = time.time() - t_batch_start
    # nb_processed = len([e for e in lmsg if e[1] == 0])
    # avg = t_batch_dur / nb_processed if nb_processed else -1
    # return BatchResults(results, t_batch_dur, avg, nb_processed, lmsg)


def to_wav(file: str, callback: Callable, start_sec: float = 0, stop_sec: float = 0):
    """
    Convert media to temp wav 16k file and return features
    """
    base, _ = os.path.splitext(os.path.basename(file))

    with tempfile.TemporaryDirectory() as tmpdir_name:
        # build ffmpeg command line
        tmp_wav = tmpdir_name + '/' + base + '.wav'
        args = ['ffmpeg', '-y', '-i', file, '-ar', '16000', '-ac', '1']

        if start_sec != 0:
            args += ['-ss', '%f' % start_sec]
        if stop_sec != 0:
            args += ['-to', '%f' % stop_sec]

        args += [tmp_wav]

        # launch ffmpeg
        p = Popen(args, stdout=PIPE, stderr=PIPE)
        output, error = p.communicate()
        assert p.returncode == 0, error

        return callback(tmp_wav)


def show_image_buffer(buf):
    im = Image.open(buf)
    im.show()
    buf.close()


def draw_result(file: str, result: list[ResultFrame]):
    """
    Draw segmentation result

    :param file: Audio file
    :param result: Segmentation result
    :return: Result image in bytes (please close it after use)
    """
    def wav_callback(wavfile: str):
        sample_rate, audio = scipy.io.wavfile.read(wavfile)
        _time = np.linspace(0, len(audio) / sample_rate, num=len(audio))

        fig: Figure = plt.gcf()
        ax: Axes = plt.gca()

        # Plot audio
        plt.plot(_time, audio, color='white')

        # Set size
        # fig.set_dpi(400)
        fig.set_size_inches(18, 6)

        # Cutoff frequency so that the plot looks centered
        cutoff = min(abs(min(audio)), abs(max(audio)))
        ax.set_ylim([-cutoff, cutoff])
        ax.set_xlim([result[0].start, result[-1].end])

        # Draw segmentation areas
        colors = {'female': '#F5A9B8', 'male': '#5BCEFA', 'default': 'gray'}
        for r in result:
            color = colors[r.gender] if r.gender in colors else colors['default']
            ax.axvspan(r.start, r.end - 0.01, alpha=.5, color=color)

        # Savefig to bytes
        buf = io.BytesIO()
        plt.axis('off')
        plt.savefig(buf, bbox_inches='tight', pad_inches=0, transparent=False)
        buf.seek(0)
        plt.clf()
        plt.close()
        return buf

    return to_wav(file, wav_callback)


def get_result_percentages(result: list[ResultFrame]) -> tuple[float, float, float, float]:
    """
    Get percentages

    :param result: Result
    :return: %female, %male, %other, %female-vs-female+male
    """
    # Count total and categorical durations
    total_dur = 0
    durations: dict[str, int] = {f.gender: 0 for f in result}
    for f in result:
        dur = f.end - f.start
        durations[f.gender] += dur
        total_dur += dur

    # Convert durations to ratios
    for d in durations:
        durations[d] /= total_dur

    # Return results
    f = durations.get('female', 0)
    m = durations.get('male', 0)

    fm_total = f + m
    pf = 0 if fm_total == 0 else f / fm_total

    return f, m, 1 - f - m, pf


def test():
    # results: BatchResults = BatchResults(
    #     [Result([ResultFrame('female', 0.0, 10.48), ResultFrame('male', 10.48, 12.780000000000001)],
    #             '../test.csv')],
    #     1.7032792568206787, 1.7032792568206787, 1,
    #     [('../test.csv', 0)])

    warnings.filterwarnings("ignore")
    audio_file = '../test.flac'

    # Warmup run
    results = segment(audio_file)
    print(results)

    # # Actual run
    # results = process(seg, ['../test.flac'])
    # print(results)

    # Benchmark
    # iterations = 60
    # total_time = 0
    # audio_len = float(subprocess.getoutput(f'ffprobe -i {audio_file} -show_entries format=duration -v quiet -of csv="p=0"'))
    # print(f'Audio length: {audio_len}')
    #
    # for i in range(iterations):
    #     results = process(seg, [audio_file])
    #     total_time += results.time_full
    #
    # time_per_second = total_time / iterations / audio_len
    # print(f'Benchmark result: {total_time}s / {iterations} iterations = {time_per_second} seconds of processing per second in audio')
    # print(f'Score: {1 / time_per_second}')

    # Draw results
    # with draw_result(audio_file, results.results[0]) as buf:
    #     show_image_buffer(buf)
    #     print(get_result_percentages(results.results[0]))


if __name__ == '__main__':
    test()
    pass