"""
Data inspection and validation utility
Helps verify dataset structure and statistics before training
"""

import os
import json
from pathlib import Path
from collections import defaultdict
import argparse


def inspect_dataset(dataset_path: Path):
    """Inspect dataset (supports data.json or legacy Kathbath format)"""
    stats = {
        'audio_files': 0,
        'transcriptions': 0,
        'matched': 0,
        'unmatched_audio': 0,
        'unmatched_text': 0,
        'avg_text_length': 0,
        'min_text_length': float('inf'),
        'max_text_length': 0,
    }

    # Try loading from data.json first
    data_file = dataset_path / "data.json"
    if data_file.exists():
        try:
            with open(data_file, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            stats['transcriptions'] = len(data)
            text_lengths = []
            matched_count = 0
            
            # Check each sample
            for item in data:
                audio_path = dataset_path / item.get('audioFilename', '')
                text = item.get('text', '')
                
                if audio_path.exists():
                    matched_count += 1
                
                if text:
                    text_lengths.append(len(text))
            
            stats['matched'] = matched_count
            stats['audio_files'] = matched_count # Approximation since we don't scan all files
            stats['unmatched_audio'] = 0 # Not easily calculated in this mode
            stats['unmatched_text'] = len(data) - matched_count
            
            if text_lengths:
                stats['avg_text_length'] = sum(text_lengths) / len(text_lengths)
                stats['min_text_length'] = min(text_lengths)
                stats['max_text_length'] = max(text_lengths)
                
            return stats
            
        except Exception as e:
            print(f"  ❌ Error reading data.json in {dataset_path.name}: {e}")
            # Fallback to legacy check if data.json fails
    
    # Legacy Kathbath format check
    audio_dir = dataset_path / "audio"
    if not audio_dir.exists():
        # Check for audios/ as fallback
        audio_dir = dataset_path / "audios"
        
    transcript_file = dataset_path / "transcription.txt"
    
    if not audio_dir.exists():
        print(f"  ❌ Missing audio directory: {audio_dir}")
        return None
    
    if not transcript_file.exists():
        # If no data.json and no transcription.txt, it's problematic
        if not data_file.exists():
            print(f"  ❌ Missing transcription file: {transcript_file}")
            return None
        return None # Should have been handled by data.json block
    
    # Count audio files
    audio_files = list(audio_dir.glob("*.wav"))
    
    # Count transcriptions
    with open(transcript_file, 'r', encoding='utf-8') as f:
        transcriptions = [line.strip() for line in f if line.strip()]
    
    stats['audio_files'] = len(audio_files)
    stats['transcriptions'] = len(transcriptions)
    
    # Parse transcriptions
    trans_dict = {}
    text_lengths = []
    
    for line in transcriptions:
        parts = line.split('\t', 1)
        if len(parts) == 2:
            audio_id, text = parts
            trans_dict[audio_id] = text
            text_lengths.append(len(text))
    
    # Match with audio files
    audio_ids = {f.stem for f in audio_files}
    trans_ids = set(trans_dict.keys())
    
    stats['matched'] = len(audio_ids & trans_ids)
    stats['unmatched_audio'] = len(audio_ids - trans_ids)
    stats['unmatched_text'] = len(trans_ids - audio_ids)
    
    if text_lengths:
        stats['avg_text_length'] = sum(text_lengths) / len(text_lengths)
        stats['min_text_length'] = min(text_lengths)
        stats['max_text_length'] = max(text_lengths)
    
    return stats


def inspect_directory(base_path: Path, dir_type: str):
    """Inspect train or test directory"""
    print(f"\n{'='*80}")
    print(f"{dir_type.upper()} DATA INSPECTION")
    print(f"{'='*80}\n")
    
    if not base_path.exists():
        print(f"❌ Directory not found: {base_path}")
        return {}
    
    datasets = {}
    
    # Find all dataset directories
    for item in base_path.iterdir():
        if item.is_dir() and not item.name.startswith('.'):
            print(f"📁 {item.name}")
            stats = inspect_dataset(item)
            
            if stats:
                datasets[item.name] = stats
                
                # Print statistics
                print(f"  ✅ Audio files: {stats['audio_files']}")
                print(f"  ✅ Transcriptions: {stats['transcriptions']}")
                print(f"  ✅ Matched samples: {stats['matched']}")
                
                if stats['unmatched_audio'] > 0:
                    print(f"  ⚠️  Unmatched audio files: {stats['unmatched_audio']}")
                
                if stats['unmatched_text'] > 0:
                    print(f"  ⚠️  Unmatched transcriptions: {stats['unmatched_text']}")
                
                if stats['avg_text_length'] > 0:
                    print(f"  📊 Avg text length: {stats['avg_text_length']:.1f} chars")
                    print(f"  📊 Text length range: {stats['min_text_length']}-{stats['max_text_length']} chars")
                
                print()
    
    return datasets


def calculate_total_stats(train_datasets, test_datasets):
    """Calculate overall statistics"""
    print(f"\n{'='*80}")
    print("OVERALL STATISTICS")
    print(f"{'='*80}\n")
    
    # Training stats
    total_train_samples = sum(d['matched'] for d in train_datasets.values())
    total_train_datasets = len(train_datasets)
    
    print(f"Training:")
    print(f"  Total datasets: {total_train_datasets}")
    print(f"  Total samples: {total_train_samples}")
    
    # Language breakdown
    lang_counts = defaultdict(int)
    for name, stats in train_datasets.items():
        # Extract language from dataset name
        name_lower = name.lower()
        if 'hindi' in name_lower:
            lang_counts['Hindi'] += stats['matched']
        elif 'bengali' in name_lower or 'bengali' in name_lower:
            lang_counts['Bengali'] += stats['matched']
        elif 'marathi' in name_lower:
            lang_counts['Marathi'] += stats['matched']
        elif 'odia' in name_lower:
            lang_counts['Odia'] += stats['matched']
    
    print(f"\n  Language breakdown:")
    for lang, count in sorted(lang_counts.items()):
        percentage = (count / total_train_samples * 100) if total_train_samples > 0 else 0
        print(f"    {lang:15s}: {count:5d} samples ({percentage:.1f}%)")
    
    # Test stats
    if test_datasets:
        total_test_samples = sum(d['matched'] for d in test_datasets.values())
        total_test_datasets = len(test_datasets)
        
        print(f"\nTest:")
        print(f"  Total datasets: {total_test_datasets}")
        print(f"  Total samples: {total_test_samples}")
        
        # Test conditions
        print(f"\n  Test conditions:")
        for name, stats in sorted(test_datasets.items()):
            print(f"    {name:40s}: {stats['matched']:5d} samples")
    
    print(f"\n{'='*80}\n")


def main():
    parser = argparse.ArgumentParser(description="Inspect training and test data")
    parser.add_argument(
        "--base_path",
        type=str,
        default=".",
        help="Base path containing train/ and test/ directories"
    )
    parser.add_argument(
        "--output",
        type=str,
        help="Save statistics to JSON file"
    )
    
    args = parser.parse_args()
    
    base_path = Path(args.base_path)
    
    # Inspect training data
    train_path = base_path / "train"
    train_datasets = inspect_directory(train_path, "train")
    
    # Inspect test data
    test_path = base_path / "test"
    test_datasets = inspect_directory(test_path, "test")
    
    # Calculate overall statistics
    calculate_total_stats(train_datasets, test_datasets)
    
    # Check for common issues
    print("⚠️  WARNINGS:")
    issues = []
    
    for name, stats in {**train_datasets, **test_datasets}.items():
        if stats['unmatched_audio'] > 0:
            issues.append(f"  • {name}: {stats['unmatched_audio']} audio files without transcriptions")
        
        if stats['unmatched_text'] > 0:
            issues.append(f"  • {name}: {stats['unmatched_text']} transcriptions without audio files")
        
        match_rate = stats['matched'] / max(stats['audio_files'], stats['transcriptions']) if max(stats['audio_files'], stats['transcriptions']) > 0 else 0
        if match_rate < 0.95:
            issues.append(f"  • {name}: Low match rate ({match_rate*100:.1f}%)")
    
    if issues:
        print("\n" + "\n".join(issues))
    else:
        print("\n  ✅ No issues found!")
    
    # Recommendations
    print(f"\n{'='*80}")
    print("RECOMMENDATIONS")
    print(f"{'='*80}\n")
    
    if lang_counts := defaultdict(int):
        for name, stats in train_datasets.items():
            name_lower = name.lower()
            if 'marathi' in name_lower:
                lang_counts['marathi'] += stats['matched']
            elif 'odia' in name_lower:
                lang_counts['odia'] += stats['matched']
            elif 'bengali' in name_lower:
                lang_counts['bengali'] += stats['matched']
            elif 'hindi' in name_lower:
                lang_counts['hindi'] += stats['matched']
        
        # Find imbalanced languages
        if lang_counts:
            max_samples = max(lang_counts.values())
            
            for lang, count in lang_counts.items():
                ratio = max_samples / count if count > 0 else 0
                if ratio > 2:
                    print(f"  • Consider increasing augmentation factor for {lang.capitalize()}")
                    print(f"    Current samples: {count}, Suggested factor: {int(ratio)}")
    
    print(f"\n  • Review config.py augmentation_factors based on data distribution")
    print(f"  • Check that all audio files are valid WAV format (16kHz recommended)")
    print(f"  • Ensure transcriptions use correct Unicode encoding (UTF-8)")
    
    # Save to file if requested
    if args.output:
        output_data = {
            'train': train_datasets,
            'test': test_datasets,
            'summary': {
                'total_train_samples': sum(d['matched'] for d in train_datasets.values()),
                'total_test_samples': sum(d['matched'] for d in test_datasets.values()),
                'train_datasets': len(train_datasets),
                'test_datasets': len(test_datasets),
            }
        }
        
        with open(args.output, 'w', encoding='utf-8') as f:
            json.dump(output_data, f, indent=2, ensure_ascii=False)
        
        print(f"\n✅ Statistics saved to: {args.output}")


if __name__ == "__main__":
    main()