"""
Tahap 2: Normalize Landmark

Script ini menormalisasi landmark yang sudah diekstrak
dengan translation dan scaling untuk konsistensi.

Sesuai README:
- Input: dataset/landmark/
- Output: dataset/normalized/
- Proses:
  - Membaca seluruh landmark
  - Menggeser koordinat ke titik referensi
  - Scaling
  - Menyimpan landmark hasil normalisasi

Catatan: Struktur folder pada normalized/ akan dibuat identik dengan folder landmark/ dan raw/.
"""

import sys
import numpy as np
from pathlib import Path
from tqdm import tqdm
from typing import List, Tuple, Optional
import argparse

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

from utils.normalization import LandmarkNormalizer
from utils.config import (
    LANDMARK_DIR,
    NORMALIZED_DIR,
    CATEGORIES,
    WINDOW_SIZE,
    FEATURE_SIZE
)


class LandmarkNormalizerPipeline:
    """
    Pipeline untuk normalisasi landmark sequences.
    """

    def __init__(self, scale_factor: Optional[float] = None):
        """
        Inisialisasi normalizer pipeline.

        Args:
            scale_factor: Optional fixed scale factor
        """
        self.normalizer = LandmarkNormalizer(scale_factor=scale_factor)

    def normalize_sequence(self, landmarks: np.ndarray) -> np.ndarray:
        """
        Normalisasi satu sequence landmark.

        Args:
            landmarks: Array shape (30, 306)

        Returns:
            Normalized sequence
        """
        if landmarks is None or len(landmarks) == 0:
            return np.zeros((WINDOW_SIZE, FEATURE_SIZE))

        normalized_sequence = []

        for frame_features in landmarks:
            # Split frame ke components
            left_hand = frame_features[:63].reshape(21, 3)
            right_hand = frame_features[63:126].reshape(21, 3)
            face = frame_features[126:].reshape(60, 3)

            # Normalize
            norm_frame = self.normalizer.normalize_frame(left_hand, right_hand, face)
            normalized_sequence.append(norm_frame)

        return np.array(normalized_sequence)


def get_landmark_files(directory: Path) -> List[Path]:
    """
    Get semua landmark .npy files dalam directory.

    Args:
        directory: Root directory

    Returns:
        List of .npy file paths
    """
    if not directory.exists():
        return []

    return sorted(directory.rglob("*.npy"))


def normalize_category(
    landmark_dir: Path,
    normalized_dir: Path,
    category: str
) -> Tuple[int, int]:
    """
    Normalisasi landmarks untuk satu kategori.

    Args:
        landmark_dir: Directory landmark mentah
        normalized_dir: Directory output normalized
        category: Nama kategori

    Returns:
        Tuple (success_count, fail_count)
    """
    category_landmark_dir = landmark_dir / category

    if not category_landmark_dir.exists():
        print(f"Warning: Category directory tidak ditemukan: {category_landmark_dir}")
        return 0, 0

    # Get all landmark files
    landmark_files = get_landmark_files(category_landmark_dir)

    if not landmark_files:
        print(f"Warning: Tidak ada landmark file ditemukan untuk kategori {category}")
        return 0, 0

    print(f"\n{'='*60}")
    print(f"Normalizing landmarks untuk kategori: {category}")
    print(f"Total files: {len(landmark_files)}")
    print(f"{'='*60}\n")

    # Initialize normalizer
    normalizer = LandmarkNormalizerPipeline()

    success_count = 0
    fail_count = 0

    for landmark_path in tqdm(landmark_files, desc=f"Normalizing {category}"):
        try:
            # Load landmark
            sequence = np.load(landmark_path)

            # Get label dari path
            # Format: landmark/{category}/{label}/file.npy
            relative_path = landmark_path.relative_to(category_landmark_dir)
            parts = relative_path.parts

            if len(parts) >= 2:
                label = parts[0]  # Folder label

                # Create output path (struktur identik)
                output_dir = normalized_dir / category / label
                output_dir.mkdir(parents=True, exist_ok=True)

                output_path = output_dir / landmark_path.name

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

                # Normalize
                normalized_sequence = normalizer.normalize_sequence(sequence)

                # Save
                np.save(output_path, normalized_sequence)
                success_count += 1
            else:
                fail_count += 1

        except Exception as e:
            print(f"Error processing {landmark_path}: {e}")
            fail_count += 1

    return success_count, fail_count


def normalize_all_categories(
    landmark_dir: Path = LANDMARK_DIR,
    normalized_dir: Path = NORMALIZED_DIR,
    categories: List[str] = None
):
    """
    Normalisasi landmarks untuk semua kategori.

    Args:
        landmark_dir: Directory landmark mentah
        normalized_dir: Directory output normalized
        categories: List kategori yang akan diproses
    """
    if categories is None:
        categories = CATEGORIES

    print("="*60)
    print("TAHAP 2: NORMALIZE LANDMARK")
    print("="*60)
    print(f"Input: {landmark_dir}")
    print(f"Output: {normalized_dir}")
    print(f"Categories: {categories}")
    print("="*60)

    total_success = 0
    total_fail = 0

    for category in categories:
        success, fail = normalize_category(
            landmark_dir=landmark_dir,
            normalized_dir=normalized_dir,
            category=category
        )
        total_success += success
        total_fail += fail

    print("\n" + "="*60)
    print("NORMALIZATION COMPLETE")
    print("="*60)
    print(f"Total berhasil: {total_success}")
    print(f"Total gagal: {total_fail}")
    if total_success + total_fail > 0:
        print(f"Success rate: {total_success/(total_success+total_fail)*100:.2f}%")
    print("="*60)


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Normalisasi landmark untuk gesture recognition"
    )
    parser.add_argument(
        "--input",
        type=str,
        default=None,
        help="Directory landmark (default: dataset/landmark)"
    )
    parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="Directory output normalized (default: dataset/normalized)"
    )
    parser.add_argument(
        "--categories",
        nargs="+",
        choices=["abjad_bisindo", "abjad_sibi", "angka", "kata"],
        default=None,
        help="Kategori yang akan diproses"
    )

    args = parser.parse_args()

    # Setup paths
    landmark_dir = Path(args.input) if args.input else LANDMARK_DIR
    normalized_dir = Path(args.output) if args.output else NORMALIZED_DIR

    # Normalize
    normalize_all_categories(
        landmark_dir=landmark_dir,
        normalized_dir=normalized_dir,
        categories=args.categories
    )


if __name__ == "__main__":
    main()