import librosa
import numpy as np


def temporal_difference(audio_path, method=1):
    audio, sr = librosa.load(audio_path, sr=None)

    mel = librosa.feature.melspectrogram(
        y=audio, sr=sr, n_fft=1024, hop_length=256, n_mels=80
    )

    log_mel = librosa.power_to_db(mel, ref=np.max)

    if method == 1:
        d = np.abs(log_mel[:, :-1] - log_mel[:, 1:])
        return d.mean()

    tmp = log_mel.sum(0)
    d = np.abs(tmp[:-1] - tmp[1:])

    return d.mean()


def solution(train_df, test_df, method=1):
    train_df = train_df.copy()
    test_df = test_df.copy()

    test_df["mean_diff"] = [
        temporal_difference(path, method) for path in test_df["audio_path"]
    ]

    if method < 3:
        return -test_df["mean_diff"]

    train_df["mean_diff"] = [
        temporal_difference(path, method) for path in train_df["audio_path"]
    ]

    speaker_stats = (
        train_df.groupby("speaker_id")["mean_diff"]
        .agg(["mean", "std"])
        .rename(columns={"mean": "speaker_mean", "std": "speaker_std"})
    )

    test_df = test_df.join(speaker_stats, on="speaker_id")

    test_df["z_score"] = (test_df["mean_diff"] - test_df["speaker_mean"]) / (
        test_df["speaker_std"] + 1e-8
    )

    return -test_df["z_score"]