# -*- coding: utf-8 -*-
import itertools
from io import BytesIO
from datetime import timedelta

import pytest
import numpy as np

from ffsubsync.sklearn_shim import make_pipeline
from ffsubsync.speech_transformers import SubtitleSpeechTransformer
from ffsubsync.subtitle_parser import GenericSubtitleParser
from ffsubsync.subtitle_transformers import SubtitleShifter

fake_srt = b"""1
00:00:00,178 --> 00:00:01,1416
<i>Previously on "Your favorite TV show..."</i>

2
00:00:01,1828 --> 00:00:04,549
Oh hi, Mark.

3
00:00:04,653 --> 00:00:03,3062
You are tearing me apart, Lisa!
"""

# Occasionally some srt files have timestamps whose 'milliseconds'
# field has more than 3 digits... Ideally we should test that these
# are handled properly with dedicated tests, but in the interest of
# development speed I've opted to sprinkle in a few >3 digit
# millisecond fields into the dummy string above in order to exercise
# this case integration-test style in the below unit tests.


@pytest.mark.parametrize("start_seconds", [0, 2, 4, 6])
def test_start_seconds(start_seconds):
    parser_zero = GenericSubtitleParser(start_seconds=0)
    parser_zero.fit(BytesIO(fake_srt))
    parser = GenericSubtitleParser(start_seconds=start_seconds)
    parser.fit(BytesIO(fake_srt))
    expected = [
        sub
        for sub in parser_zero.subs_
        if sub.start >= timedelta(seconds=start_seconds)
    ]
    assert all(esub == psub for esub, psub in zip(expected, parser.subs_))


@pytest.mark.parametrize("max_seconds", [1, 1.5, 2.0, 2.5])
def test_max_seconds(max_seconds):
    parser = GenericSubtitleParser(max_subtitle_seconds=max_seconds)
    parser.fit(BytesIO(fake_srt))
    assert max(sub.end - sub.start for sub in parser.subs_) <= timedelta(
        seconds=max_seconds
    )


@pytest.mark.parametrize("encoding", ["utf-8", "ascii", "latin-1"])
def test_same_encoding(encoding):
    parser = GenericSubtitleParser(encoding=encoding)
    offseter = SubtitleShifter(1)
    pipe = make_pipeline(parser, offseter)
    pipe.fit(BytesIO(fake_srt))
    assert parser.subs_._encoding == encoding
    assert offseter.subs_._encoding == parser.subs_._encoding
    assert offseter.subs_.set_encoding("same")._encoding == encoding
    assert offseter.subs_.set_encoding("utf-8")._encoding == "utf-8"


@pytest.mark.parametrize("offset", [1, 1.5, -2.3])
def test_offset(offset):
    parser = GenericSubtitleParser()
    offseter = SubtitleShifter(offset)
    pipe = make_pipeline(parser, offseter)
    pipe.fit(BytesIO(fake_srt))
    for sub_orig, sub_offset in zip(parser.subs_, offseter.subs_):
        assert (
            abs(
                sub_offset.start.total_seconds()
                - sub_orig.start.total_seconds()
                - offset
            )
            < 1e-6
        )
        assert (
            abs(sub_offset.end.total_seconds() - sub_orig.end.total_seconds() - offset)
            < 1e-6
        )


@pytest.mark.parametrize(
    "sample_rate,start_seconds", itertools.product([10, 20, 100, 300], [0, 2, 4, 6])
)
def test_speech_extraction(sample_rate, start_seconds):
    parser = GenericSubtitleParser(start_seconds=start_seconds)
    extractor = SubtitleSpeechTransformer(
        sample_rate=sample_rate, start_seconds=start_seconds
    )
    pipe = make_pipeline(parser, extractor)
    bitstring = pipe.fit_transform(BytesIO(fake_srt)).astype(bool)
    bitstring_shifted_left = np.append(bitstring[1:], [False])
    bitstring_shifted_right = np.append([False], bitstring[:-1])
    bitstring_cumsum = np.cumsum(bitstring)
    consec_ones_end_pos = np.nonzero(
        bitstring_cumsum
        * (bitstring ^ bitstring_shifted_left)
        * (bitstring_cumsum != np.cumsum(bitstring_shifted_right))
    )[0]
    prev = 0
    for pos, sub in zip(consec_ones_end_pos, parser.subs_):
        start = int(round(sub.start.total_seconds() * sample_rate))
        duration = sub.end.total_seconds() - sub.start.total_seconds()
        stop = start + int(round(duration * sample_rate))
        assert bitstring_cumsum[pos] - prev == stop - start
        prev = bitstring_cumsum[pos]


def test_max_time_found():
    parser = GenericSubtitleParser()
    extractor = SubtitleSpeechTransformer(sample_rate=100)
    pipe = make_pipeline(parser, extractor)
    pipe.fit(BytesIO(fake_srt))
    assert extractor.max_time_ == 6.062
