#!/usr/bin/env python3
"""
Visualization script for spectral features.

This script allows testing and visualizing the spectral feature calculations
before integrating them into the training pipeline. It shows:
- Original waveform
- Loudness sequence (default 25Hz)
- Spectral centroid sequence (default 25Hz)
- Spectral complexity sequence (default 25Hz)

Usage:
    # Full audio visualization
    python visualize_spectral_features.py --audio_path /path/to/audio.wav

    # Test cropping (mimics training pipeline)
    python visualize_spectral_features.py --audio_path /path/to/audio.wav --start_s 10 --end_s 30

    # Generate and display tags
    python visualize_spectral_features.py --audio_path /path/to/audio.wav --show_tags
"""

import argparse
import os
import sys

import matplotlib.pyplot as plt
import numpy as np
from scipy import signal

# Add parent directory to path for imports
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from spectral_features import (
    calculate_loudness_seq,
    calculate_spectral_centroid_seq,
    calculate_spectral_complexity_seq,
)

# Import suno_utils if available, otherwise provide fallback
try:
    from suno_utils.audio import Audio
except ImportError:
    print("Warning: suno_utils not available. Using basic audio loading.")
    import soundfile as sf

    class Audio:
        @staticmethod
        def from_file(path):
            data, sr = sf.read(path)
            audio_obj = type("Audio", (), {})()
            audio_obj.array_float = data.T if data.ndim == 2 else data[np.newaxis, :]
            audio_obj.sample_rate = sr
            audio_obj.n_channels = audio_obj.array_float.shape[0]
            return audio_obj

        @staticmethod
        def sum(audios):
            # Simple implementation
            return audios[0] if audios else None


def load_audio(audio_path):
    """Load audio file and return audio data."""
    audio = Audio.from_file(audio_path)
    return audio


def calculate_features(
    audio_data, sample_rate=24000, target_rate=25, smoothing_kernel_size=7, normalize=False
):
    """Calculate all spectral features from audio data.

    Args:
        audio_data: Audio array (can be stereo or mono)
        sample_rate: Sample rate of audio
        target_rate: Target resolution in Hz (default 25Hz)
        smoothing_kernel_size: Size of smoothing kernel (default 7)
        normalize: If True, normalize each feature to [0, 1] range

    Returns:
        dict: Dictionary containing all calculated features
    """
    # Resample to 24kHz if needed
    if sample_rate != 24000:
        print(f"Warning: Audio is at {sample_rate}Hz, features expect 24kHz")
        # TODO: Implement resampling if needed
        # For now, assume audio is already at correct rate

    # Convert to mono for visualization
    if audio_data.ndim > 1:
        audio_mono = np.mean(audio_data, axis=0)
    else:
        audio_mono = audio_data

    # Calculate loudness first (for use in silence detection if normalizing)
    loudness = calculate_loudness_seq(
        audio_mono,
        sample_rate=sample_rate,
        target_rate=target_rate,
        smoothing_kernel_size=smoothing_kernel_size,
        normalize=normalize,
    )

    # Calculate spectral features, passing loudness for silence detection if normalizing
    features = {
        "loudness": loudness,
        "spectral_centroid": calculate_spectral_centroid_seq(
            audio_mono,
            sample_rate=sample_rate,
            target_rate=target_rate,
            smoothing_kernel_size=smoothing_kernel_size,
            normalize=normalize,
            loudness_seq=loudness if normalize else None,
        ),
        "spectral_complexity": calculate_spectral_complexity_seq(
            audio_mono,
            sample_rate=sample_rate,
            target_rate=target_rate,
            smoothing_kernel_size=smoothing_kernel_size,
            normalize=normalize,
            loudness_seq=loudness if normalize else None,
        ),
    }

    return features, audio_mono


def crop_features_and_audio(audio_mono, features, start_s, end_s, sample_rate=24000, target_rate=25):
    """Crop both audio and features to specified time range.

    This mimics what happens in the training pipeline.

    Args:
        audio_mono: Mono audio array
        features: Dict of feature arrays
        start_s: Start time in seconds
        end_s: End time in seconds
        sample_rate: Audio sample rate
        target_rate: Feature resolution in Hz (default 25Hz)

    Returns:
        tuple: (cropped_audio, cropped_features)
    """
    # Crop audio
    start_sample = int(start_s * sample_rate)
    end_sample = int(end_s * sample_rate)
    cropped_audio = audio_mono[start_sample:end_sample]

    # Crop features (at target_rate resolution)
    start_frame = int(start_s * target_rate)
    end_frame = int(end_s * target_rate)

    cropped_features = {}
    for key, feature_array in features.items():
        cropped_features[key] = feature_array[start_frame:end_frame]

    # Verify alignment
    expected_frames = len(cropped_audio) / sample_rate * target_rate
    actual_frames = len(cropped_features["loudness"])
    if abs(expected_frames - actual_frames) > 1:
        print(
            f"Warning: Feature alignment issue. Expected {expected_frames:.1f} frames, got {actual_frames}"
        )

    return cropped_audio, cropped_features


def generate_tag_examples(features, normalized=False):
    """Generate example tags from features (simplified version).

    Args:
        features: Dict of feature arrays
        normalized: Whether features are normalized to [0, 1] range

    Returns:
        str: Example tag string
    """
    tags = []

    # Loudness tags
    loudness = features["loudness"]
    if normalized:
        activity = np.sum(loudness > 0.1) / len(loudness) * 100
    else:
        activity = np.sum(loudness > 0.01) / len(loudness) * 100
    tags.append(f"activity:{int(round(activity / 10) * 10)}%")

    # Spectral centroid tags
    centroid = features["spectral_centroid"]
    mean_centroid = np.mean(centroid)

    if normalized:
        # Centroid is in [0, 1] range
        if mean_centroid > 0.7:
            tags.append("centroid:bright")
        elif mean_centroid < 0.3:
            tags.append("centroid:dark")
        else:
            tags.append("centroid:balanced")
        tags.append(f"centroid_mean:{mean_centroid:.2f}")

        # Example span with normalized values (first 10 frames)
        if len(centroid) >= 10:
            span_values = centroid[:10]
            values_str = ",".join([f"{v:.2f}" for v in span_values])
            tags.append(f"centroid[0:10]:[{values_str}]")
    else:
        # Centroid is in Hz
        if mean_centroid > 8000:
            tags.append("centroid:bright")
        elif mean_centroid < 3000:
            tags.append("centroid:dark")
        else:
            tags.append("centroid:balanced")
        tags.append(f"centroid_mean:{int(mean_centroid)}hz")

        # Example span with Hz values (first 10 frames)
        if len(centroid) >= 10:
            span_values = centroid[:10]
            values_str = ",".join([f"{int(round(v, -2))}hz" for v in span_values])
            tags.append(f"centroid[0:10]:[{values_str}]")

    # Spectral complexity tags (always in 0-1 range, but normalization stretches it)
    complexity = features["spectral_complexity"]
    mean_complexity = np.mean(complexity)
    if mean_complexity > 0.7:
        tags.append("complexity:rich")
    elif mean_complexity < 0.4:
        tags.append("complexity:simple")
    else:
        tags.append("complexity:moderate")
    tags.append(f"complexity_mean:{mean_complexity:.2f}")

    return "{" + ";".join(tags) + "}"


def visualize_features(
    audio_path,
    start_s=None,
    end_s=None,
    show_tags=False,
    output_path=None,
    target_rate=25,
    smoothing_kernel_size=7,
    normalize=False,
):
    """Visualize audio and spectral features.

    Args:
        audio_path: Path to audio file
        start_s: Optional start time for cropping
        end_s: Optional end time for cropping
        show_tags: Whether to show generated tags
        output_path: Optional path to save figure
        target_rate: Target resolution in Hz (default 25Hz)
        smoothing_kernel_size: Size of smoothing kernel (default 7)
        normalize: If True, normalize each feature to [0, 1] range
    """
    print(f"Loading audio from: {audio_path}")
    audio = load_audio(audio_path)

    # Get audio data
    audio_data = audio.array_float
    sample_rate = audio.sample_rate

    # Handle both mono and stereo audio shapes
    if audio_data.ndim == 1:
        duration_s = len(audio_data) / sample_rate
        n_channels = 1
    else:
        duration_s = audio_data.shape[1] / sample_rate
        n_channels = audio_data.shape[0]

    print(f"Audio: {n_channels} channels, {sample_rate}Hz, {duration_s:.2f}s")

    # Calculate features on full audio
    norm_str = " (normalized)" if normalize else ""
    print(
        f"Calculating features at {target_rate}Hz with kernel size {smoothing_kernel_size}{norm_str}..."
    )
    features, audio_mono = calculate_features(
        audio_data,
        sample_rate,
        target_rate=target_rate,
        smoothing_kernel_size=smoothing_kernel_size,
        normalize=normalize,
    )

    # Apply cropping if specified
    if start_s is not None or end_s is not None:
        start_s = start_s or 0.0
        end_s = end_s or (len(audio_mono) / sample_rate)
        print(f"Cropping to {start_s:.2f}s - {end_s:.2f}s")
        audio_mono, features = crop_features_and_audio(
            audio_mono, features, start_s, end_s, sample_rate, target_rate=target_rate
        )
        time_offset = start_s
    else:
        time_offset = 0.0

    # Generate tags if requested
    tag_string = ""
    if show_tags:
        tag_string = generate_tag_examples(features, normalized=normalize)
        print(f"\nGenerated tags:\n{tag_string}\n")

    # Create visualization
    fig, axes = plt.subplots(4, 1, figsize=(14, 10))
    fig.suptitle(f"Spectral Features: {os.path.basename(audio_path)}", fontsize=14, fontweight="bold")

    # Time arrays - calculate actual rate for each feature based on its length
    audio_time = np.arange(len(audio_mono)) / sample_rate + time_offset
    audio_duration = len(audio_mono) / sample_rate

    # Calculate actual rates from feature lengths and audio duration
    loudness_actual_rate = len(features["loudness"]) / audio_duration
    centroid_actual_rate = len(features["spectral_centroid"]) / audio_duration
    complexity_actual_rate = len(features["spectral_complexity"]) / audio_duration

    # Create separate time arrays for each feature based on their actual lengths
    loudness_time = np.arange(len(features["loudness"])) / loudness_actual_rate + time_offset
    centroid_time = np.arange(len(features["spectral_centroid"])) / centroid_actual_rate + time_offset
    complexity_time = (
        np.arange(len(features["spectral_complexity"])) / complexity_actual_rate + time_offset
    )

    # Row 1: Spectrogram
    nperseg = min(2048, len(audio_mono) // 4)  # Window size for spectrogram
    f, t, Sxx = signal.spectrogram(audio_mono, sample_rate, nperseg=nperseg, noverlap=nperseg // 2)
    spec_db = 10 * np.log10(Sxx + 1e-10)  # Convert to dB

    # Adjust time axis for offset
    t_adjusted = t + time_offset

    axes[0].pcolormesh(
        t_adjusted,
        f / 1000,
        spec_db,
        shading="gouraud",
        cmap="viridis",
        vmin=np.percentile(spec_db, 5),
        vmax=np.percentile(spec_db, 95),
    )
    axes[0].set_ylabel("Frequency (kHz)")
    axes[0].set_title("Spectrogram", fontsize=10, fontweight="bold")
    axes[0].set_ylim(0, 12)  # Show up to 12 kHz
    axes[0].set_xlim(audio_time[0], audio_time[-1])

    # Row 2: Loudness
    axes[1].plot(loudness_time, features["loudness"], linewidth=2, color="green")

    if normalize:
        axes[1].set_ylabel("Normalized Loudness (0-1)")
        axes[1].set_title(
            f"Loudness @ {loudness_actual_rate:.1f}Hz (Normalized)", fontsize=10, fontweight="bold"
        )
        axes[1].set_ylim(-0.05, 1.05)
    else:
        axes[1].axhline(
            y=0.01, color="red", linestyle="--", linewidth=1, alpha=0.5, label="Activity threshold"
        )
        axes[1].set_ylabel("RMS Loudness")
        axes[1].set_title(f"Loudness @ {loudness_actual_rate:.1f}Hz", fontsize=10, fontweight="bold")
        axes[1].legend(loc="upper right", fontsize=8)

    axes[1].grid(True, alpha=0.3)
    axes[1].set_xlim(audio_time[0], audio_time[-1])

    # Row 3: Spectral Centroid
    axes[2].plot(centroid_time, features["spectral_centroid"], linewidth=2, color="orange")

    if normalize:
        axes[2].set_ylabel("Normalized Brightness (0-1)")
        axes[2].set_title(
            f"Spectral Centroid @ {centroid_actual_rate:.1f}Hz (Normalized)",
            fontsize=10,
            fontweight="bold",
        )
        axes[2].set_ylim(-0.05, 1.05)
    else:
        axes[2].set_ylabel("Frequency (Hz)")
        axes[2].set_title(
            f"Spectral Centroid @ {centroid_actual_rate:.1f}Hz (Brightness)",
            fontsize=10,
            fontweight="bold",
        )
        # Add reference lines for brightness categories (only in non-normalized mode)
        axes[2].axhline(
            y=8000, color="yellow", linestyle="--", linewidth=1, alpha=0.3, label="Bright threshold"
        )
        axes[2].axhline(
            y=3000, color="purple", linestyle="--", linewidth=1, alpha=0.3, label="Dark threshold"
        )
        axes[2].legend(loc="upper right", fontsize=8)

    axes[2].grid(True, alpha=0.3)
    axes[2].set_xlim(audio_time[0], audio_time[-1])

    # Row 4: Spectral Complexity
    axes[3].plot(complexity_time, features["spectral_complexity"], linewidth=2, color="purple")
    axes[3].set_ylabel("Complexity (0-1)")
    axes[3].set_xlabel("Time (seconds)")

    if normalize:
        axes[3].set_title(
            f"Spectral Complexity @ {complexity_actual_rate:.1f}Hz (Normalized)",
            fontsize=10,
            fontweight="bold",
        )
    else:
        axes[3].set_title(
            f"Spectral Complexity @ {complexity_actual_rate:.1f}Hz (Entropy)",
            fontsize=10,
            fontweight="bold",
        )
        # Add reference lines for complexity categories (only in non-normalized mode)
        axes[3].axhline(
            y=0.7, color="red", linestyle="--", linewidth=1, alpha=0.3, label="Rich threshold"
        )
        axes[3].axhline(
            y=0.4, color="blue", linestyle="--", linewidth=1, alpha=0.3, label="Simple threshold"
        )
        axes[3].legend(loc="upper right", fontsize=8)

    axes[3].grid(True, alpha=0.3)
    axes[3].set_ylim(-0.05, 1.05)
    axes[3].set_xlim(audio_time[0], audio_time[-1])

    # Add tags as text if requested
    if show_tags:
        fig.text(
            0.5,
            0.02,
            f"Tags: {tag_string}",
            ha="center",
            fontsize=9,
            bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.5),
            wrap=True,
        )

    plt.tight_layout(rect=[0, 0.03, 1, 0.97])

    if output_path:
        print(f"Saving figure to: {output_path}")
        plt.savefig(output_path, dpi=150, bbox_inches="tight")

    plt.show()


def main():
    parser = argparse.ArgumentParser(description="Visualize spectral features from audio file")
    parser.add_argument("--audio_path", type=str, required=True, help="Path to audio file")
    parser.add_argument(
        "--start_s", type=float, default=None, help="Start time in seconds (for cropping)"
    )
    parser.add_argument("--end_s", type=float, default=None, help="End time in seconds (for cropping)")
    parser.add_argument("--show_tags", action="store_true", help="Generate and show example tags")
    parser.add_argument("--output_path", type=str, default=None, help="Path to save figure")

    args = parser.parse_args()

    if not os.path.exists(args.audio_path):
        print(f"Error: Audio file not found: {args.audio_path}")
        sys.exit(1)

    visualize_features(
        args.audio_path,
        start_s=args.start_s,
        end_s=args.end_s,
        show_tags=args.show_tags,
        output_path=args.output_path,
    )


if __name__ == "__main__":
    main()