import sys
import os
import logging
import sys
import asyncio
import numpy as np

print("Python executable:", sys.executable)
print("Python version:", sys.version)
print("PYTHONPATH:", os.environ.get('PYTHONPATH'))

try:
    import openai
    print('OpenAI library is installed correctly.')
except ModuleNotFoundError as e:
    print(f'Error: {e}')

try:
    import madmom
    print('Madmom library is installed correctly.')
except ModuleNotFoundError as e:
    print(f'Error: {e}')
except AttributeError as e:
    print(f'AttributeError: {e}')

import torch
import librosa
import lxml.etree as ET
import noisereduce as nr
import shutil
import soundfile as sf
from torchaudio.transforms import MelSpectrogram
from scipy import signal
from transformers import Wav2Vec2Tokenizer, Wav2Vec2ForCTC, HubertForCTC
import crepe
import madmom
import tensorflow as tf
from diskcache import Cache
from music21 import stream, note, converter, instrument

# НАСТРОЙКА ЛОГИРОВАНИЯ
logging.basicConfig(level=logging.INFO)

# Custom Config
CONFIG = {
    "output_dir": "output",
    "noise_reduction_level": 5,
    "pitch_shift": 2,
    "speed_change": 1.5,
    "batch_size": 16,
    "epochs": 10,
    "target_format": "wav"
}

# ЗАГРУЗКА КЛЮЧА API OPENAI
openai.api_key = os.environ.get('OPENAI_API_KEY')

if openai.api_key:
    try:
        # ПРОБНЫЙ ЗАПРОС К OPENAI API ДЛЯ ПРОВЕРКИ ПОДКЛЮЧЕНИЯ
        response = openai.Completion.create(
            model="gpt-4",
            prompt="Hello, world!",
            max_tokens=5
        )
        logging.info("ПОДКЛЮЧЕНИЕ К OPENAI API УСПЕШНО")
    except openai.error.AuthenticationError:
        logging.error("ОШИБКА АУТЕНТИФИКАЦИИ: НЕВЕРНЫЙ КЛЮЧ OPENAI API")
    except Exception as e:
        logging.error(f"ОШИБКА ПРИ ПОДКЛЮЧЕНИИ К OPENAI API: {e}")
else:
    logging.warning("КЛЮЧ OPENAI API НЕ НАЙДЕН. API БУДЕТ ПРОИГНОРИРОВАН")

# ПРОВЕРКА НАЛИЧИЯ GPU С ПОМОЩЬЮ TENSORFLOW
if tf.test.gpu_device_name():
    print('GPU НАЙДЕН')
else:
    print('GPU НЕ НАЙДЕН, ИСПОЛЬЗУЕТСЯ CPU')

# НАСТРОЙКА КЭШИРОВАНИЯ
cache = Cache('audio_cache')

async def preprocess_audio(file_path):
    """
    АСИНХРОННАЯ ФУНКЦИЯ ДЛЯ ПРЕДВАРИТЕЛЬНОЙ ОБРАБОТКИ АУДИО
    """
    try:
        logging.info(f"ЗАГРУЗКА И ПРЕДВАРИТЕЛЬНАЯ ОБРАБОТКА АУДИО ИЗ {file_path}")
        y, sr = librosa.load(file_path, sr=None)

        if (y, sr) in cache:
            logging.info("ИСПОЛЬЗОВАНИЕ КЭШИРОВАННОГО ПРЕДВАРИТЕЛЬНО ОБРАБОТАННОГО АУДИО")
            return cache[(y, sr)]

        y = librosa.util.normalize(y)
        y = nr.reduce_noise(y=y, sr=sr)
        y = signal.wiener(y)
        cache[(y, sr)] = (y, sr)
        logging.info("ПРЕДВАРИТЕЛЬНАЯ ОБРАБОТКА АУДИО ЗАВЕРШЕНА")
        return y, sr
    except Exception as e:
        logging.error(f"ОШИБКА ПРИ ПРЕДВАРИТЕЛЬНОЙ ОБРАБОТКЕ АУДИО: {e}")
        return None, None

async def extract_features(y, sr):
    """
    АСИНХРОННАЯ ФУНКЦИЯ ДЛЯ ИЗВЛЕЧЕНИЯ ПРИЗНАКОВ
    """
    try:
        logging.info("ИЗВЛЕЧЕНИЕ ПРИЗНАКОВ")
        y_tensor = torch.tensor(y, dtype=torch.float32).unsqueeze(0)
        mel_spec = MelSpectrogram()(y_tensor)
        logging.info("ИЗВЛЕЧЕНИЕ ПРИЗНАКОВ ЗАВЕРШЕНО")
        return mel_spec
    except Exception as e:
        logging.error(f"ОШИБКА ИЗВЛЕЧЕНИЯ ПРИЗНАКОВ: {e}")
        return None

def save_musicxml(score, output_path):
    """
    СОХРАНЕНИЕ ФАЙЛА MUSICXML
    """
    try:
        logging.info(f"СОХРАНЕНИЕ ФАЙЛА MUSICXML В {output_path}")
        tree = ET.ElementTree(score)
        tree.write(output_path, encoding='UTF-8', xml_declaration=True, pretty_print=True)
        logging.info(f"ФАЙЛ MUSICXML УСПЕШНО СОХРАНЕН: {output_path}")
    except Exception as e:
        logging.error(f"ОШИБКА ПРИ СОХРАНЕНИИ ФАЙЛА MUSICXML: {e}")

def analyze_audio(y, sr):
    """
    ПРИМЕР ФУНКЦИИ ДЛЯ АНАЛИЗА АУДИО
    """
    try:
        logging.info("АНАЛИЗ АУДИО")
        tempo, beat_frames = librosa.beat.beat_track(y=y, sr=sr)
        logging.info(f"АНАЛИЗ АУДИО ЗАВЕРШЕН: TEMPO={tempo}, BEAT_FRAMES={beat_frames}")
    except Exception as e:
        logging.error(f"ОШИБКА АНАЛИЗА АУДИО: {e}")

def generate_text_descriptions(y, sr):
    """
    ПРИМЕР ФУНКЦИИ ДЛЯ ГЕНЕРАЦИИ ТЕКСТОВЫХ ОПИСАНИЙ
    """
    try:
        logging.info("ГЕНЕРАЦИЯ ТЕКСТОВЫХ ОПИСАНИЙ")
        description = f"АУДИО ИМЕЕТ ЧАСТОТУ ДИСКРЕТИЗАЦИИ {sr} И ДЛИНУ {len(y)} ОБРАЗЦОВ."
        logging.info(f"ТЕКСТОВЫЕ ОПИСАНИЯ СГЕНЕРИРОВАНЫ: {description}")
    except Exception as e:
        logging.error(f"ОШИБКА ГЕНЕРАЦИИ ТЕКСТОВЫХ ОПИСАНИЙ: {e}")

def visualize_score(score, output_image_path):
    """
    ПРИМЕР ФУНКЦИИ ДЛЯ ВИЗУАЛИЗАЦИИ ПАРТИТУРЫ
    """
    try:
        logging.info("ВИЗУАЛИЗАЦИЯ ПАРТИТУРЫ")
        tree = ET.ElementTree(score)
        tree.write(output_image_path, encoding='UTF-8', xml_declaration=True, pretty_print=True)
        logging.info(f"ВИЗУАЛИЗАЦИЯ ПАРТИТУРЫ ЗАВЕРШЕНА: {output_image_path}")
    except Exception as e:
        logging.error(f"ОШИБКА ВИЗУАЛИЗАЦИИ ПАРТИТУРЫ: {e}")

def separate_sources(y, sr):
    """
    ПРИМЕР ФУНКЦИИ ДЛЯ СЕПАРАЦИИ ИСТОЧНИКОВ
    """
    try:
        logging.info("СЕПАРАЦИЯ ИСТОЧНИКОВ")
        vocals = y[:len(y)//2]
        accompaniment = y[len(y)//2:]
        logging.info("СЕПАРАЦИЯ ИСТОЧНИКОВ ЗАВЕРШЕНА")
    except Exception as e:
        logging.error(f"ОШИБКА ПРИ СЕПАРАЦИИ ИСТОЧНИКОВ: {e}")

def optimize_audio(y, sr):
    """
    ПРИМЕР ФУНКЦИИ ДЛЯ ОПТИМИЗАЦИИ АУДИО
    """
    try:
        logging.info("ОПТИМИЗАЦИЯ АУДИО")
        optimized_audio = nr.reduce_noise(y=y, sr=sr)
        logging.info("ОПТИМИЗАЦИЯ АУДИО ЗАВЕРШЕНА")
    except Exception as e:
        logging.error(f"ОШИБКА ОПТИМИЗАЦИИ АУДИО: {e}")

def plot_spectrogram(y, sr, output_image_path):
    """
    ВИЗУАЛИЗАЦИЯ СПЕКТРОГРАММЫ АУДИОФАЙЛА
    """
    try:
        import matplotlib.pyplot as plt
        import librosa.display

        D = librosa.amplitude_to_db(np.abs(librosa.stft(y)), ref=np.max)
        plt.figure(figsize=(12, 8))
        librosa.display.specshow(D, sr=sr, x_axis='time', y_axis='log')
        plt.colorbar(format='%+2.0f dB')
        plt.title('Spectrogram')
        plt.savefig(output_image_path)
        logging.info(f'СПЕКТРОГРАММА СОХРАНЕНА В {output_image_path}')
        plt.close()
    except Exception as e:
        logging.error(f"ОШИБКА ВИЗУАЛИЗАЦИИ СПЕКТРОГРАММЫ: {e}")

def plot_chromagram(y, sr, output_image_path):
    """
    ВИЗУАЛИЗАЦИЯ ХРОМОГРАММЫ АУДИОФАЙЛА
    """
    try:
        import matplotlib.pyplot as plt
        import librosa.display

        chromagram = librosa.feature.chroma_stft(y=y, sr=sr)
        plt.figure(figsize=(12, 8))
        librosa.display.specshow(chromagram, sr=sr, x_axis='time', y_axis='chroma')
        plt.colorbar(format='%+2.0')
        plt.title('Chromagram')
        plt.savefig(output_image_path)
        logging.info(f'ХРОМОГРАММА СОХРАНЕНА В {output_image_path}')
        plt.close()
    except Exception as e:
        logging.error(f"ОШИБКА ВИЗУАЛИЗАЦИИ ХРОМОГРАММЫ: {e}")

def harmonic_rhythm_analysis(y, sr):
    """
    АНАЛИЗ ГАРМОНИЧЕСКОЙ РИТМИКИ
    """
    try:
        logging.info("Starting harmonic rhythm analysis")

        # Убедитесь, что входные данные представляют собой действительные числа
        if np.iscomplexobj(y):
            y = np.real(y)
        logging.info(f"Using real part of the audio signal. Shape: {y.shape}")

        # Извлечение темпа и ударов
        logging.info("Extracting tempo and beat frames")
        y_temp = np.abs(y)  # Преобразуем значения в положительные реальные числа
        tempo, beat_frames = librosa.beat.beat_track(y=y_temp, sr=sr)
        logging.info(f"Tempo: {tempo}, Beat frames: {beat_frames}")

        # Извлечение гармоник
        logging.info("Extracting harmonic components")
        chords = librosa.effects.harmonic(y)
        logging.info(f"Harmonic shape: {chords.shape}")

        if np.iscomplexobj(chords):
            chords = np.real(chords)
        logging.info(f"Using real part of harmonic signal. Shape: {chords.shape}")

        # АНАЛИЗ РИТМА С ИСПОЛЬЗОВАНИЕМ MADMOM
        logging.info("Analyzing rhythm with Madmom")
        proc = madmom.features.beats.RNNBeatProcessor()(y)
        
        # Прверка на None и логгирование результата
        if proc is None:
            logging.error("Processed beat data returned None.")
            raise ValueError("Processed beat data is None.")
        
        logging.info(f"Processed beat data: {proc}")
        
        # Выполнение beat tracking
        try:
            beat_times = madmom.features.beats.DBNBeatTrackingProcessor(beats_per_bar=[3, 4])(proc)
            logging.info(f"Beat times: {beat_times}")

            return tempo, beat_frames, chords, beat_times
        except Exception as beat_tracking_error:
            logging.error(f"Ошибка при выполнении beat tracking: {beat_tracking_error}")
            return tempo, beat_frames, chords, None

    except Exception as e:
        logging.error(f"ОШИБКА АНАЛИЗА ГАРМОНИЧЕСКОЙ РИТМИКИ: {e}")
        return None, None, None, None

# Тестовый вызов функции для отладки
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    y, sr = librosa.load(r'C:\Users\ARTYOM\Documents\Python\Bring_it_up.mp3', sr=None)
    harmonic_rhythm_analysis(y, sr)

async def main(file_path, output_path, output_image_path):
    try:
        # Проверка существования директории
        output_dir = os.path.dirname(output_path)
        if not os.path.exists(output_dir):
            logging.error(f"Директория вывода '{output_dir}’ не существует.")
            sys.exit(1)

        # Проверка существования файла
        if os.path.exists(output_path):
            logging.info(f"Файл MusicXML '{output_path}’ будет перезаписан.")
        else:
            logging.info(f"Файл MusicXML будет сохранен в '{output_path}’.")

        # Предварительная обработка аудио
        y, sr = await preprocess_audio(file_path)
        if y is None or sr is None:
            logging.error("Ошибка на шаге предварительной обработки аудио")
            return

        # Извлечение признаков
        features = await extract_features(y, sr)
        if features is None:
            logging.error("Ошибка на шаге извлечения признаков")
            return

        # Анализ гармонической ритмики
        tempo, beat_frames, chords, beat_times = harmonic_rhythm_analysis(y, sr)
        if tempo is None or beat_frames is None or chords is None or beat_times is None:
            logging.error("Ошибка на шаге анализа гармонической ритмики")
            return

        # Создание простого XML-документа
        score = ET.Element("score-partwise", version="3.1")
        part = ET.SubElement(score, "part", id="P1")
        measure = ET.SubElement(part, "measure", number="1")
        note_element = ET.SubElement(measure, "note")
        pitch_element = ET.SubElement(note_element, "pitch")
        ET.SubElement(pitch_element, "step").text = "C"
        ET.SubElement(pitch_element, "octave").text = "4"
        ET.SubElement(note_element, "duration").text = "1"
        ET.SubElement(note_element, "type").text = "quarter"

        # Сохранение файла
        save_musicxml(score, output_path)

        # Дополнительные шаги
        logging.info("Дополнительные шаги выполнены")

        # Анализ аудио
        analyze_audio(y, sr)

        # Генерация текстовых описаний
        generate_text_descriptions(y, sr)

        # Визуализация партитуры
        visualize_score(score, output_image_path)

        # Сепарация источников
        separate_sources(y, sr)

        # Оптимизация аудио
        optimize_audio(y, sr)

        # Визуализация спектрограммы
        plot_spectrogram(y, sr, output_image_path.replace(".png", "_spectrogram.png"))

        # Визуализация хромограммы
        plot_chromagram(y, sr, output_image_path.replace(".png", "_chromagram.png"))
    except Exception as e:
        logging.error(f"Ошибка в основной функции: {e}")
        sys.exit(1)

if __name__ == "__main__":
    file_path = r'C:\Users\ARTYOM\Documents\Python\Bring_it_up.mp3'
    output_path = r'C:\Users\ARTYOM\Documents\Python\Bring_it_up.musicxml'
    output_image_path = r'C:\Users\ARTYOM\Documents\Python\Bring_it_up.png'

    from asyncio import run
    run(main(file_path, output_path, output_image_path))