"""
Section Sample Mapper - Maps samples to song sections with coherence validation.

This engine maps 12 samples per category to 5 song sections (intro, verse, chorus, bridge, outro)
with specific energy levels and sample distributions for professional track composition.

Usage:
    from engines.section_sample_mapper import SectionSampleMapper
    
    mapper = SectionSampleMapper()
    section_map = mapper.map_samples_to_sections(samples, role="drums")
    chorus_samples = mapper.get_samples_for_section("chorus", samples)
    coherence = mapper.validate_section_coherence(section_map)
"""

import logging
from typing import Dict, List, Any, Optional
from dataclasses import dataclass, field

logger = logging.getLogger("SectionSampleMapper")


@dataclass
class SectionMapping:
    """Configuration for mapping samples to a specific section."""
    name: str
    indices: List[int]
    energy: float
    description: str = ""


class SectionSampleMapper:
    """
    Maps 12 samples per category to 5 song sections with coherence validation.
    
    Section Configuration:
    - intro: indices [0, 1, 2] (3 samples, low energy 0.3)
    - verse: indices [3, 4, 5, 6] (4 samples, medium energy 0.6)
    - chorus: indices [7, 8, 9, 10] (4 samples, high energy 0.9)
    - bridge: indices [10, 11] (2 samples, transition energy 0.5)
    - outro: indices [11, 0] (2 samples, fade out 0.2)
    
    Args:
        None - Uses built-in configuration
        
    Attributes:
        sections (Dict): Mapping of section names to SectionMapping configs
        total_samples (int): Total samples per category (12)
    """
    
    def __init__(self):
        """Initialize section configuration with 12-sample mapping."""
        self.total_samples: int = 12
        self.sections: Dict[str, SectionMapping] = {
            "intro": SectionMapping(
                name="intro",
                indices=[0, 1, 2],
                energy=0.3,
                description="Low energy introduction with 3 samples"
            ),
            "verse": SectionMapping(
                name="verse",
                indices=[3, 4, 5, 6],
                energy=0.6,
                description="Medium energy verse with 4 samples"
            ),
            "chorus": SectionMapping(
                name="chorus",
                indices=[7, 8, 9, 10],
                energy=0.9,
                description="High energy chorus with 4 samples"
            ),
            "bridge": SectionMapping(
                name="bridge",
                indices=[10, 11],
                energy=0.5,
                description="Transition energy bridge with 2 samples"
            ),
            "outro": SectionMapping(
                name="outro",
                indices=[11, 0],
                energy=0.2,
                description="Fade out outro with 2 samples"
            ),
        }
        logger.info("[SectionSampleMapper] Initialized with %d sections", len(self.sections))
    
    def map_samples_to_sections(
        self, 
        samples: List[str], 
        role: str
    ) -> Dict[str, List[str]]:
        """
        Maps 12 samples to 5 song sections based on configuration.
        
        Distributes samples according to energy curve:
        - intro: low energy samples [0-2]
        - verse: building samples [3-6]
        - chorus: peak energy samples [7-10]
        - bridge: transition samples [10-11]
        - outro: bookend samples [11, 0]
        
        Args:
            samples: List of 12 sample paths or identifiers
            role: Sample role/category (e.g., "drums", "bass", "synth")
            
        Returns:
            Dict mapping section names to lists of samples for that section
            
        Raises:
            ValueError: If samples list doesn't contain exactly 12 items
            
        Example:
            >>> mapper = SectionSampleMapper()
            >>> samples = [f"kick_{i}.wav" for i in range(12)]
            >>> section_map = mapper.map_samples_to_sections(samples, "kick")
            >>> section_map["intro"]
            ['kick_0.wav', 'kick_1.wav', 'kick_2.wav']
        """
        if len(samples) != self.total_samples:
            raise ValueError(
                f"Expected {self.total_samples} samples, got {len(samples)}"
            )
        
        section_map: Dict[str, List[str]] = {}
        
        for section_name, config in self.sections.items():
            # Map indices to samples (handle wrap-around for outro)
            section_samples: List[str] = []
            for idx in config.indices:
                if idx < len(samples):
                    section_samples.append(samples[idx])
            section_map[section_name] = section_samples
            
            logger.debug(
                "[SectionSampleMapper] Mapped %d samples to %s (energy %.1f)",
                len(section_samples),
                section_name,
                config.energy
            )
        
        logger.info(
            "[SectionSampleMapper] Mapped %d %s samples to %d sections",
            len(samples),
            role,
            len(section_map)
        )
        
        return section_map
    
    def get_samples_for_section(
        self, 
        section: str, 
        samples: List[str]
    ) -> List[str]:
        """
        Returns samples for a specific section.
        
        Args:
            section: Section name ("intro", "verse", "chorus", "bridge", "outro")
            samples: List of 12 sample paths or identifiers
            
        Returns:
            List of samples mapped to the specified section
            
        Raises:
            ValueError: If section name is invalid or samples count != 12
            
        Example:
            >>> samples = [f"snare_{i}.wav" for i in range(12)]
            >>> chorus_samples = mapper.get_samples_for_section("chorus", samples)
            >>> chorus_samples
            ['snare_7.wav', 'snare_8.wav', 'snare_9.wav', 'snare_10.wav']
        """
        if section not in self.sections:
            raise ValueError(
                f"Invalid section '{section}'. "
                f"Valid sections: {list(self.sections.keys())}"
            )
        
        if len(samples) != self.total_samples:
            raise ValueError(
                f"Expected {self.total_samples} samples, got {len(samples)}"
            )
        
        config = self.sections[section]
        section_samples: List[str] = []
        
        for idx in config.indices:
            if idx < len(samples):
                section_samples.append(samples[idx])
        
        logger.debug(
            "[SectionSampleMapper] Retrieved %d samples for %s section",
            len(section_samples),
            section
        )
        
        return section_samples
    
    def validate_section_coherence(
        self, 
        section_map: Dict[str, List[str]]
    ) -> float:
        """
        Returns coherence score (0.0-1.0) for the section mapping.
        
        Validates:
        - All required sections present
        - No duplicate samples across sections (except intentional overlaps)
        - Proper energy progression (intro < verse < chorus)
        - Sample count per section matches configuration
        
        Args:
            section_map: Dict mapping section names to sample lists
            
        Returns:
            Float coherence score between 0.0 and 1.0
            1.0 = perfect coherence
            0.0 = no coherence (validation failed)
            
        Example:
            >>> section_map = mapper.map_samples_to_sections(samples, "kick")
            >>> score = mapper.validate_section_coherence(section_map)
            >>> score
            0.95
        """
        if not section_map:
            logger.warning("[SectionSampleMapper] Empty section map, coherence = 0.0")
            return 0.0
        
        score: float = 1.0
        checks_passed: int = 0
        total_checks: int = 4
        
        # Check 1: All required sections present
        required_sections = set(self.sections.keys())
        present_sections = set(section_map.keys())
        if required_sections.issubset(present_sections):
            checks_passed += 1
        else:
            missing = required_sections - present_sections
            logger.warning("[SectionSampleMapper] Missing sections: %s", missing)
            score -= 0.25
        
        # Check 2: Sample counts match configuration
        count_valid = True
        for section_name, config in self.sections.items():
            if section_name in section_map:
                expected_count = len(config.indices)
                actual_count = len(section_map[section_name])
                if actual_count != expected_count:
                    count_valid = False
                    logger.warning(
                        "[SectionSampleMapper] %s: expected %d samples, got %d",
                        section_name,
                        expected_count,
                        actual_count
                    )
        if count_valid:
            checks_passed += 1
        else:
            score -= 0.25
        
        # Check 3: Energy progression (intro < verse < chorus)
        energy_valid = True
        energy_order = ["intro", "verse", "chorus"]
        energies = []
        for section_name in energy_order:
            if section_name in self.sections:
                energies.append(self.sections[section_name].energy)
        
        if len(energies) >= 2:
            for i in range(len(energies) - 1):
                if energies[i] >= energies[i + 1]:
                    energy_valid = False
                    break
        
        if energy_valid:
            checks_passed += 1
        else:
            score -= 0.25
            logger.warning("[SectionSampleMapper] Energy progression invalid")
        
        # Check 4: Minimal sample reuse (intentional overlaps allowed)
        all_samples: List[str] = []
        for section_samples in section_map.values():
            all_samples.extend(section_samples)
        
        unique_samples = set(all_samples)
        if len(unique_samples) > 0:
            reuse_ratio = len(all_samples) / len(unique_samples)
            # Some overlap is intentional (bridge shares with chorus/outro)
            if reuse_ratio <= 1.3:  # Allow up to 30% overlap
                checks_passed += 1
            else:
                score -= 0.15
                logger.warning(
                    "[SectionSampleMapper] High sample reuse ratio: %.2f",
                    reuse_ratio
                )
        
        # Calculate final score based on checks
        final_score = checks_passed / total_checks
        
        logger.info(
            "[SectionSampleMapper] Coherence validation: %.2f (%d/%d checks passed)",
            final_score,
            checks_passed,
            total_checks
        )
        
        return final_score
    
    def get_section_config(self) -> Dict[str, Any]:
        """
        Returns the section mapping configuration.
        
        Returns:
            Dict with section names as keys and configuration dicts as values.
            Each config contains: indices, energy, description
            
        Example:
            >>> config = mapper.get_section_config()
            >>> config["chorus"]
            {'indices': [7, 8, 9, 10], 'energy': 0.9, 'description': 'High energy chorus...'}
        """
        config: Dict[str, Any] = {}
        
        for section_name, mapping in self.sections.items():
            config[section_name] = {
                "indices": mapping.indices,
                "energy": mapping.energy,
                "description": mapping.description,
                "sample_count": len(mapping.indices)
            }
        
        return {
            "sections": config,
            "total_samples": self.total_samples,
            "section_count": len(self.sections)
        }
    
    def get_section_energy(self, section: str) -> float:
        """
        Get the energy level for a specific section.
        
        Args:
            section: Section name
            
        Returns:
            Energy level (0.0-1.0)
            
        Raises:
            ValueError: If section name is invalid
        """
        if section not in self.sections:
            raise ValueError(
                f"Invalid section '{section}'. "
                f"Valid sections: {list(self.sections.keys())}"
            )
        
        return self.sections[section].energy
    
    def get_all_section_energies(self) -> Dict[str, float]:
        """
        Get energy levels for all sections.
        
        Returns:
            Dict mapping section names to energy levels
        """
        return {
            name: config.energy 
            for name, config in self.sections.items()
        }


# Convenience function
def create_section_mapper() -> SectionSampleMapper:
    """
    Factory function to create a SectionSampleMapper instance.
    
    Returns:
        SectionSampleMapper instance with default configuration
    """
    return SectionSampleMapper()
