pipelines/data_pipeline.py

"""End-to-end data processing pipeline."""

import argparse
import logging
import sys
from pathlib import Path
from typing import List

sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from src.utils.logger import setup_logger

logger = setup_logger("data_pipeline")


def process_audio(input_dir: str, output_dir: str, config: dict = None) -> List[str]:
    """Process audio files."""
    from src.data_processing.audio_processor import AudioProcessor
    
    logger.info(f"Processing audio from {input_dir}")
    
    processor = AudioProcessor(
        sample_rate=config.get("sample_rate", 16000),
        n_fft=config.get("n_fft", 512),
        hop_length=config.get("hop_length", 160),
    )
    
    results = processor.batch_process(input_dir, output_dir)
    
    logger.info(f"Processed {len(results)} audio files")
    return [r["output_path"] for r in results]


def process_text(input_path: str, output_path: str, config: dict = None) -> str:
    """Process text data."""
    from src.data_processing.text_normalizer import MyanmarTextNormalizer
    
    logger.info(f"Processing text from {input_path}")
    
    normalizer = MyanmarTextNormalizer(
        custom_rules_path=config.get("custom_rules") if config else None
    )
    
    import pandas as pd
    
    if input_path.endswith(".csv"):
        df = pd.read_csv(input_path)
        texts = df["text"].tolist()
    else:
        raise ValueError(f"Unsupported input format: {input_path}")
    
    normalized_texts = normalizer.normalize_corpus(texts)
    df["text_normalized"] = normalized_texts
    
    df.to_csv(output_path, index=False)
    logger.info(f"Saved normalized text to {output_path}")
    
    return output_path


def augment_data(input_path: str, output_path: str, config: dict = None) -> str:
    """Augment training data."""
    from src.augmentation.synonym_replacer import MyanmarSynonymReplacer
    from src.augmentation.perturbator import TextPerturbator
    
    logger.info(f"Augmenting data from {input_path}")
    
    replacer = MyanmarSynonymReplacer()
    perturbator = TextPerturbator()
    
    import pandas as pd
    import json
    
    if input_path.endswith(".csv"):
        df = pd.read_csv(input_path)
        samples = df.to_dict("records")
    elif input_path.endswith(".json"):
        with open(input_path, "r") as f:
            samples = json.load(f)
    else:
        raise ValueError(f"Unsupported format: {input_path}")
    
    augmented = []
    
    for sample in samples:
        text = sample.get("text", "")
        
        # Synonym replacement
        aug_text, replacements = replacer.augment_text(
            text,
            replace_prob=config.get("synonym_prob", 0.3) if config else 0.3,
        )
        if replacements:
            aug_sample = sample.copy()
            aug_sample["text"] = aug_text
            aug_sample["augmentation_type"] = "synonym"
            aug_sample["replacements"] = replacements
            augmented.append(aug_sample)
        
        # Perturbation
        aug_text, perturbations = perturbator.apply_random_perturbations(
            text,
            n_perturbations=config.get("n_perturbations", 2) if config else 2,
        )
        if perturbations:
            aug_sample = sample.copy()
            aug_sample["text"] = aug_text
            aug_sample["augmentation_type"] = "perturbation"
            aug_sample["perturbations"] = [p.value for p in perturbations]
            augmented.append(aug_sample)
    
    # Save augmented data
    output_df = pd.DataFrame(augmented)
    output_df.to_csv(output_path, index=False)
    
    logger.info(f"Generated {len(augmented)} augmented samples")
    return output_path


def split_data(input_path: str, output_dir: str, config: dict = None) -> dict:
    """Split data into train/val/test."""
    from sklearn.model_selection import train_test_split
    
    logger.info(f"Splitting data from {input_path}")
    
    import pandas as pd
    
    df = pd.read_csv(input_path)
    
    train_ratio = config.get("train_ratio", 0.8) if config else 0.8
    val_ratio = config.get("val_ratio", 0.1) if config else 0.1
    
    # First split: train vs rest
    train_df, temp_df = train_test_split(
        df, train_size=train_ratio, random_state=42
    )
    
    # Second split: val vs test
    val_size = val_ratio / (1 - train_ratio)
    val_df, test_df = train_test_split(
        temp_df, train_size=val_size, random_state=42
    )
    
    # Save splits
    Path(output_dir).mkdir(parents=True, exist_ok=True)
    
    train_df.to_csv(f"{output_dir}/train.csv", index=False)
    val_df.to_csv(f"{output_dir}/val.csv", index=False)
    test_df.to_csv(f"{output_dir}/test.csv", index=False)
    
    logger.info(f"Saved: train={len(train_df)}, val={len(val_df)}, test={len(test_df)}")
    
    return {
        "train": f"{output_dir}/train.csv",
        "val": f"{output_dir}/val.csv",
        "test": f"{output_dir}/test.csv",
    }


def run_pipeline(input_path: str, output_dir: str, config: dict = None) -> dict:
    """Run the full data processing pipeline."""
    logger.info("Starting data processing pipeline...")
    
    Path(output_dir).mkdir(parents=True, exist_ok=True)
    
    # Step 1: Normalize text
    normalized_path = f"{output_dir}/normalized.csv"
    process_text(input_path, normalized_path, config)
    
    # Step 2: Augment data
    augmented_path = f"{output_dir}/augmented.csv"
    augment_data(normalized_path, augmented_path, config)
    
    # Step 3: Split data
    splits = split_data(augmented_path, f"{output_dir}/splits", config)
    
    logger.info("Pipeline complete!")
    
    return {
        "normalized": normalized_path,
        "augmented": augmented_path,
        "splits": splits,
    }


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Data processing pipeline")
    parser.add_argument("--input", type=str, required=True, help="Input data file")
    parser.add_argument("--output", type=str, default="data/processed", help="Output directory")
    
    args = parser.parse_args()
    
    run_pipeline(args.input, args.output)