import librosa
import numpy as np

class ChordAnalyzer:
    def __init__(self):
        # Template Chord - Prioritize basic triads for accuracy
        self.templates = self._generate_chord_templates()

    def _generate_chord_templates(self):
        """
        Membuat template chroma untuk berbagai jenis chord.
        12 Nada: C, C#, D, D#, E, F, F#, G, G#, A, A#, B
        PRIORITIZED: Basic Major/Minor triads have higher matching priority
        """
        templates = {}
        roots = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']
        
        # Define chord qualities with PRIORITY WEIGHTS
        # Higher weight = more likely to be matched
        # Format: (intervals, priority_boost)
        qualities = {
            # PRIORITY 1: Basic triads (most common, highest priority)
            '': ([0, 4, 7], 1.3),           # Major - HIGHEST priority
            'm': ([0, 3, 7], 1.3),          # Minor - HIGHEST priority
            
            # PRIORITY 2: Power chord & Suspended
            '5': ([0, 7], 1.1),             # Power chord
            'sus4': ([0, 5, 7], 1.0),
            'sus2': ([0, 2, 7], 1.0),
            
            # PRIORITY 3: 7th chords
            'maj7': ([0, 4, 7, 11], 0.95),
            'm7': ([0, 3, 7, 10], 0.95),
            '7': ([0, 4, 7, 10], 0.95),     # Dominant 7
            
            # PRIORITY 4: Extended & Other (lower priority to avoid false matches)
            'dim': ([0, 3, 6], 0.9),
            'aug': ([0, 4, 8], 0.9),
            '6': ([0, 4, 7, 9], 0.85),
            'm6': ([0, 3, 7, 9], 0.85),
            'add9': ([0, 4, 7, 2], 0.85),
            'madd9': ([0, 3, 7, 2], 0.85),
        }

        for i, root in enumerate(roots):
            for quality, (intervals, priority) in qualities.items():
                # Build chroma vector with weighted notes
                vec = np.zeros(12)
                for j, interval in enumerate(intervals):
                    idx = (i + interval) % 12
                    # Root = 2.0, Fifth = 1.5, Third = 1.2, Others = 1.0
                    if j == 0:  # Root
                        weight = 2.0
                    elif interval == 7:  # Fifth
                        weight = 1.5
                    elif interval in [3, 4]:  # Third (major or minor)
                        weight = 1.2
                    else:
                        weight = 1.0
                    vec[idx] = weight
                
                # Apply priority boost
                vec *= priority
                
                chord_name = f"{root}{quality}"
                # Normalize vector
                norm = np.linalg.norm(vec)
                if norm > 0:
                    vec /= norm
                templates[chord_name] = vec
                
        return templates

    def analyze(self, audio_path: str, sr=22050):
        """
        Menganalisis file audio dan mengembalikan progresi chord dengan timestamp.
        """
        print(f"Analyzing chords for: {audio_path}")
        try:
            y, sr = librosa.load(audio_path, sr=sr)
            
            # Harmonic-Percussive Source Separation
            y_harmonic, _ = librosa.effects.hpss(y)
            
            # Compute Chroma CQT
            chroma = librosa.feature.chroma_cqt(y=y_harmonic, sr=sr, bins_per_octave=24)
            
            # 1. TEMPORAL SMOOTHING (Median Filter)
            # Increased filter size for more stability (21 frames ~= 0.5s)
            import scipy.ndimage
            chroma = scipy.ndimage.median_filter(chroma, size=(1, 21))
            chroma = librosa.util.normalize(chroma)
            
            num_frames = chroma.shape[1]
            
            # Template matching
            template_names = list(self.templates.keys())
            template_matrix = np.array([self.templates[name] for name in template_names])
            scores = np.dot(template_matrix, chroma)
            
            max_indices = np.argmax(scores, axis=0)
            max_scores = np.max(scores, axis=0)
            
            # 2. POST-PROCESSING (Merge Short Segments)
            current_chord = None
            start_time = 0.0
            
            THRESHOLD = 0.6  # Lower threshold for basic chord detection
            MIN_DURATION = 0.8  # Chord must last 0.8s to be valid
            
            raw_segments = []
            
            # First Pass: Collect segments
            for i in range(num_frames):
                idx = max_indices[i]
                score = max_scores[i]
                timestamp = librosa.frames_to_time(i, sr=sr)
                chord_name = template_names[idx] if score > THRESHOLD else "N.C."
                
                if chord_name != current_chord:
                    if current_chord is not None:
                         raw_segments.append({
                            "chord": current_chord,
                            "start": start_time,
                            "end": timestamp,
                            "duration": timestamp - start_time
                        })
                    current_chord = chord_name
                    start_time = timestamp
            
            # Append last
            if current_chord is not None:
                 end_time = librosa.get_duration(y=y, sr=sr)
                 raw_segments.append({
                    "chord": current_chord,
                    "start": start_time,
                    "end": end_time,
                    "duration": end_time - start_time
                })

            # Second Pass: Merge short segments to neighbor
            final_results = []
            if not raw_segments: return []

            # Simple heuristic: If segment < MIN_DURATION, merge to previous if possible
            for seg in raw_segments:
                if not final_results:
                    final_results.append(seg)
                    continue

                prev = final_results[-1]
                
                # Jika segmen sekarang terlalu pendek, "makan" oleh segmen sebelumnya (atau abaikan)
                # TAPI jika chord-nya SAMA dengan sebelumnya, gabung saja.
                if seg["chord"] == prev["chord"]:
                    prev["end"] = seg["end"]
                    prev["duration"] += seg["duration"]
                elif seg["duration"] < MIN_DURATION:
                    # Merge to previous (extend previous to cover this short blip)
                    prev["end"] = seg["end"]
                    prev["duration"] += seg["duration"]
                else:
                    final_results.append(seg)
            
            # Format output logic (remove internal keys if needed)
            formatted_results = []
            for r in final_results:
                formatted_results.append({
                    "chord": r["chord"],
                    "start": round(r["start"], 2),
                    "end": round(r["end"], 2)
                })

            return formatted_results

        except Exception as e:
            print(f"Chord Analysis Error: {e}")
            return []