"""
TensorFlow Audio Processing Module

TensorFlow-based audio processing and spectrogram generation utilities for soundhub

License: BSD 3-Clause
"""
def safe_load_tensorflow(legacy_keras):
    if legacy_keras:
        os.environ['TF_USE_LEGACY_KERAS'] = str(int(legacy_keras))
    import tensorflow as tf
    return tf


import os
import ast
from pathlib import Path
from typing import Union, Optional, Tuple, List, Dict, Any
import numpy as np
import pandas as pd
import librosa
from soundhub_utils.constants import SUPPORTED_AUDIO_FORMATS
from soundhub_utils.audio import _shared
from soundhub_utils.audio._shared import AudioData
import sox_tensorflow.processor
from cocina.config_handler import ConfigHandler
from importlib import reload
ch = ConfigHandler(__file__)
tf = safe_load_tensorflow(ch.get('TF_USE_LEGACY_KERAS'))


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

    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 TensorFlow tensor
    audio_tensor = tf.convert_to_tensor(audio_np, dtype=tf.float32)

    # 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 spectrogram(
    audio_array: Union['tf.Tensor', np.ndarray],
    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,
    mel: bool = False,
    frame_length: Optional[int] = None,
    frame_step: Optional[int] = None,
    colormap: Optional[str] = 'viridis',
    create_parents: bool = True,
    overwrite: bool = True
) -> Union['tf.Tensor', str]:
    """
    Convert audio array to spectrogram using TensorFlow.

    Args:
        audio_array: Audio data as TensorFlow tensor or numpy array
        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
        mel: If True, compute mel-spectrogram instead of regular spectrogram
        frame_length: STFT frame length (defaults to config value or 512)
        frame_step: STFT frame step (defaults to config value or 256)
        colormap: Matplotlib colormap ('gray', 'viridis', etc.) or None for raw grayscale
        create_parents: Create parent directories if needed
        overwrite: Overwrite existing file if it exists

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

    Raises:
        ValueError: If invalid parameters provided
        FileExistsError: If file exists and overwrite=False
    """
    if sample_rate is None:
        sample_rate = _shared.DEFAULT_SAMPLE_RATE

    # Convert to TensorFlow tensor if needed
    if isinstance(audio_array, np.ndarray):
        audio_tensor = tf.convert_to_tensor(audio_array, dtype=tf.float32)
    else:
        audio_tensor = audio_array

    # 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]
    if mel:
        return mel_spectrogram(
            audio_tensor, shape, dest, None, None, None, sample_rate,
            frame_length=frame_length, frame_step=frame_step,
            create_parents=create_parents, overwrite=overwrite
        )

    # Get STFT parameters from function args, config, or defaults
    if frame_length is None:
        frame_length = ch.get('STFT_FRAME_LENGTH', 512)
    if frame_step is None:
        frame_step = ch.get('STFT_FRAME_STEP', 96)

    stft = tf.signal.stft(
        audio_tensor,
        frame_length=frame_length,
        frame_step=frame_step,
        window_fn=tf.signal.hann_window
    )

    # Convert to power spectrogram (magnitude squared)
    power_spectrogram = tf.square(tf.abs(stft))

    # Convert to dB scale with 90 dB dynamic range (matching SoX -z 90)
    # SoX formula: 10 * log10(power) with 90 dB range
    spectrogram_db = 10.0 * tf.math.log(tf.maximum(power_spectrogram, 1e-9)) / tf.math.log(10.0)

    # Apply 90 dB dynamic range clipping (like SoX -z 90)
    db_max = tf.reduce_max(spectrogram_db)
    db_min = db_max - 90.0  # 90 dB dynamic range
    spectrogram_tensor = tf.clip_by_value(spectrogram_db, db_min, db_max)


    # Resize to target shape if needed
    if tuple(spectrogram_tensor.shape) != shape:
        # Ensure shape is a tuple or list of integers
        if isinstance(shape, str):
            # Handle string representation like "(257, 1000)"
            shape = ast.literal_eval(shape)

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

        # Add batch and channel dimensions for resize
        spec_resized = tf.image.resize(
            tf.expand_dims(tf.expand_dims(spectrogram_tensor, 0), -1),
            tf.convert_to_tensor([target_height, target_width], dtype=tf.int32)
        )
        spectrogram_tensor = tf.squeeze(spec_resized)

    if dest is None:
        return spectrogram_tensor

    # Save spectrogram as PNG
    _shared.save_spectrogram_png(spectrogram_tensor, dest, create_parents, overwrite, colormap)
    return str(dest)


def mel_spectrogram(
    audio_array: Union['tf.Tensor', np.ndarray],
    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,
    num_mel_bins: int = 128,
    frame_length: Optional[int] = None,
    frame_step: Optional[int] = None,
    colormap: Optional[str] = 'viridis',
    create_parents: bool = True,
    overwrite: bool = True
) -> Union['tf.Tensor', str]:
    """
    Convert audio array to mel-spectrogram using tf.keras.utils.audio.

    Args:
        audio_array: Audio data as TensorFlow tensor or numpy array
        shape: Output shape of mel-spectrogram (height, width)
        dest: Optional file path to save mel-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
        num_mel_bins: Number of mel frequency bins
        frame_length: STFT frame length (defaults to config value or 512)
        frame_step: STFT frame step (defaults to config value or 256)
        colormap: Matplotlib colormap ('gray', 'viridis', etc.) or None for raw grayscale
        create_parents: Create parent directories if needed
        overwrite: Overwrite existing file if it exists

    Returns:
        Mel-spectrogram tensor if dest is None, otherwise path to saved PNG file

    Raises:
        ValueError: If invalid parameters provided
        FileExistsError: If file exists and overwrite=False
    """
    if sample_rate is None:
        sample_rate = _shared.DEFAULT_SAMPLE_RATE

    # Convert to TensorFlow tensor if needed
    if isinstance(audio_array, np.ndarray):
        audio_tensor = tf.convert_to_tensor(audio_array, dtype=tf.float32)
    else:
        audio_tensor = audio_array

    # 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]

    # Use TensorFlow's built-in mel-spectrogram computation
    # First compute STFT with configurable parameters
    if frame_length is None:
        frame_length = ch.get('STFT_FRAME_LENGTH', 512)
    if frame_step is None:
        frame_step = ch.get('STFT_FRAME_STEP', 96)

    stft = tf.signal.stft(
        audio_tensor,
        frame_length=frame_length,
        frame_step=frame_step,
        window_fn=tf.signal.hann_window
    )

    # Convert to power spectrogram
    power_spectrogram = tf.real(stft * tf.math.conj(stft))

    # Create mel-scale filterbank
    linear_to_mel_weight_matrix = tf.signal.linear_to_mel_weight_matrix(
        num_mel_bins=num_mel_bins,
        num_spectrogram_bins=power_spectrogram.shape[-2],
        sample_rate=sample_rate,
        lower_edge_hertz=0.0,
        upper_edge_hertz=sample_rate / 2.0
    )

    # Apply mel-scale filterbank
    mel_spec = tf.tensordot(power_spectrogram, linear_to_mel_weight_matrix, 1)

    # Resize to target shape if needed
    if tuple(mel_spec.shape) != shape:
        # Ensure shape is a tuple or list of integers
        if isinstance(shape, str):
            # Handle string representation like "(257, 1000)"
            shape = ast.literal_eval(shape)

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

        # Add batch and channel dimensions for resize
        mel_resized = tf.image.resize(
            tf.expand_dims(tf.expand_dims(mel_spec, 0), -1),
            tf.convert_to_tensor([target_height, target_width], dtype=tf.int32)
        )
        mel_spec = tf.squeeze(mel_resized)

    if dest is None:
        return mel_spec

    # Save mel-spectrogram as PNG
    _shared.save_spectrogram_png(mel_spec, dest, create_parents, overwrite, colormap)
    return str(dest)


# TensorFlow backend no longer provides spectrograms() method.
# Use the main audio.spectrograms() function with TensorFlow-specific spectrogram functions:
#   from soundhub_utils.audio import spectrograms
#   from soundhub_utils.audio import tensorflow_audio
#   spectrograms('audio.wav', 'output/', (257, 1000), tensorflow_audio.spectrogram)
#   spectrograms('audio.wav', 'output/', (128, 1000), tensorflow_audio.mel_spectrogram)


def inference(
    model: Union[str, tf.keras.Model],
    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]:
    """
    Run TensorFlow model inference on spectrogram files (PNG or TFRecord).

    Automatically detects input format based on file extension.

    Args:
        model: Either path to .h5 model file or loaded TensorFlow model
        files: List of file paths (PNG or TFRecord) or DataFrame with file paths
        spectrogram_shape: Expected spectrogram shape (height, width)
        confidence_threshold: Minimum confidence for predictions
        top_k: Number of top predictions to return per file
        class_names: List of class names (if None, uses indices)
        return_format: Output format ('dict' or 'dataframe')

    Returns:
        Dictionary or DataFrame containing inference results with predictions,
        confidences, and metadata for each input file

    Raises:
        FileNotFoundError: If model file not found
        ValueError: If invalid parameters provided
        IOError: If inference fails
    """
    # Set default spectrogram shape if not provided
    if spectrogram_shape is None:
        spectrogram_shape = _shared.DEFAULT_SPECTROGRAM_SHAPE

    # Load model if path provided
    if isinstance(model, str):
        if not os.path.exists(model):
            raise FileNotFoundError(f"Model file not found: {model}")
        tf_model = tf.keras.models.load_model(model)
    else:
        tf_model = model

    # Convert files input to list of paths
    if isinstance(files, pd.DataFrame):
        if 'Filename' in files.columns:
            file_paths = files['Filename'].tolist()
        elif len(files.columns) == 1:
            file_paths = files.iloc[:, 0].tolist()
        else:
            raise ValueError("DataFrame must have 'Filename' column or single column")
    else:
        file_paths = files

    # Validate file paths
    for file_path in file_paths:
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"Spectrogram file not found: {file_path}")

    # Auto-detect input format (TFRecord vs PNG)
    from soundhub_utils.utils.tf import is_tfrecord_path
    is_tfrecord = is_tfrecord_path(file_paths)

    if is_tfrecord:
        # Load from TFRecord file(s)
        from soundhub_utils.utils.tf import load_spectrograms_from_tfrecord
        dataset = load_spectrograms_from_tfrecord(
            file_paths,
            spectrogram_shape,
            batch_size=100,
            compression=ch.get('TFRECORD_COMPRESSION', 'GZIP')
        )

        # Run inference on batches and collect results
        all_predictions = []
        all_metadata = []

        for spectrograms_batch, metadata_batch in dataset:
            batch_predictions = tf_model.predict(spectrograms_batch, verbose=0)
            all_predictions.append(batch_predictions)
            all_metadata.append(metadata_batch)

        # Concatenate all predictions
        predictions = np.concatenate(all_predictions, axis=0)

        # Build file paths from metadata for results processing
        # (TFRecord stores segment_index and filename, reconstruct paths if needed)
        result_file_paths = []
        for metadata_batch in all_metadata:
            for i in range(len(metadata_batch['segment_index'])):
                filename = metadata_batch['filename'][i].numpy().decode('utf-8')
                segment_idx = metadata_batch['segment_index'][i].numpy()
                # Create pseudo-path for results tracking
                result_file_paths.append(f"{filename}_segment{segment_idx:04d}")

    else:
        # Load spectrograms from PNG files
        from soundhub_utils.utils.tf import load_spectrograms_tf
        input_data = load_spectrograms_tf(file_paths, spectrogram_shape)
        all_predictions = []
        for batch in input_data:
            batch_predictions = tf_model(batch, training=False).numpy()
            all_predictions.append(batch_predictions)
        predictions = np.concatenate(all_predictions, axis=0)

        result_file_paths = file_paths

    # Process results
    results = _shared.process_inference_results(
        predictions, result_file_paths, confidence_threshold, top_k, class_names
    )


    if return_format == 'dataframe':
        return pd.DataFrame(results)
    else:
        return results




