"""
Tahap 3: Split Dataset

Script ini membagi dataset menjadi 80% training dan 20% testing
secara stratified untuk setiap kelas.

Sesuai README:
- Input: dataset/normalized/
- Output: dataset/train/ dan dataset/test/
- Pembagian dilakukan untuk setiap kelas
- Struktur folder dibuat sama persis dengan folder normalized/

Persentase:
- Train: 80%
- Test: 20%
"""

import sys
import json
import shutil
import numpy as np
from pathlib import Path
from tqdm import tqdm
from typing import List, Tuple, Dict
import argparse
from collections import defaultdict, Counter

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from utils.config import (
    NORMALIZED_DIR,
    TRAIN_DIR,
    TEST_DIR,
    CATEGORIES,
    TRAIN_RATIO
)


class DatasetSplitter:
    """
    Splitter untuk membagi dataset menjadi train dan test sets.
    Pembagian dilakukan untuk setiap kelas secara stratified.
    """

    def __init__(self, train_ratio: float = TRAIN_RATIO, random_seed: int = 42):
        """
        Inisialisasi splitter.

        Args:
            train_ratio: Rasio data untuk training (0.8 = 80%)
            random_seed: Seed untuk random
        """
        self.train_ratio = train_ratio

        # Set random seed
        np.random.seed(random_seed)

    def split_samples(
        self,
        samples: List[Path],
        labels: List[str]
    ) -> Tuple[List[Path], List[str], List[Path], List[str]]:
        """
        Split samples menjadi train dan test sets secara stratified.

        Args:
            samples: List of sample paths
            labels: List of corresponding labels

        Returns:
            Tuple (train_samples, train_labels, test_samples, test_labels)
        """
        # Group samples by label
        label_to_samples = defaultdict(list)
        for sample, label in zip(samples, labels):
            label_to_samples[label].append(sample)

        train_samples = []
        train_labels = []
        test_samples = []
        test_labels = []

        # Split setiap kelas secara proporsional
        for label, class_samples in label_to_samples.items():
            # Shuffle samples
            indices = np.arange(len(class_samples))
            np.random.shuffle(indices)

            # Calculate split point
            n_train = int(len(class_samples) * self.train_ratio)
            n_train = max(1, n_train)  # Minimal 1 sample untuk train

            # Split
            train_indices = indices[:n_train]
            test_indices = indices[n_train:]

            # Jika test set kosong, ambil 1 dari train
            if len(test_indices) == 0 and len(train_indices) > 1:
                test_indices = [train_indices[-1]]
                train_indices = train_indices[:-1]

            train_samples.extend([class_samples[i] for i in train_indices])
            train_labels.extend([label] * len(train_indices))

            test_samples.extend([class_samples[i] for i in test_indices])
            test_labels.extend([label] * len(test_indices))

        return train_samples, train_labels, test_samples, test_labels

    def copy_samples(
        self,
        samples: List[Path],
        labels: List[str],
        output_dir: Path,
        category: str,
        split: str
    ) -> int:
        """
        Copy samples ke output directory.

        Args:
            samples: List of sample paths
            labels: List of labels
            output_dir: Root output directory
            category: Category name
            split: 'train' atau 'test'

        Returns:
            Number of copied files
        """
        copied_count = 0

        for sample_path, label in tqdm(
            zip(samples, labels),
            total=len(samples),
            desc=f"Copying {split}"
        ):
            try:
                # Create output path (struktur sama persis dengan normalized/)
                # Format: {train|test}/{category}/{label}/file.npy
                # output_dir sudah berupa dataset/train atau dataset/test
                output_path = output_dir / category / label / sample_path.name

                # Skip jika sudah ada
                if output_path.exists():
                    continue

                # Create directory
                output_path.parent.mkdir(parents=True, exist_ok=True)

                # Copy file
                shutil.copy2(sample_path, output_path)
                copied_count += 1

            except Exception as e:
                print(f"Error copying {sample_path}: {e}")

        return copied_count


def get_normalized_files(normalized_dir: Path, category: str) -> Tuple[List[Path], List[str]]:
    """
    Get semua file normalized dan labels-nya.

    Args:
        normalized_dir: Root normalized directory
        category: Category name

    Returns:
        Tuple (list of files, list of labels)
    """
    category_dir = normalized_dir / category

    if not category_dir.exists():
        return [], []

    files = []
    labels = []

    # Iterate semua subdirectories (labels)
    for label_dir in category_dir.iterdir():
        if label_dir.is_dir():
            label = label_dir.name

            # Get semua .npy files
            for npy_file in label_dir.glob("*.npy"):
                files.append(npy_file)
                labels.append(label)

    return files, labels


def split_category(
    normalized_dir: Path,
    train_dir: Path,
    test_dir: Path,
    category: str,
    train_ratio: float = TRAIN_RATIO
) -> Dict:
    """
    Split dataset untuk satu kategori.

    Args:
        normalized_dir: Directory normalized data
        train_dir: Directory untuk data train
        test_dir: Directory untuk data test
        category: Category name
        train_ratio: Train ratio

    Returns:
        Statistics dictionary
    """
    print(f"\n{'='*60}")
    print(f"Splitting dataset untuk kategori: {category}")
    print(f"{'='*60}\n")

    # Get files dan labels
    files, labels = get_normalized_files(normalized_dir, category)

    if not files:
        print(f"Warning: Tidak ada file ditemukan untuk kategori {category}")
        return {}

    print(f"Total samples: {len(files)}")

    # Count per class
    class_counts = Counter(labels)
    print(f"Classes: {len(class_counts)}")
    for label, count in sorted(class_counts.items()):
        print(f"  - {label}: {count} samples")

    # Split
    splitter = DatasetSplitter(train_ratio=train_ratio)
    train_files, train_labels, test_files, test_labels = splitter.split_samples(files, labels)

    print(f"\nAfter split:")
    print(f"  Train: {len(train_files)} samples ({len(train_files)/len(files)*100:.1f}%)")
    print(f"  Test: {len(test_files)} samples ({len(test_files)/len(files)*100:.1f}%)")

    # Copy files
    print("\nCopying files...")
    train_copied = splitter.copy_samples(
        train_files, train_labels, train_dir, category, "train"
    )
    test_copied = splitter.copy_samples(
        test_files, test_labels, test_dir, category, "test"
    )
    print(f"Copied: {train_copied} train, {test_copied} test files")

    # Generate label encoder
    unique_labels = sorted(set(labels))
    label_encoder = {label: idx for idx, label in enumerate(unique_labels)}

    # Save label encoder di folder models/
    model_dir = Path(__file__).parent.parent / "models"
    encoder_path = model_dir / f"{category}_label_encoder.json"
    encoder_path.parent.mkdir(parents=True, exist_ok=True)
    with open(encoder_path, 'w') as f:
        json.dump(label_encoder, f, indent=2)
    print(f"\nSaved label encoder to: {encoder_path}")

    # Return statistics
    return {
        "category": category,
        "train_ratio": train_ratio,
        "total_samples": len(files),
        "train_samples": len(train_files),
        "test_samples": len(test_files),
        "num_classes": len(unique_labels),
        "classes": unique_labels
    }


def split_all_categories(
    normalized_dir: Path = NORMALIZED_DIR,
    train_dir: Path = TRAIN_DIR,
    test_dir: Path = TEST_DIR,
    categories: List[str] = None,
    train_ratio: float = TRAIN_RATIO
):
    """
    Split semua kategori dataset.

    Args:
        normalized_dir: Directory normalized data
        train_dir: Directory untuk data train
        test_dir: Directory untuk data test
        categories: List kategori
        train_ratio: Train ratio
    """
    if categories is None:
        categories = CATEGORIES

    print("="*60)
    print("TAHAP 3: SPLIT DATASET")
    print("="*60)
    print(f"Input: {normalized_dir}")
    print(f"Output:")
    print(f"  - Train: {train_dir}")
    print(f"  - Test: {test_dir}")
    print(f"Categories: {categories}")
    print(f"Train ratio: {train_ratio*100}%")
    print(f"Test ratio: {(1-train_ratio)*100}%")
    print("="*60)

    all_split_info = {}

    for category in categories:
        split_info = split_category(
            normalized_dir=normalized_dir,
            train_dir=train_dir,
            test_dir=test_dir,
            category=category,
            train_ratio=train_ratio
        )
        if split_info:
            all_split_info[category] = split_info

    # Print summary table
    print("\n" + "="*60)
    print("SPLIT SUMMARY")
    print("="*60)
    print(f"{'Kategori':<20} {'Total':<10} {'Train':<10} {'Test':<10} {'Kelas':<10}")
    print("-"*60)
    for cat, info in all_split_info.items():
        print(f"{cat:<20} {info['total_samples']:<10} {info['train_samples']:<10} "
              f"{info['test_samples']:<10} {info['num_classes']:<10}")
    print("-"*60)


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Split dataset untuk training dan testing"
    )
    parser.add_argument(
        "--input",
        type=str,
        default=None,
        help="Directory normalized data (default: dataset/normalized)"
    )
    parser.add_argument(
        "--train-dir",
        type=str,
        default=None,
        help="Directory untuk data train (default: dataset/train)"
    )
    parser.add_argument(
        "--test-dir",
        type=str,
        default=None,
        help="Directory untuk data test (default: dataset/test)"
    )
    parser.add_argument(
        "--categories",
        nargs="+",
        choices=["abjad_bisindo", "abjad_sibi", "angka", "kata"],
        default=None,
        help="Kategori yang akan diproses"
    )
    parser.add_argument(
        "--train-ratio",
        type=float,
        default=TRAIN_RATIO,
        help=f"Rasio training (default: {TRAIN_RATIO})"
    )

    args = parser.parse_args()

    # Setup paths
    normalized_dir = Path(args.input) if args.input else NORMALIZED_DIR
    train_dir = Path(args.train_dir) if args.train_dir else TRAIN_DIR
    test_dir = Path(args.test_dir) if args.test_dir else TEST_DIR

    # Split
    split_all_categories(
        normalized_dir=normalized_dir,
        train_dir=train_dir,
        test_dir=test_dir,
        categories=args.categories,
        train_ratio=args.train_ratio
    )


if __name__ == "__main__":
    main()