"""
Neural Network Model 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: Deep learning models for advanced music genre classification
License: Educational Purpose Only
"""

import os
import sys
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score
import joblib
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

try:
    import tensorflow as tf
    from tensorflow import keras
    from tensorflow.keras import layers, models
    from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, ReduceLROnPlateau
    TENSORFLOW_AVAILABLE = True
except ImportError:
    TENSORFLOW_AVAILABLE = False
    print("TensorFlow not available. Install with: pip install tensorflow")


class DeepMusicClassifier:
    """
    Deep Learning Music Genre Classifier
    
    Author: Molla Samser (RSK World)
    Website: https://rskworld.in
    Email: help@rskworld.in
    """
    
    def __init__(self, input_shape=None, num_classes=8, model_type='cnn'):
        """
        Initialize Deep Learning Classifier
        
        Args:
            input_shape: Shape of input features
            num_classes: Number of genre classes
            model_type: 'cnn', 'lstm', 'dense', or 'hybrid'
        """
        if not TENSORFLOW_AVAILABLE:
            raise ImportError("TensorFlow is required for neural network models")
        
        self.input_shape = input_shape
        self.num_classes = num_classes
        self.model_type = model_type
        self.model = None
        self.label_encoder = LabelEncoder()
        self.scaler = StandardScaler()
        self.audio_processor = AudioProcessor()
        self.feature_extractor = FeatureExtractor()
        self.history = None
        
    def build_dense_model(self):
        """
        Build Dense Neural Network
        
        Author: Molla Samser (RSK World)
        """
        model = models.Sequential([
            layers.Input(shape=self.input_shape),
            layers.Dense(512, activation='relu'),
            layers.Dropout(0.3),
            layers.BatchNormalization(),
            
            layers.Dense(256, activation='relu'),
            layers.Dropout(0.3),
            layers.BatchNormalization(),
            
            layers.Dense(128, activation='relu'),
            layers.Dropout(0.2),
            layers.BatchNormalization(),
            
            layers.Dense(64, activation='relu'),
            layers.Dropout(0.2),
            
            layers.Dense(self.num_classes, activation='softmax')
        ])
        
        return model
    
    def build_cnn_model(self):
        """
        Build Convolutional Neural Network
        
        Author: Molla Samser (RSK World)
        """
        # Reshape input for CNN
        input_layer = layers.Input(shape=self.input_shape)
        
        # Expand dims for CNN
        x = layers.Reshape((self.input_shape[0], 1))(input_layer)
        
        # Conv blocks
        x = layers.Conv1D(64, 3, activation='relu', padding='same')(x)
        x = layers.BatchNormalization()(x)
        x = layers.MaxPooling1D(2)(x)
        x = layers.Dropout(0.2)(x)
        
        x = layers.Conv1D(128, 3, activation='relu', padding='same')(x)
        x = layers.BatchNormalization()(x)
        x = layers.MaxPooling1D(2)(x)
        x = layers.Dropout(0.2)(x)
        
        x = layers.Conv1D(256, 3, activation='relu', padding='same')(x)
        x = layers.BatchNormalization()(x)
        x = layers.GlobalAveragePooling1D()(x)
        
        # Dense layers
        x = layers.Dense(128, activation='relu')(x)
        x = layers.Dropout(0.3)(x)
        output = layers.Dense(self.num_classes, activation='softmax')(x)
        
        model = models.Model(inputs=input_layer, outputs=output)
        return model
    
    def build_lstm_model(self):
        """
        Build LSTM Neural Network
        
        Author: Molla Samser (RSK World)
        """
        input_layer = layers.Input(shape=self.input_shape)
        
        # Reshape for LSTM
        x = layers.Reshape((self.input_shape[0], 1))(input_layer)
        
        # LSTM layers
        x = layers.LSTM(128, return_sequences=True)(x)
        x = layers.Dropout(0.3)(x)
        
        x = layers.LSTM(64, return_sequences=True)(x)
        x = layers.Dropout(0.3)(x)
        
        x = layers.LSTM(32)(x)
        x = layers.Dropout(0.2)(x)
        
        # Dense layers
        x = layers.Dense(64, activation='relu')(x)
        x = layers.Dropout(0.2)(x)
        output = layers.Dense(self.num_classes, activation='softmax')(x)
        
        model = models.Model(inputs=input_layer, outputs=output)
        return model
    
    def build_hybrid_model(self):
        """
        Build Hybrid CNN-LSTM Model
        
        Author: Molla Samser (RSK World)
        Website: https://rskworld.in
        """
        input_layer = layers.Input(shape=self.input_shape)
        
        # Reshape for CNN
        x = layers.Reshape((self.input_shape[0], 1))(input_layer)
        
        # CNN layers
        x = layers.Conv1D(64, 3, activation='relu', padding='same')(x)
        x = layers.BatchNormalization()(x)
        x = layers.MaxPooling1D(2)(x)
        
        # LSTM layers
        x = layers.LSTM(64, return_sequences=True)(x)
        x = layers.Dropout(0.3)(x)
        x = layers.LSTM(32)(x)
        
        # Dense layers
        x = layers.Dense(64, activation='relu')(x)
        x = layers.Dropout(0.2)(x)
        output = layers.Dense(self.num_classes, activation='softmax')(x)
        
        model = models.Model(inputs=input_layer, outputs=output)
        return model
    
    def build_model(self):
        """Build model based on type"""
        print(f"\nBuilding {self.model_type.upper()} model...")
        
        model_builders = {
            'dense': self.build_dense_model,
            'cnn': self.build_cnn_model,
            'lstm': self.build_lstm_model,
            'hybrid': self.build_hybrid_model
        }
        
        self.model = model_builders.get(self.model_type, self.build_dense_model)()
        
        # Compile model
        self.model.compile(
            optimizer=keras.optimizers.Adam(learning_rate=0.001),
            loss='sparse_categorical_crossentropy',
            metrics=['accuracy']
        )
        
        print(f"Model built successfully!")
        return self.model
    
    def train(self, X_train, y_train, X_val=None, y_val=None, epochs=100, batch_size=32):
        """
        Train the neural network
        
        Author: Molla Samser
        Company: RSK World
        """
        print(f"\nTraining {self.model_type.upper()} model...")
        print(f"Training samples: {len(X_train)}")
        if X_val is not None:
            print(f"Validation samples: {len(X_val)}")
        
        # Encode labels
        y_train_encoded = self.label_encoder.fit_transform(y_train)
        if y_val is not None:
            y_val_encoded = self.label_encoder.transform(y_val)
        
        # Scale features
        X_train_scaled = self.scaler.fit_transform(X_train)
        if X_val is not None:
            X_val_scaled = self.scaler.transform(X_val)
        
        # Set input shape if not set
        if self.input_shape is None:
            self.input_shape = (X_train_scaled.shape[1],)
            self.build_model()
        
        # Callbacks
        callbacks = [
            EarlyStopping(
                monitor='val_loss' if X_val is not None else 'loss',
                patience=15,
                restore_best_weights=True,
                verbose=1
            ),
            ReduceLROnPlateau(
                monitor='val_loss' if X_val is not None else 'loss',
                factor=0.5,
                patience=7,
                min_lr=0.00001,
                verbose=1
            )
        ]
        
        # Train
        validation_data = (X_val_scaled, y_val_encoded) if X_val is not None else None
        
        self.history = self.model.fit(
            X_train_scaled, y_train_encoded,
            validation_data=validation_data,
            epochs=epochs,
            batch_size=batch_size,
            callbacks=callbacks,
            verbose=1
        )
        
        print("\nTraining completed!")
    
    def evaluate(self, X_test, y_test):
        """Evaluate the model"""
        y_test_encoded = self.label_encoder.transform(y_test)
        X_test_scaled = self.scaler.transform(X_test)
        
        # Predictions
        y_pred_proba = self.model.predict(X_test_scaled)
        y_pred = np.argmax(y_pred_proba, axis=1)
        
        # Metrics
        accuracy = accuracy_score(y_test_encoded, y_pred)
        
        print(f"\n{'='*70}")
        print(f"Neural Network Evaluation Results")
        print(f"{'='*70}")
        print(f"Model Type: {self.model_type.upper()}")
        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='saved_models'):
        """Save trained model"""
        os.makedirs(save_dir, exist_ok=True)
        
        model_path = os.path.join(save_dir, f'{self.model_type}_nn_model.h5')
        encoder_path = os.path.join(save_dir, f'{self.model_type}_label_encoder.pkl')
        scaler_path = os.path.join(save_dir, f'{self.model_type}_scaler.pkl')
        
        self.model.save(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
    """
    import argparse
    
    parser = argparse.ArgumentParser(description='Train Deep Learning Music Classifier')
    parser.add_argument('--model', type=str, default='dense',
                        choices=['dense', 'cnn', 'lstm', 'hybrid'],
                        help='Model architecture')
    parser.add_argument('--epochs', type=int, default=100,
                        help='Number of training epochs')
    parser.add_argument('--batch_size', type=int, default=32,
                        help='Batch size')
    
    args = parser.parse_args()
    
    print("="*70)
    print("Music Genre Classification - Deep Learning")
    print("Author: Molla Samser | Company: RSK World")
    print("Website: https://rskworld.in")
    print("="*70)
    
    # Load data (placeholder - you would load actual features here)
    print("\nNote: This is a template. Load your actual feature data.")
    print("© 2026 RSK World - Founded by Molla Samser")


if __name__ == "__main__":
    main()