"""
Model Training Script for Music Genre Classification

Project: Music Classification Dataset
Author: Molla Samser
Company: RSK World
Designer & Tester: Rima Khatun
Website: https://rskworld.in
Email: help@rskworld.in, support@rskworld.in
Phone: +91 93305 39277
Description: Train machine learning models for music genre classification
License: Educational Purpose Only
"""

import os
import sys
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.svm import SVC
from sklearn.neighbors import KNeighborsClassifier
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score
import joblib
import argparse
from tqdm import tqdm
import warnings
warnings.filterwarnings('ignore')

# Add parent directory to path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from utils.audio_processor import AudioProcessor
from utils.feature_extractor import FeatureExtractor


class MusicClassifier:
    """
    Music Genre Classifier
    
    Author: Molla Samser (RSK World)
    Website: https://rskworld.in
    Email: help@rskworld.in
    """
    
    def __init__(self, model_type: str = 'random_forest'):
        """
        Initialize classifier
        
        Args:
            model_type: Type of model ('random_forest', 'svm', 'knn')
        """
        self.model_type = model_type
        self.model = self._get_model(model_type)
        self.label_encoder = LabelEncoder()
        self.scaler = StandardScaler()
        self.audio_processor = AudioProcessor()
        self.feature_extractor = FeatureExtractor()
        
    def _get_model(self, model_type: str):
        """Get model based on type"""
        models = {
            'random_forest': RandomForestClassifier(n_estimators=100, random_state=42),
            'svm': SVC(kernel='rbf', random_state=42),
            'knn': KNeighborsClassifier(n_neighbors=5)
        }
        return models.get(model_type, models['random_forest'])
    
    def extract_features_from_csv(self, csv_path: str) -> tuple:
        """
        Extract features from audio files listed in CSV
        
        Args:
            csv_path: Path to CSV file with audio file paths
            
        Returns:
            Features array and labels array
        """
        df = pd.read_csv(csv_path)
        features_list = []
        labels_list = []
        
        print(f"Extracting features from {len(df)} audio files...")
        
        for idx, row in tqdm(df.iterrows(), total=len(df)):
            try:
                # Load audio
                audio_path = row['path']
                if not os.path.exists(audio_path):
                    print(f"File not found: {audio_path}")
                    continue
                
                audio = self.audio_processor.load_audio(audio_path)
                
                # Extract features
                features = self.feature_extractor.extract_all_features(audio)
                flat_features = self.feature_extractor.flatten_features(features)
                
                features_list.append(flat_features)
                labels_list.append(row['genre'])
                
            except Exception as e:
                print(f"Error processing {row['filename']}: {str(e)}")
                continue
        
        return np.array(features_list), np.array(labels_list)
    
    def train(self, X_train, y_train):
        """
        Train the model
        
        Args:
            X_train: Training features
            y_train: Training labels
        """
        print(f"\nTraining {self.model_type} model...")
        
        # Encode labels
        y_train_encoded = self.label_encoder.fit_transform(y_train)
        
        # Scale features
        X_train_scaled = self.scaler.fit_transform(X_train)
        
        # Train model
        self.model.fit(X_train_scaled, y_train_encoded)
        
        print("Training completed!")
    
    def evaluate(self, X_test, y_test):
        """
        Evaluate the model
        
        Args:
            X_test: Test features
            y_test: Test labels
        """
        # Encode labels
        y_test_encoded = self.label_encoder.transform(y_test)
        
        # Scale features
        X_test_scaled = self.scaler.transform(X_test)
        
        # Make predictions
        y_pred = self.model.predict(X_test_scaled)
        
        # Calculate metrics
        accuracy = accuracy_score(y_test_encoded, y_pred)
        
        print(f"\n{'='*50}")
        print(f"Model Evaluation Results")
        print(f"{'='*50}")
        print(f"Accuracy: {accuracy * 100:.2f}%\n")
        
        print("Classification Report:")
        print(classification_report(
            y_test_encoded, 
            y_pred, 
            target_names=self.label_encoder.classes_
        ))
        
        print("\nConfusion Matrix:")
        print(confusion_matrix(y_test_encoded, y_pred))
    
    def save_model(self, save_dir: str = 'saved_models'):
        """
        Save trained model
        
        Args:
            save_dir: Directory to save model
        """
        os.makedirs(save_dir, exist_ok=True)
        
        model_path = os.path.join(save_dir, f'{self.model_type}_model.pkl')
        encoder_path = os.path.join(save_dir, 'label_encoder.pkl')
        scaler_path = os.path.join(save_dir, 'scaler.pkl')
        
        joblib.dump(self.model, model_path)
        joblib.dump(self.label_encoder, encoder_path)
        joblib.dump(self.scaler, scaler_path)
        
        print(f"\nModel saved to {save_dir}/")


def main():
    """
    Main training function
    
    Author: Molla Samser
    Company: RSK World
    Website: https://rskworld.in
    """
    parser = argparse.ArgumentParser(description='Train Music Genre Classifier')
    parser.add_argument('--model', type=str, default='random_forest',
                        choices=['random_forest', 'svm', 'knn'],
                        help='Model type to train')
    parser.add_argument('--train_csv', type=str, default='../data/train_data.csv',
                        help='Path to training CSV file')
    parser.add_argument('--test_csv', type=str, default='../data/test_data.csv',
                        help='Path to test CSV file')
    
    args = parser.parse_args()
    
    print("="*70)
    print("Music Genre Classification - Model Training")
    print("Author: Molla Samser | Company: RSK World")
    print("Website: https://rskworld.in")
    print("="*70)
    
    # Initialize classifier
    classifier = MusicClassifier(model_type=args.model)
    
    # Extract features
    print("\n[1/4] Extracting training features...")
    X_train, y_train = classifier.extract_features_from_csv(args.train_csv)
    
    print("\n[2/4] Extracting test features...")
    X_test, y_test = classifier.extract_features_from_csv(args.test_csv)
    
    # Train model
    print("\n[3/4] Training model...")
    classifier.train(X_train, y_train)
    
    # Evaluate model
    print("\n[4/4] Evaluating model...")
    classifier.evaluate(X_test, y_test)
    
    # Save model
    classifier.save_model()
    
    print("\n" + "="*70)
    print("Training pipeline completed successfully!")
    print("© 2026 RSK World - Founded by Molla Samser")
    print("="*70)


if __name__ == "__main__":
    main()