"""
PyTorch Audio Processing Module

PyTorch-based audio processing and spectrogram generation utilities for soundhub

License: BSD 3-Clause
"""

import ast
import os
from pathlib import Path
from typing import Union, Optional, Tuple, List, Dict, Any
import numpy as np
import matplotlib.pyplot as plt
import torch
import pandas as pd
import librosa
from soundhub_utils import names
from soundhub_utils.constants import SUPPORTED_AUDIO_FORMATS
from soundhub_utils.audio import _shared
from soundhub_utils.audio._shared import AudioData
from cocina.printer import Printer
from cocina.config_handler import ConfigHandler
printer = Printer()
ch = ConfigHandler(__file__)


#
# DEFAULTS (from ConfigHandler for module-specific config)
#


#
# MAIN
#
def load_audio(
    src: Union[str, Path],
    sample_rate: Optional[int] = None,
    file_type: Optional[str] = None
) -> AudioData:
    """
    Load audio file using PyTorch and librosa.

    Note: Uses librosa for loading (same as tensorflow_audio) to avoid torchaudio
    hanging issues with large files.

    Args:
        src: Path to audio file (FLAC, WAV, etc.)
        sample_rate: Target sample rate in Hz. If None, uses original sample rate
        file_type: Audio file type. If None, auto-detected from extension

    Returns:
        AudioData containing tensor, sample rate, duration, and channels

    Raises:
        IOError: If file cannot be loaded
        ValueError: If invalid file type or sample rate specified
    """
    src_path = Path(src)

    # Auto-detect file type if not provided
    if file_type is None:
        file_type = src_path.suffix.lower()

    if file_type not in SUPPORTED_AUDIO_FORMATS:
        raise ValueError(f"Unsupported audio format: {file_type}")

    # Load audio using librosa (handles FLAC, WAV, MP3, etc.)
    # librosa automatically handles resampling when sr parameter is provided
    audio_np, loaded_sr = librosa.load(str(src_path), sr=sample_rate, mono=True)

    # Convert to PyTorch tensor
    audio_tensor = torch.from_numpy(audio_np).float()

    # Calculate duration
    duration = len(audio_np) / loaded_sr

    return AudioData(
        tensor=audio_tensor,
        sample_rate=loaded_sr,
        duration=duration,
        channels=1  # librosa loads as mono
    )


def to_array(audio_data: AudioData, to_numpy: bool = False) -> Union[torch.Tensor, np.ndarray]:
    """
    Extract audio tensor from AudioData.

    Args:
        audio_data: AudioData container
        to_numpy: If True, return numpy array; otherwise return PyTorch tensor

    Returns:
        Audio array as numpy array or PyTorch tensor
    """
    if to_numpy:
        return audio_data.tensor.numpy()
    return audio_data.tensor


def pysox_spectrogram(
    audio_array: Union[torch.Tensor, np.ndarray, AudioData],
    shape: Tuple[int, int],
    dest: Optional[Union[str, Path]] = None,
    segment: Optional[int] = None,
    segment_duration: Optional[float] = None,
    segment_overlap: Optional[float] = None,
    sample_rate: Optional[int] = None,
    create_parents: bool = True,
    overwrite: bool = True
) -> Union[torch.Tensor, str]:
    """
    Create Sox-matching spectrogram using native PyTorch with absolute dB normalization.

    Uses absolute dB normalization (Method 2) for stable, consistent Sox matching.
    Achieves ~91.8% within ±5 pixels, stable across different sample sizes.

    Args:
        audio_array: Audio data as PyTorch tensor, numpy array, or AudioData
        shape: Output shape of spectrogram (height, width)
        dest: Optional file path to save spectrogram as PNG
        segment: Segment number for extracting specific audio portion
        segment_duration: Duration of segment in seconds
        segment_overlap: Overlap between segments in seconds
        sample_rate: Audio sample rate in Hz (required for tensor/array inputs)
        create_parents: Create parent directories if needed
        overwrite: Overwrite existing file if it exists

    Returns:
        Spectrogram tensor (uint8) if dest is None, otherwise path to saved PNG file

    Raises:
        ValueError: If invalid parameters provided
        FileExistsError: If file exists and overwrite=False

    Notes:
        - Uses exact Sox algorithm: n_fft=512, hop_length=20
        - Block averaging matches Sox's implementation
        - dB_range=80.0, brightness_correction=-96 for Sox compatibility
        - Output is uint8 grayscale (4-231 range like Sox)
        - Absolute dB normalization (Method 2) is stable across sample sizes
        - See /workspace/py-sox-testing/sh/NORMALIZATION_FIX.md for details
    """
    # Extract AudioData if provided
    if isinstance(audio_array, AudioData):
        audio_tensor = audio_array.tensor
        if sample_rate is None:
            sample_rate = audio_array.sample_rate
    elif isinstance(audio_array, np.ndarray):
        audio_tensor = torch.from_numpy(audio_array).float()
    else:
        audio_tensor = audio_array

    if sample_rate is None:
        sample_rate = _shared.DEFAULT_SAMPLE_RATE

    # Ensure shape is a tuple or list of integers
    if isinstance(shape, str):
        shape = ast.literal_eval(shape)

    target_height, target_width = shape[0], shape[1]

    # Extract segment if specified
    if segment is not None:
        start_idx, end_idx = _shared.calculate_segment_indices(
            len(audio_tensor), sample_rate, segment, segment_duration, segment_overlap
        )
        audio_tensor = audio_tensor[start_idx:end_idx]

    # Get Sox parameters from config
    n_fft = ch.get('SOX_N_FFT', 512)
    hop_length = ch.get('SOX_HOP_LENGTH', 20)
    dB_range = ch.get('SOX_DB_RANGE', 80.0)

    # METHOD 2: Always use absolute dB normalization (more stable across sample sizes)
    # NOTE: Ignoring reference_db parameter - using absolute dB regardless
    brightness_correction = -96  # Optimized for Method 2 (absolute dB)

    # Compute STFT with Hann window and center padding
    window = torch.hann_window(n_fft)
    D = torch.stft(
        audio_tensor,
        n_fft=n_fft,
        hop_length=hop_length,
        win_length=n_fft,
        window=window,
        center=True,
        return_complex=True
    )

    # Compute power spectrum
    S_power = torch.abs(D) ** 2  # Shape: (freq_bins, time_frames)

    # Block averaging (Sox's exact algorithm) - NATIVE PYTORCH
    n_frames = S_power.shape[1]
    frames_per_block = n_frames / target_width

    # Pre-allocate output tensor
    S_power_accumulated = torch.zeros((S_power.shape[0], target_width), dtype=torch.float32)

    # Block averaging using PyTorch operations
    for col_idx in range(target_width):
        start_frame = col_idx * frames_per_block
        end_frame = (col_idx + 1) * frames_per_block

        start_idx = int(start_frame)
        end_idx = int(torch.ceil(torch.tensor(end_frame)).item())

        # Ensure we don't exceed bounds
        end_idx = min(end_idx, n_frames)

        # Accumulate power (stays in PyTorch)
        block = S_power[:, start_idx:end_idx]
        S_power_accumulated[:, col_idx] = torch.mean(block, dim=1)

    # Convert to dB
    S_db = 10.0 * torch.log10(S_power_accumulated + 1e-10)

    # METHOD 2: Absolute dB normalization (no reference subtraction)
    # - More stable across different sample sizes
    # - No first-pass needed to calculate global_peak_db
    # - Achieved 91.8% within ±5px in testing
    S_db_norm = S_db  # Use absolute dB values directly

    # Map to pixel values
    S_pixels = 4.0 + torch.clamp((S_db_norm + dB_range) / dB_range * 227.0, 0.0, 227.0)
    S_pixels = S_pixels + brightness_correction
    S_pixels = torch.clamp(S_pixels, 4.0, 231.0)

    # Flip vertically
    S_pixels = torch.flipud(S_pixels)

    # Convert to uint8
    spectrogram_tensor = S_pixels.to(torch.uint8)

    # Verify dimensions
    assert spectrogram_tensor.shape == (target_height, target_width), \
        f"Output shape {spectrogram_tensor.shape} != expected ({target_height}, {target_width})"

    if dest is None:
        return spectrogram_tensor

    # Convert to numpy and save using shared function
    spec_np = spectrogram_tensor.cpu().numpy()
    _shared.save_spectrogram_png(spec_np, dest, create_parents, overwrite, colormap=None)
    return str(dest)


def spectrograms(
    src: Union[str, Path, AudioData, torch.Tensor, np.ndarray],
    dest_folder: str,
    spectrogram_shape: Tuple[int, int],
    basename: Optional[str] = None,
    directory: Optional[str] = None,
    segment_duration: Optional[float] = None,
    segment_overlap: Optional[float] = None,
    sample_rate: Optional[int] = None,
    spectrogram_type: Optional[str] = None,
    create_parents: bool = True,
    overwrite: bool = True,
    output_format: Optional[str] = None,
    tfrecord_compression: Optional[str] = None,
    tfrecord_chunk_size: Optional[int] = None,
    metadata: Optional[Dict[str, Any]] = None
) -> Union[List[Union[torch.Tensor, str]], str]:
    """
    Generate multiple spectrograms from audio source with segmentation.

    Args:
        src: Audio source - file path, AudioData, PyTorch tensor, or numpy array
        dest_folder: Folder to save spectrogram files
        spectrogram_shape: Output shape for each spectrogram (height, width)
        basename: Base name for output files (auto-determined if None)
        directory: Directory name within dest_folder (optional)
        segment_duration: Duration of each segment in seconds
        segment_overlap: Overlap between segments in seconds
        sample_rate: Audio sample rate in Hz (required for tensor/array inputs)
        spectrogram_type: Type of spectrogram ('sox' only for PyTorch backend)
        create_parents: Create parent directories if needed
        overwrite: Overwrite existing files if they exist
        output_format: Output format - 'png' or 'tfrecord' (default from config)
        tfrecord_compression: TFRecord compression - 'GZIP', 'ZLIB', or None
        tfrecord_chunk_size: Split TFRecord into chunks of this size (None = single file)
        metadata: Additional metadata to store in TFRecord (e.g., audio_uri, start_time)

    Returns:
        If output_format='png': List of spectrogram tensors or file paths
        If output_format='tfrecord': List of TFRecord file paths

    Raises:
        ValueError: If invalid parameters provided
        IOError: If audio loading or file operations fail
        FileExistsError: If files exist and overwrite=False
    """
    # Validate spectrogram_type (PyTorch only supports pysox and sox for now)
    if spectrogram_type not in [None, 'sox', 'pysox']:
        raise ValueError(f"PyTorch backend only supports spectrogram_type=None, 'sox', or 'pysox', got: {spectrogram_type}")
    printer.message(1)

    # Get output format from config if not specified
    if output_format is None:
        output_format = ch.get('TFRECORD_OUTPUT_FORMAT', 'png')
    if tfrecord_compression is None:
        tfrecord_compression = ch.get('TFRECORD_COMPRESSION', 'GZIP')
    if tfrecord_chunk_size is None:
        tfrecord_chunk_size = ch.get('TFRECORD_CHUNK_SIZE', None)

    if segment_duration is None:
        segment_duration = _shared.DEFAULT_SEGMENT_DURATION
    if segment_overlap is None:
        segment_overlap = 0.0
    printer.message(2)

    # Handle different source types
    _audio_path = None
    _native_sr = None
    _total_frames = None

    if isinstance(src, (str, Path)):
        # Streaming mode: read only file header, don't load audio into memory
        audio_info = _shared.get_audio_info(src)
        actual_sample_rate = sample_rate if sample_rate is not None else audio_info.sample_rate
        audio_tensor = None
        _audio_path = str(src)
        _native_sr = audio_info.sample_rate
        _total_frames = audio_info.frames
        audio_length_seconds = audio_info.duration
        if basename is None:
            basename = Path(src).stem
    elif isinstance(src, AudioData):
        # Use AudioData directly
        audio_tensor = src.tensor
        actual_sample_rate = src.sample_rate
        audio_length_seconds = len(audio_tensor) / actual_sample_rate
        if basename is None:
            basename = "audio_data"
    elif isinstance(src, np.ndarray):
        # Convert numpy array to tensor
        audio_tensor = torch.from_numpy(src).float()
        if sample_rate is None:
            raise ValueError("sample_rate must be provided for numpy array input")
        actual_sample_rate = sample_rate
        audio_length_seconds = len(audio_tensor) / actual_sample_rate
        if basename is None:
            basename = "audio_array"
    elif isinstance(src, torch.Tensor):
        # Use tensor directly
        audio_tensor = src
        if sample_rate is None:
            raise ValueError("sample_rate must be provided for tensor input")
        actual_sample_rate = sample_rate
        audio_length_seconds = len(audio_tensor) / actual_sample_rate
        if basename is None:
            basename = "audio_tensor"
    else:
        raise ValueError(f"Unsupported source type: {type(src)}")
    printer.message(3)

    # Calculate number of segments
    num_segments = _shared.calculate_num_segments(
        audio_length_seconds,
        segment_duration,
        segment_overlap
    )
    printer.message(4)

    # Select spectrogram function based on type
    if spectrogram_type == 'pysox':
        spec_func = pysox_spectrogram
    elif spectrogram_type == 'sox':
        spec_func = lambda *args, **kwargs: _shared.sox_spectrogram(*args, **kwargs, backend='pytorch')
    else:
        spec_func = pysox_spectrogram  # Default to pysox

    # Use shared spectrogram generation (handles both PNG and TFRecord)
    return _shared.generate_spectrograms(
        spec_func=spec_func,
        audio_data=audio_tensor,
        num_segments=num_segments,
        spectrogram_shape=spectrogram_shape,
        dest_folder=dest_folder,
        basename=basename,
        directory=directory,
        segment_duration=segment_duration,
        segment_overlap=segment_overlap,
        sample_rate=actual_sample_rate,
        output_format=output_format,
        tfrecord_compression=tfrecord_compression,
        tfrecord_chunk_size=tfrecord_chunk_size,
        metadata=metadata,
        create_parents=create_parents,
        overwrite=overwrite,
        audio_path=_audio_path,
        native_sample_rate=_native_sr,
        total_frames=_total_frames,
    )


def inference(
    model: str,
    files: Union[List[str], pd.DataFrame],
    spectrogram_shape: Optional[Tuple[int, int]] = None,
    confidence_threshold: float = 0.1,
    top_k: int = 5,
    class_names: Optional[List[str]] = None,
    return_format: str = 'dict'
) -> Union[Dict[str, Any], pd.DataFrame]:
    """
    TODO: inference for pytorch models
    """
    raise NotImplementedError
