"""
Variation Controller Engine

Controls variation strategies for different instrument roles in music production.
Provides intelligent sample rotation, pattern sequencing, and coherence validation.
"""

from typing import Dict, List, Optional, Any, Tuple
import random


class VariationController:
    """
    Controls variation strategies for different instrument roles.
    
    Supports multiple variation strategies:
    - round_robin: Cycle through samples in order
    - pattern: Follow predefined patterns with variations
    - layered: Use multiple samples simultaneously
    - section_based: Change samples based on song sections
    - random: Random selection with coherence constraints
    """
    
    # Default strategies for each instrument role
    DEFAULT_STRATEGIES: Dict[str, Dict[str, Any]] = {
        "kick": {
            "strategy": "round_robin",
            "params": {
                "rotation": "every_4_bars",  # Changes every 4 bars
                "shuffle": False,
                "offset": 0
            }
        },
        "snare": {
            "strategy": "pattern",
            "params": {
                "pattern": [0, 0, 1, 0],  # Dembow with variation on beat 3
                "fill_probability": 0.3,
                "fill_positions": [7, 15, 23, 31]  # Standard fill positions
            }
        },
        "hihat": {
            "strategy": "layered",
            "params": {
                "layer_count": 3,  # 2-3 simultaneous
                "velocity_variation": 0.2,
                "timing_jitter": 0.05
            }
        },
        "bass": {
            "strategy": "section_based",
            "params": {
                "section_map": {
                    "intro": 0,
                    "verse": 0,
                    "chorus": 1,
                    "bridge": 2,
                    "outro": 0
                }
            }
        },
        "perc": {
            "strategy": "random",
            "params": {
                "coherence_threshold": 0.7,
                "max_repetition": 2,
                "seed": None
            }
        },
        "fx": {
            "strategy": "section_based",
            "params": {
                "section_map": {
                    "intro": 0,
                    "build": 1,
                    "drop": 2,
                    "break": 3,
                    "outro": 0
                },
                "transition_only": True
            }
        },
        "chords": {
            "strategy": "section_based",
            "params": {
                "section_map": {
                    "intro": 0,
                    "verse": 0,
                    "chorus": 1,
                    "bridge": 2,
                    "outro": 0
                },
                "variation_intensity": 0.5
            }
        },
        "melody": {
            "strategy": "pattern",
            "params": {
                "pattern": [0, 0, 1, 0, 0, 1, 2, 0],  # Phrase-based variation
                "phrase_length": 4,
                "development": True
            }
        }
    }
    
    def __init__(self):
        """Initialize the variation controller with default strategies."""
        self._strategies: Dict[str, Dict[str, Any]] = {}
        self._variation_history: Dict[str, List[Tuple[int, int]]] = {}  # role -> [(bar, sample_index), ...]
        
        # Initialize with defaults
        for role, config in self.DEFAULT_STRATEGIES.items():
            self._strategies[role] = {
                "strategy": config["strategy"],
                "params": config["params"].copy()
            }
            self._variation_history[role] = []
    
    def configure_variation(self, role: str, strategy: str, params: Dict[str, Any]) -> None:
        """
        Configure variation strategy for a specific role.
        
        Args:
            role: Instrument role (kick, snare, hihat, bass, perc, fx, chords, melody)
            strategy: Strategy type (round_robin, pattern, layered, section_based, random)
            params: Strategy-specific parameters
        """
        if strategy not in ["round_robin", "pattern", "layered", "section_based", "random"]:
            raise ValueError(f"Unknown strategy: {strategy}")
        
        self._strategies[role] = {
            "strategy": strategy,
            "params": params.copy()
        }
        # Reset history for this role
        self._variation_history[role] = []
    
    def get_strategy(self, role: str) -> Dict[str, Any]:
        """
        Get current strategy configuration for a role.
        
        Args:
            role: Instrument role
            
        Returns:
            Dict with 'strategy' and 'params' keys
        """
        if role not in self._strategies:
            return {"strategy": "pattern", "params": {}}
        return self._strategies[role].copy()
    
    def generate_injection_plan(
        self, 
        role: str, 
        samples: List[str], 
        sections: List[Dict[str, Any]]
    ) -> Dict[str, Any]:
        """
        Generate a complete injection plan for a role across all sections.
        
        Args:
            role: Instrument role
            samples: List of sample file paths or identifiers
            sections: List of section dicts with 'name', 'start_bar', 'duration_bars'
            
        Returns:
            Injection plan dict with positions and sample assignments
        """
        if not samples:
            return {"role": role, "plan": [], "total_positions": 0}
        
        strategy_config = self._strategies.get(role, {"strategy": "pattern", "params": {}})
        strategy = strategy_config["strategy"]
        params = strategy_config["params"]
        
        # Calculate total song length
        if sections:
            song_length = max(
                s.get("start_bar", 0) + s.get("duration_bars", 8) 
                for s in sections
            )
        else:
            song_length = 32  # Default 32 bars
        
        plan = {
            "role": role,
            "strategy": strategy,
            "plan": [],
            "total_positions": 0,
            "samples_used": set()
        }
        
        # Generate based on strategy
        if strategy == "round_robin":
            plan["plan"] = self._generate_round_robin_plan(
                samples, song_length, params
            )
        elif strategy == "pattern":
            plan["plan"] = self._generate_pattern_plan(
                samples, song_length, sections, params
            )
        elif strategy == "layered":
            plan["plan"] = self._generate_layered_plan(
                samples, song_length, sections, params
            )
        elif strategy == "section_based":
            plan["plan"] = self._generate_section_based_plan(
                samples, sections, params
            )
        elif strategy == "random":
            plan["plan"] = self._generate_random_plan(
                samples, song_length, params
            )
        
        plan["total_positions"] = len(plan["plan"])
        plan["samples_used"] = list(plan["samples_used"])
        
        # Record in history
        for entry in plan["plan"]:
            self._variation_history[role].append(
                (entry.get("bar", 0), entry.get("sample_index", 0))
            )
        
        return plan
    
    def _generate_round_robin_plan(
        self, 
        samples: List[str], 
        song_length: int, 
        params: Dict[str, Any]
    ) -> List[Dict[str, Any]]:
        """Generate plan for round_robin strategy."""
        rotation = params.get("rotation", "every_4_bars")
        shuffle = params.get("shuffle", False)
        offset = params.get("offset", 0)
        
        positions = self.calculate_rotation_positions(rotation, song_length)
        plan = []
        
        sample_indices = list(range(len(samples)))
        if shuffle:
            random.shuffle(sample_indices)
        
        for i, bar in enumerate(positions):
            if bar >= song_length:
                break
            sample_idx = sample_indices[(i + offset) % len(samples)]
            plan.append({
                "bar": bar,
                "sample_index": sample_idx,
                "sample_path": samples[sample_idx],
                "variation_type": "round_robin"
            })
        
        return plan
    
    def _generate_pattern_plan(
        self, 
        samples: List[str], 
        song_length: int,
        sections: List[Dict[str, Any]],
        params: Dict[str, Any]
    ) -> List[Dict[str, Any]]:
        """Generate plan for pattern strategy."""
        pattern = params.get("pattern", [0, 0, 1, 0])
        fill_probability = params.get("fill_probability", 0.3)
        fill_positions = params.get("fill_positions", [7, 15, 23, 31])
        phrase_length = params.get("phrase_length", 4)
        
        extended_pattern = self.build_pattern_sequence(pattern, song_length)
        plan = []
        
        for bar in range(song_length):
            sample_idx = extended_pattern[bar] % len(samples)
            variation_type = "pattern"
            
            # Check for fill
            if bar in fill_positions and random.random() < fill_probability:
                if len(samples) > 1:
                    sample_idx = (sample_idx + 1) % len(samples)
                    variation_type = "fill"
            
            # Check for phrase boundary variation
            if bar % phrase_length == 0 and bar > 0:
                if len(samples) > 2:
                    sample_idx = (sample_idx + 2) % len(samples)
                    variation_type = "phrase"
            
            plan.append({
                "bar": bar,
                "sample_index": sample_idx,
                "sample_path": samples[sample_idx],
                "variation_type": variation_type
            })
        
        return plan
    
    def _generate_layered_plan(
        self, 
        samples: List[str], 
        song_length: int,
        sections: List[Dict[str, Any]],
        params: Dict[str, Any]
    ) -> List[Dict[str, Any]]:
        """Generate plan for layered strategy."""
        layer_count = min(params.get("layer_count", 3), len(samples))
        velocity_variation = params.get("velocity_variation", 0.2)
        timing_jitter = params.get("timing_jitter", 0.05)
        
        plan = []
        
        for bar in range(song_length):
            # Select layer_count samples
            selected_indices = random.sample(range(len(samples)), layer_count)
            
            layers = []
            for idx in selected_indices:
                velocity = random.uniform(0.7, 1.0) if random.random() < velocity_variation else 0.85
                timing_offset = random.uniform(-timing_jitter, timing_jitter) if random.random() < 0.3 else 0.0
                
                layers.append({
                    "sample_index": idx,
                    "sample_path": samples[idx],
                    "velocity": velocity,
                    "timing_offset": timing_offset
                })
            
            plan.append({
                "bar": bar,
                "layers": layers,
                "layer_count": layer_count,
                "variation_type": "layered"
            })
        
        return plan
    
    def _generate_section_based_plan(
        self, 
        samples: List[str], 
        sections: List[Dict[str, Any]],
        params: Dict[str, Any]
    ) -> List[Dict[str, Any]]:
        """Generate plan for section_based strategy."""
        section_map = params.get("section_map", {
            "intro": 0, "verse": 0, "chorus": 1, "bridge": 2, "outro": 0
        })
        transition_only = params.get("transition_only", False)
        variation_intensity = params.get("variation_intensity", 0.5)
        
        plan = []
        
        for section in sections:
            section_name = section.get("name", "verse").lower()
            start_bar = section.get("start_bar", 0)
            duration = section.get("duration_bars", 8)
            
            # Get sample index for this section
            mapped_idx = section_map.get(section_name, 0)
            sample_idx = mapped_idx % len(samples)
            
            # Add variation within section based on intensity
            for bar_offset in range(duration):
                bar = start_bar + bar_offset
                current_idx = sample_idx
                
                # Subtle variation within section
                if variation_intensity > 0.5 and bar_offset % 4 == 3:
                    if len(samples) > 1:
                        current_idx = (sample_idx + 1) % len(samples)
                
                plan.append({
                    "bar": bar,
                    "sample_index": current_idx,
                    "sample_path": samples[current_idx],
                    "section": section_name,
                    "variation_type": "section" if bar_offset == 0 else "section_sub"
                })
        
        return plan
    
    def _generate_random_plan(
        self, 
        samples: List[str], 
        song_length: int, 
        params: Dict[str, Any]
    ) -> List[Dict[str, Any]]:
        """Generate plan for random strategy."""
        coherence_threshold = params.get("coherence_threshold", 0.7)
        max_repetition = params.get("max_repetition", 2)
        seed = params.get("seed")
        
        if seed is not None:
            random.seed(seed)
        
        plan = []
        last_sample = 0
        repetition_count = 0
        
        for bar in range(song_length):
            # Random selection with coherence constraint
            if repetition_count >= max_repetition:
                # Force change
                available = [i for i in range(len(samples)) if i != last_sample]
                sample_idx = random.choice(available) if available else last_sample
                repetition_count = 0
            else:
                # Random with probability of coherence
                if random.random() < coherence_threshold:
                    sample_idx = last_sample
                    repetition_count += 1
                else:
                    available = [i for i in range(len(samples)) if i != last_sample]
                    sample_idx = random.choice(available) if available else last_sample
                    repetition_count = 0 if sample_idx != last_sample else repetition_count + 1
            
            last_sample = sample_idx
            
            plan.append({
                "bar": bar,
                "sample_index": sample_idx,
                "sample_path": samples[sample_idx],
                "variation_type": "random",
                "coherence": coherence_threshold
            })
        
        return plan
    
    def calculate_rotation_positions(self, rotation: str, song_length: int) -> List[int]:
        """
        Calculate bar positions for rotation-based strategies.
        
        Args:
            rotation: Rotation type (every_1_bar, every_2_bars, every_4_bars, every_8_bars)
            song_length: Total length in bars
            
        Returns:
            List of bar positions where rotation occurs
        """
        interval_map = {
            "every_1_bar": 1,
            "every_1_bars": 1,
            "every_2_bars": 2,
            "every_4_bars": 4,
            "every_8_bars": 8,
            "every_16_bars": 16
        }
        
        interval = interval_map.get(rotation, 4)
        return list(range(0, song_length, interval))
    
    def build_pattern_sequence(self, pattern: List[int], length: int) -> List[int]:
        """
        Extend a pattern to fill a given length.
        
        Args:
            pattern: Base pattern as list of sample indices
            length: Target length
            
        Returns:
            Extended pattern list
        """
        if not pattern:
            return [0] * length
        
        extended = []
        pattern_len = len(pattern)
        
        for i in range(length):
            extended.append(pattern[i % pattern_len])
        
        return extended
    
    def validate_variation_coherence(
        self, 
        plan: Dict[str, Any], 
        samples: List[Dict[str, Any]]
    ) -> float:
        """
        Validate the coherence of a variation plan.
        
        Args:
            plan: Injection plan dict
            samples: List of sample metadata dicts with features
            
        Returns:
            Coherence score between 0.0 and 1.0
        """
        if not plan.get("plan") or len(plan["plan"]) < 2:
            return 1.0  # Single sample is perfectly coherent
        
        plan_entries = plan["plan"]
        coherence_scores = []
        
        # Check adjacent sample compatibility
        for i in range(1, len(plan_entries)):
            prev_entry = plan_entries[i - 1]
            curr_entry = plan_entries[i]
            
            prev_idx = prev_entry.get("sample_index", 0)
            curr_idx = curr_entry.get("sample_index", 0)
            
            # Same sample = perfect coherence
            if prev_idx == curr_idx:
                coherence_scores.append(1.0)
                continue
            
            # Different samples - check compatibility
            if prev_idx < len(samples) and curr_idx < len(samples):
                prev_sample = samples[prev_idx]
                curr_sample = samples[curr_idx]
                
                # Compare spectral features if available
                prev_features = prev_sample.get("spectral_features", {})
                curr_features = curr_sample.get("spectral_features", {})
                
                if prev_features and curr_features:
                    # Calculate spectral centroid similarity
                    prev_centroid = prev_features.get("centroid_mean", 1000)
                    curr_centroid = curr_features.get("centroid_mean", 1000)
                    
                    # Normalize difference
                    max_diff = 5000  # 5kHz max reasonable difference
                    diff = abs(prev_centroid - curr_centroid)
                    similarity = max(0, 1 - (diff / max_diff))
                    coherence_scores.append(similarity)
                else:
                    # No features available, assume moderate coherence
                    coherence_scores.append(0.7)
            else:
                coherence_scores.append(0.5)
        
        # Calculate overall coherence
        if coherence_scores:
            avg_coherence = sum(coherence_scores) / len(coherence_scores)
            
            # Penalize excessive repetition (boring) or excessive change (chaotic)
            unique_transitions = len(set(
                (plan_entries[i-1].get("sample_index"), plan_entries[i].get("sample_index"))
                for i in range(1, len(plan_entries))
            ))
            total_transitions = len(plan_entries) - 1
            
            diversity_ratio = unique_transitions / max(1, total_transitions)
            
            # Ideal diversity is around 0.3-0.5
            if diversity_ratio < 0.2:
                # Too repetitive
                avg_coherence *= 0.8
            elif diversity_ratio > 0.7:
                # Too chaotic
                avg_coherence *= 0.9
            
            return round(avg_coherence, 2)
        
        return 1.0
    
    def get_variation_history(self, role: str) -> List[Tuple[int, int]]:
        """
        Get the variation history for a role.
        
        Args:
            role: Instrument role
            
        Returns:
            List of (bar, sample_index) tuples
        """
        return self._variation_history.get(role, []).copy()
    
    def clear_history(self, role: Optional[str] = None) -> None:
        """
        Clear variation history.
        
        Args:
            role: Specific role to clear, or None for all
        """
        if role is None:
            for r in self._variation_history:
                self._variation_history[r] = []
        else:
            self._variation_history[role] = []
    
    def reset_to_defaults(self) -> None:
        """Reset all strategies to defaults."""
        self.__init__()
