"""
Extract SEQUENTIAL embeddings for Arabic Native vs Non-Native 1D-CNN classifier.
Outputs a 2D array (Time_Steps, 1024) instead of a flattened 1D array.
Uses XLS-R 53 EARLY-TO-MID layers (4,6,8,10).
"""

import pandas as pd
import numpy as np
import torch
import librosa
import os
from transformers import Wav2Vec2Model, AutoFeatureExtractor
from tqdm import tqdm
from scipy.signal import butter, lfilter

# Prosody-focused layers
# model_name = "facebook/wav2vec2-large-xlsr-53"
model_name = "facebook/wav2vec2-lv-60-espeak-cv-ft"
PROSODY_LAYERS = (4, 6, 8, 10)

_feature_extractor = None
_model = None
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def _get_wav2vec():
    global _feature_extractor, _model
    if _model is None:
        _feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
        _model = Wav2Vec2Model.from_pretrained(model_name)
        _model.eval()
        _model.to(device)
    return _feature_extractor, _model


# The exact feature size of the Wav2Vec2 layers
WAV2VEC_DIM = 1024


def get_embeddings(audio_path):
    """
    Extracts the chronological 2D sequence: (Time, 1024)
    No averaging. No squashing.
    """
    feature_extractor, model = _get_wav2vec()
    try:
        speech, sr = librosa.load(audio_path, sr=16000)

        input_values = feature_extractor(
            speech, return_tensors="pt", sampling_rate=16000
        ).input_values.to(device)

        with torch.no_grad():
            outputs = model(input_values, output_hidden_states=True)

        h = outputs.hidden_states
        num_layers = len(h) - 1
        layer_idx = [min(max(l, 1), num_layers) for l in PROSODY_LAYERS]
        embs = [h[i] for i in layer_idx]

        # Stack the target layers and average them together (Layer Average, NOT Time Average)
        # Shape becomes (Time, 1024)
        time_seq = torch.stack(embs, dim=0).mean(dim=0).squeeze(0)

        # Convert directly to numpy and return the 2D matrix
        emb = time_seq.cpu().numpy().astype(np.float32)
        return emb

    except Exception as e:
        print(f"Error on {audio_path}: {e}")
        return None


def get_embeddings_from_audio(speech, sr=16000):
    """Extract 2D sequence (Time, 1024) from raw audio array. For test-time use."""
    feature_extractor, model = _get_wav2vec()
    try:
        input_values = feature_extractor(
            speech, return_tensors="pt", sampling_rate=sr
        ).input_values.to(device)
        with torch.no_grad():
            outputs = model(input_values, output_hidden_states=True)
        h = outputs.hidden_states
        num_layers = len(h) - 1
        layer_idx = [min(max(l, 1), num_layers) for l in PROSODY_LAYERS]
        embs = [h[i] for i in layer_idx]
        time_seq = torch.stack(embs, dim=0).mean(dim=0).squeeze(0)
        return time_seq.cpu().numpy().astype(np.float32)
    except Exception:
        return None


if __name__ == "__main__":
    import argparse

    ap = argparse.ArgumentParser(
        description="Extract 2D SEQUENTIAL embeddings for 1D-CNN."
    )
    ap.add_argument(
        "--manifest",
        default="manifests/processed_manifest.csv",
        help="Input manifest CSV (chunk_path, original_id, label, language)",
    )
    ap.add_argument(
        "--output-dir", default="embeddings", help="Folder for .npy embeddings"
    )
    ap.add_argument(
        "--output-manifest",
        default=None,
        help="Output manifest CSV",
    )
    args = ap.parse_args()

    manifest_path = args.manifest
    output_folder = args.output_dir

    if args.output_manifest:
        output_manifest_path = args.output_manifest
    else:
        if manifest_path.strip() == "manifests/processed_manifest.csv":
            output_manifest_path = "manifests/final_training_manifest.csv"
        else:
            d = os.path.dirname(manifest_path)
            base = (
                os.path.basename(manifest_path)
                .replace(".csv", "")
                .replace("_processed", "_embeddings")
            )
            output_manifest_path = os.path.join(d, base + ".csv")

    os.makedirs(output_folder, exist_ok=True)
    df = pd.read_csv(manifest_path)
    embedding_paths = []

    print(
        f"Extracting 2D Sequential Embeddings for {len(df)} chunks (Shape: Time x {WAV2VEC_DIM})..."
    )

    for idx, row in tqdm(df.iterrows(), total=len(df)):
        chunk_path = str(row["chunk_path"]).replace("\\", "/")
        emb = get_embeddings(chunk_path)

        if emb is not None:
            emb_filename = os.path.basename(chunk_path).replace(".wav", ".npy")
            emb_path = os.path.join(output_folder, emb_filename)
            np.save(emb_path, emb)
            embedding_paths.append(emb_path)
        else:
            embedding_paths.append(None)

    df["embedding_path"] = embedding_paths
    df.dropna(subset=["embedding_path"], inplace=True)
    df.to_csv(output_manifest_path, index=False)

    # Save the dimension size so the training script knows what to expect (use separate file if custom manifest to avoid overwriting main embedding_dim.txt)
    dim_name = "embedding_dim_cnn.txt" if args.output_manifest else "embedding_dim.txt"
    dim_file = os.path.join(os.path.dirname(output_manifest_path), dim_name)
    with open(dim_file, "w") as f:
        f.write(str(WAV2VEC_DIM))

    print(
        f"Done! 2D Embeddings in {output_folder}, manifest at {output_manifest_path}."
    )