"""Music classification using Essentia-TensorFlow models."""
import os
import json
from pathlib import Path
from typing import Dict, List, Optional
import numpy as np

from ..utils.logging import get_logger
from ..utils.config import settings

logger = get_logger(__name__)

# Try to import essentia
try:
    from essentia.standard import (
        MonoLoader,
        TensorflowPredictEffnetDiscogs,
        TensorflowPredict2D,
        TensorflowPredictMusiCNN
    )
    ESSENTIA_AVAILABLE = True
except ImportError:
    logger.warning("Essentia-TensorFlow not available. Classification will be limited.")
    ESSENTIA_AVAILABLE = False


class EssentiaClassifier:
    """Classifier using Essentia pre-trained models."""

    # Model URLs (for documentation)
    MODEL_URLS = {
        "genre": "https://essentia.upf.edu/models/classification-heads/genre_discogs400/genre_discogs400-discogs-effnet-1.pb",
        "mood": "https://essentia.upf.edu/models/classification-heads/mtg_jamendo_moodtheme/mtg_jamendo_moodtheme-discogs-effnet-1.pb",
        "instrument": "https://essentia.upf.edu/models/classification-heads/mtg_jamendo_instrument/mtg_jamendo_instrument-discogs-effnet-1.pb",
    }

    def __init__(self, models_path: Optional[str] = None):
        """Initialize Essentia classifier.

        Args:
            models_path: Path to models directory (default: from settings)
        """
        self.models_path = Path(models_path or settings.ESSENTIA_MODELS_PATH)
        self.models = {}
        self.class_labels = {}

        if not ESSENTIA_AVAILABLE:
            logger.warning("Essentia not available - using fallback classifications")
            return

        # Load models if available
        self._load_models()

    def _load_models(self) -> None:
        """Load Essentia TensorFlow models."""
        if not self.models_path.exists():
            logger.warning(f"Models path {self.models_path} does not exist")
            return

        # Check for embedding model first
        embedding_file = "discogs-effnet-bs64-1.pb"
        embedding_path = self.models_path / embedding_file
        if embedding_path.exists():
            logger.info(f"Loading embedding model from {embedding_path}")
            self.models["embedding"] = str(embedding_path)
        else:
            logger.warning(f"Embedding model not found: {embedding_path}")
            return  # Cannot proceed without embeddings

        # Model file names for classification heads
        model_files = {
            "genre": "genre_discogs400-discogs-effnet-1.pb",
            "mood": "mtg_jamendo_moodtheme-discogs-effnet-1.pb",
            "instrument": "mtg_jamendo_instrument-discogs-effnet-1.pb",
        }

        for model_name, model_file in model_files.items():
            model_path = self.models_path / model_file
            if model_path.exists():
                try:
                    logger.info(f"Loading {model_name} model from {model_path}")
                    # Models will be loaded on demand
                    self.models[model_name] = str(model_path)
                except Exception as e:
                    logger.error(f"Failed to load {model_name} model: {e}")
            else:
                logger.warning(f"Model file not found: {model_path}")

        # Load class labels
        self._load_class_labels()

    def _load_class_labels(self) -> None:
        """Load class labels for models."""
        # Load genre labels from Discogs400 JSON file
        genre_json_path = self.models_path / "genre_discogs400-discogs-effnet-1.json"
        if genre_json_path.exists():
            try:
                with open(genre_json_path, 'r', encoding='utf-8') as f:
                    genre_metadata = json.load(f)
                    self.class_labels["genre"] = genre_metadata.get("classes", [])
                    logger.info(f"Loaded {len(self.class_labels['genre'])} genre labels from JSON")
            except Exception as e:
                logger.error(f"Failed to load genre labels from JSON: {e}")
                self.class_labels["genre"] = []
        else:
            logger.warning(f"Genre labels JSON not found: {genre_json_path}")
            self.class_labels["genre"] = []

        self.class_labels["mood"] = [
            "action", "adventure", "advertising", "background", "ballad",
            "calm", "children", "christmas", "commercial", "cool",
            "corporate", "dark", "deep", "documentary", "drama",
            "dramatic", "dream", "emotional", "energetic", "epic",
            "fast", "film", "fun", "funny", "game",
            "groovy", "happy", "heavy", "holiday", "hopeful",
            "inspiring", "love", "meditative", "melancholic", "mellow",
            "melodic", "motivational", "movie", "nature", "party",
            "positive", "powerful", "relaxing", "retro", "romantic",
            "sad", "sexy", "slow", "soft", "soundscape",
            "space", "sport", "summer", "trailer", "travel",
            "upbeat", "uplifting"
        ]

        self.class_labels["instrument"] = [
            "accordion", "acousticbassguitar", "acousticguitar", "bass",
            "beat", "bell", "bongo", "brass", "cello",
            "clarinet", "classicalguitar", "computer", "doublebass", "drummachine",
            "drums", "electricguitar", "electricpiano", "flute", "guitar",
            "harmonica", "harp", "horn", "keyboard", "oboe",
            "orchestra", "organ", "pad", "percussion", "piano",
            "pipeorgan", "rhodes", "sampler", "saxophone", "strings",
            "synthesizer", "trombone", "trumpet", "viola", "violin",
            "voice"
        ]

    def predict_genre(self, audio_path: str) -> Dict:
        """Predict music genre.

        Args:
            audio_path: Path to audio file

        Returns:
            Dictionary with genre predictions
        """
        if not ESSENTIA_AVAILABLE or "genre" not in self.models:
            return self._fallback_genre()

        try:
            # Step 1: Extract embeddings using discogs-effnet
            audio = MonoLoader(filename=audio_path, sampleRate=16000, resampleQuality=4)()

            embedding_model = TensorflowPredictEffnetDiscogs(
                graphFilename=self.models["embedding"],
                output="PartitionedCall:1"
            )
            embeddings = embedding_model(audio)

            # Average embeddings over time
            embeddings_mean = np.mean(embeddings, axis=0)

            # Step 2: Feed embeddings to classification head
            # Discogs400 uses different node names than MTG-Jamendo
            classifier = TensorflowPredict2D(
                graphFilename=self.models["genre"],
                input="serving_default_model_Placeholder",
                output="PartitionedCall:0"
            )
            predictions = classifier(embeddings_mean.reshape(1, -1))
            predictions = predictions[0]  # Remove batch dimension

            # Get top predictions
            labels = self.class_labels.get("genre", [])
            logger.info(f"Genre predictions shape: {predictions.shape}, num_labels: {len(labels)}")

            # Ensure we don't go out of bounds
            if len(predictions) == 0:
                logger.warning("No predictions returned from genre model")
                return self._fallback_genre()

            top_indices = np.argsort(predictions)[::-1][:5]
            # Only use indices that are within the labels range
            valid_top_indices = [i for i in top_indices if i < len(labels)]

            if not valid_top_indices:
                logger.warning(f"No valid indices found. Predictions: {len(predictions)}, Labels: {len(labels)}")
                return self._fallback_genre()

            primary = labels[valid_top_indices[0]]
            secondary = [labels[i] for i in valid_top_indices[1:4]]
            confidence = float(predictions[valid_top_indices[0]])

            return {
                "primary": primary,
                "secondary": secondary,
                "confidence": confidence,
            }

        except Exception as e:
            logger.error(f"Genre prediction failed: {e}")
            return self._fallback_genre()

    def predict_mood(self, audio_path: str) -> Dict:
        """Predict mood/theme.

        Args:
            audio_path: Path to audio file

        Returns:
            Dictionary with mood predictions
        """
        if not ESSENTIA_AVAILABLE or "mood" not in self.models or "embedding" not in self.models:
            return self._fallback_mood()

        try:
            # Step 1: Extract embeddings using discogs-effnet
            audio = MonoLoader(filename=audio_path, sampleRate=16000, resampleQuality=4)()

            embedding_model = TensorflowPredictEffnetDiscogs(
                graphFilename=self.models["embedding"],
                output="PartitionedCall:1"
            )
            embeddings = embedding_model(audio)
            embeddings_mean = np.mean(embeddings, axis=0)

            # Step 2: Feed embeddings to classification head
            classifier = TensorflowPredict2D(
                graphFilename=self.models["mood"],
                input="model/Placeholder",
                output="model/Sigmoid"
            )
            predictions = classifier(embeddings_mean.reshape(1, -1))
            predictions = predictions[0]

            # Get top predictions
            labels = self.class_labels.get("mood", [])

            if len(predictions) == 0:
                return self._fallback_mood()

            top_indices = np.argsort(predictions)[::-1][:5]
            valid_top_indices = [i for i in top_indices if i < len(labels)]

            if not valid_top_indices:
                return self._fallback_mood()

            primary = labels[valid_top_indices[0]] if valid_top_indices else "unknown"
            secondary = [labels[i] for i in valid_top_indices[1:3]] if len(valid_top_indices) > 1 else []

            # Estimate arousal and valence from mood labels (simplified)
            arousal, valence = self._estimate_arousal_valence(primary)

            return {
                "primary": primary,
                "secondary": secondary,
                "arousal": arousal,
                "valence": valence,
            }

        except Exception as e:
            logger.error(f"Mood prediction failed: {e}")
            return self._fallback_mood()

    def predict_instruments(self, audio_path: str) -> List[Dict]:
        """Predict instruments.

        Args:
            audio_path: Path to audio file

        Returns:
            List of instruments with confidence scores
        """
        if not ESSENTIA_AVAILABLE or "instrument" not in self.models or "embedding" not in self.models:
            return self._fallback_instruments()

        try:
            # Step 1: Extract embeddings using discogs-effnet
            audio = MonoLoader(filename=audio_path, sampleRate=16000, resampleQuality=4)()

            embedding_model = TensorflowPredictEffnetDiscogs(
                graphFilename=self.models["embedding"],
                output="PartitionedCall:1"
            )
            embeddings = embedding_model(audio)
            embeddings_mean = np.mean(embeddings, axis=0)

            # Step 2: Feed embeddings to classification head
            classifier = TensorflowPredict2D(
                graphFilename=self.models["instrument"],
                input="model/Placeholder",
                output="model/Sigmoid"
            )
            predictions = classifier(embeddings_mean.reshape(1, -1))
            predictions = predictions[0]

            # Get instruments above threshold
            threshold = 0.1
            labels = self.class_labels.get("instrument", [])
            instruments = []

            for i, score in enumerate(predictions):
                if score > threshold and i < len(labels):
                    instruments.append({
                        "name": labels[i],
                        "confidence": float(score)
                    })

            # Sort by confidence
            instruments.sort(key=lambda x: x["confidence"], reverse=True)

            return instruments[:10]  # Top 10

        except Exception as e:
            logger.error(f"Instrument prediction failed: {e}")
            return self._fallback_instruments()

    def _estimate_arousal_valence(self, mood: str) -> tuple:
        """Estimate arousal and valence from mood label.

        Args:
            mood: Mood label

        Returns:
            Tuple of (arousal, valence) scores (0-1)
        """
        # Simplified mapping (in production, use trained model)
        arousal_map = {
            "energetic": 0.9, "powerful": 0.9, "fast": 0.9, "action": 0.9,
            "calm": 0.2, "relaxing": 0.2, "meditative": 0.1, "slow": 0.3,
            "upbeat": 0.8, "party": 0.9, "groovy": 0.7,
        }

        valence_map = {
            "happy": 0.9, "positive": 0.9, "uplifting": 0.9, "fun": 0.9,
            "sad": 0.1, "dark": 0.2, "melancholic": 0.2, "dramatic": 0.3,
            "energetic": 0.7, "calm": 0.6, "romantic": 0.7,
        }

        arousal = arousal_map.get(mood.lower(), 0.5)
        valence = valence_map.get(mood.lower(), 0.5)

        return arousal, valence

    def _fallback_genre(self) -> Dict:
        """Fallback genre when model not available."""
        return {
            "primary": "unknown",
            "secondary": [],
            "confidence": 0.0,
        }

    def _fallback_mood(self) -> Dict:
        """Fallback mood when model not available."""
        return {
            "primary": "unknown",
            "secondary": [],
            "arousal": 0.5,
            "valence": 0.5,
        }

    def _fallback_instruments(self) -> List[Dict]:
        """Fallback instruments when model not available."""
        return []
