src/preprocess_spark.py

"""
Preprocesare dataset TTS-Romanian folosind PySpark.
Acest script poate fi extins cu filtre personalizate pe date.
"""
import os
import sys
from pathlib import Path

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, length, regexp_replace, lower, trim, when, count, avg, min as spark_min, max as spark_max


def create_spark_session(app_name: str = "TTS-Romanian-Preprocess", memory: str = "8g") -> SparkSession:
    """Creează o sesiune Spark optimizată pentru prelucrare audio."""
    spark = (
        SparkSession.builder
        .appName(app_name)
        .config("spark.driver.memory", memory)
        .config("spark.executor.memory", memory)
        .config("spark.sql.adaptive.enabled", "true")
        .config("spark.sql.adaptive.coalescePartitions.enabled", "true")
        .config("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
        .getOrCreate()
    )
    spark.sparkContext.setLogLevel("WARN")
    return spark


def load_dataset_spark(spark: SparkSession, dataset_path: str):
    """Încarcă dataset-ul într-un DataFrame Spark (presupunând format Parquet/CSV/JSON)."""
    path = Path(dataset_path)
    
    if not path.exists():
        raise FileNotFoundError(f"Dataset path not found: {dataset_path}")
    
    # Încercăm să detectăm formatul
    if any(path.glob("*.parquet")):
        df = spark.read.parquet(str(path / "*.parquet"))
    elif any(path.glob("*.csv")):
        df = spark.read.option("header", "true").csv(str(path / "*.csv"))
    elif any(path.glob("*.json")):
        df = spark.read.json(str(path / "*.json"))
    else:
        # Încercăm să citim direct din director (HuggingFace datasets salvează în format specific)
        df = spark.read.parquet(str(path / "**/*.parquet"))
    
    return df


def apply_basic_filters(df):
    """Filtre de bază pentru calitatea datelor TTS."""
    print(f"[INFO] Înainte de filtre: {df.count()} înregistrări")
    
    # Eliminăm textele goale sau foarte scurte
    df = df.filter((col("text").isNotNull()) & (length(trim(col("text"))) > 5))
    
    # Eliminăm textele foarte lungi (posibile erori)
    df = df.filter(length(col("text")) < 500)
    
    # Curățăm whitespace-ul excesiv
    df = df.withColumn("text", regexp_replace(trim(col("text")), r"\s+", " "))
    
    # Eliminăm caractere non-romanesti comune (opțional, de activat dacă e nevoie)
    # df = df.withColumn("text", regexp_replace(col("text"), r"[^a-zA-ZăâîșțĂÂÎȘȚ0-9\s\.,;:!?\"'-]", ""))
    
    print(f"[INFO] După filtre de bază: {df.count()} înregistrări")
    return df


def compute_statistics(df):
    """Calculează statistici pe dataset."""
    stats = df.select(
        count("*").alias("total_records"),
        avg(length(col("text"))).alias("avg_text_length"),
        spark_min(length(col("text"))).alias("min_text_length"),
        spark_max(length(col("text"))).alias("max_text_length")
    ).collect()[0]
    
    print("\n[STATISTICI DATASET]")
    print(f"  Total înregistrări: {stats['total_records']}")
    print(f"  Lungime medie text: {stats['avg_text_length']:.2f}")
    print(f"  Lungime minimă text: {stats['min_text_length']}")
    print(f"  Lungime maximă text: {stats['max_text_length']}")
    
    return stats


def save_processed_dataset(df, output_path: str, format: str = "parquet"):
    """Salvează dataset-ul procesat."""
    os.makedirs(output_path, exist_ok=True)
    
    if format == "parquet":
        df.write.mode("overwrite").parquet(output_path)
    elif format == "csv":
        df.write.mode("overwrite").option("header", "true").csv(output_path)
    elif format == "json":
        df.write.mode("overwrite").json(output_path)
    else:
        raise ValueError(f"Format necunoscut: {format}")
    
    print(f"[INFO] Dataset salvat în: {output_path} (format: {format})")


def main():
    import argparse
    parser = argparse.ArgumentParser(description="Preprocesare TTS-Romanian cu PySpark")
    parser.add_argument("--input", required=True, help="Cale către datasetul brut")
    parser.add_argument("--output", required=True, help="Cale către directorul de output")
    parser.add_argument("--format", default="parquet", choices=["parquet", "csv", "json"])
    parser.add_argument("--memory", default="8g", help="Memorie alocată pentru Spark")
    args = parser.parse_args()
    
    spark = create_spark_session(memory=args.memory)
    
    try:
        df = load_dataset_spark(spark, args.input)
        print(f"[INFO] Coloane detectate: {df.columns}")
        df.printSchema()
        
        df = apply_basic_filters(df)
        compute_statistics(df)
        
        save_processed_dataset(df, args.output, args.format)
        
    finally:
        spark.stop()
        print("[INFO] Sesiune Spark închisă.")


if __name__ == "__main__":
    main()