#!/usr/bin/env python

import argparse
import dataclasses

import numpy as np
from sklearn.cluster import AgglomerativeClustering
from sklearn.preprocessing import normalize
from speechbrain.inference.speaker import EncoderClassifier
import torch
import torchaudio
import whisper

@dataclasses.dataclass
class Segment:
    start: np.float32
    end: np.float32
    text: str

def srt_timestamp(milis):
    h = milis // (1000 * 60 * 60)
    milis %= 1000 * 60 * 60
    m = milis // (1000 * 60)
    milis %= 1000 * 60
    s = milis // 1000
    milis %= 1000
    return f"{h:02}:{m:02}:{s:02},{milis:03}"

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("input")
    parser.add_argument("-o", "--output", default="output.srt")
    args = parser.parse_args()

    model = whisper.load_model("turbo")
    result = model.transcribe(args.input, word_timestamps=True, language="en")
    segments = []
    for seg in result["segments"]:
        words = seg["words"]
        start = words[0]["start"]
        end = words[-1]["end"]
        segments.append(Segment(start, end, seg["text"].strip()))

    print("Transcription done. Starting speaker classification")

    classifier = EncoderClassifier.from_hparams(
        source="speechbrain/spkrec-ecapa-voxceleb",
    )

    signal, sample_rate = torchaudio.load(args.input)

    print(sample_rate)
    # resample to 16khz
    if sample_rate != 16000:
        resampler = torchaudio.transforms.Resample(sample_rate, 16000)
        signal = resampler(signal)
        sample_rate = 16000

    # convert to mono if needed
    if signal.shape[0] > 1:
        signal = torch.mean(signal, dim=0, keepdim=True)

    embeddings = []
    for seg in segments:
        start_sample = int(seg.start * sample_rate)
        end_sample = int(seg.end * sample_rate)

        chunk = signal[:, start_sample:end_sample]

        # SpeechBrain expects [batch, time]
        emb = classifier.encode_batch(chunk)
        emb = emb.squeeze().detach().cpu().numpy()

        embeddings.append(emb)
    embeddings = np.vstack(embeddings)
    embeddings = normalize(embeddings)

    n_speakers = 2
    clustering = AgglomerativeClustering(
        n_clusters=n_speakers,
        metric="cosine",
        linkage="average"
    )
    labels = clustering.fit_predict(embeddings)

    for seg, label in zip(segments, labels):
        seg.text = f"Speaker {label + 1}: {seg.text}"

    with open(args.output, "w") as srt_file:
        for i, seg in enumerate(segments, 1):
            start = srt_timestamp(round(seg.start * 1000))
            end = srt_timestamp(round(seg.end * 1000))
            srt_file.write(f"{i}\n{start} --> {end}\n{seg.text}\n\n")
            print(start, end, seg.text)

    from sklearn.decomposition import PCA
    import matplotlib.pyplot as plt

    pca = PCA(n_components=2)
    points = pca.fit_transform(embeddings)

    plt.scatter(points[:,0], points[:,1])
    plt.show()

if __name__ == "__main__":
    main()