import os
import logging
import openai
import sys
import asyncio
import numpy as np
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:
        tempo, beat_frames = librosa.beat.beat_track(y=y, sr=sr)
        chords = librosa.effects.harmonic(y)
        
        # АНАЛИЗ РИТМА С ИСПОЛЬЗОВАНИЕМ MADMOM
        proc = madmom.features.beats.RNNBeatProcessor()(y)
        beat_times = madmom.features.beats.DBNBeatTrackingProcessor(beats_per_bar=[3, 4])(proc)
        
        return tempo, beat_frames, chords, beat_times
    except Exception as e:
        logging.error(f"ОШИБКА АНАЛИЗА ГАРМОНИЧЕСКОЙ РИТМИКИ: {e}", exc_info=True)
        return None, None, None, None

async def main(file_path, output_path, output_image_path):
    """
    ОСНОВНАЯ ФУНКЦИЯ ДЛЯ ОБРАБОТКИ АУДИО И СОЗДАНИЯ MUSICXML
    """
    try:
        # ПРОВЕРЬТЕ СУЩЕСТВОВАНИЕ ДИРЕКТОРИИ ДЛЯ СОХРАНЕНИЯ ФАЙЛА
        output_dir = os.path.dirname(output_path)
        if not os.path.exists(output_dir):
            logging.error(f"ДИРЕКТОРИЯ ВЫВОДА '{output_dir}' НЕ СУЩЕСТВУЕТ. ФАЙЛ MUSICXML НЕ БУДЕТ СОХРАНЕН.")
            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("ДОПОЛНИТЕЛЬНЫЕ ШАГИ ВЫПОЛНЕНЫ")

        # 1. АНАЛИЗ АУДИО
        logging.info("НАЧАЛО АНАЛИЗА АУДИО…")
        analyze_audio(y, sr)
        logging.info("АНАЛИЗ АУДИО ЗАВЕРШЕН")

        # 2. ГЕНЕРАЦИЯ ТЕКСТОВЫХ ОПИСАНИЙ
        logging.info("ГЕНЕРАЦИЯ ТЕКСТОВЫХ ОПИСАНИЙ…")
        generate_text_descriptions(y, sr)
        logging.info("ТЕКСТОВЫЕ ОПИСАНИЯ СОЗДАНЫ")

        # 3. ВИЗУАЛИЗАЦИЯ ПАРТИТУРЫ
        logging.info("ВИЗУАЛИЗАЦИЯ ПАРТИТУРЫ…")
        visualize_score(score, output_image_path)
        logging.info("ВИЗУАЛИЗАЦИЯ ПАРТИТУРЫ ЗАВЕРШЕНА")

        # 4. СЕПАРАЦИЯ ИСТОЧНИКОВ
        logging.info("НАЧАЛО СЕПАРАЦИИ ИСТОЧНИКОВ…")
        separate_sources(y, sr)
        logging.info("СЕПАРАЦИЯ ИСТОЧНИКОВ ЗАВЕРШЕНА")

        # 5. ОПТИМИЗАЦИЯ АУДИО
        logging.info("ОПТИМИЗАЦИЯ АУДИО…")
        optimize_audio(y, sr)
        logging.info("ОПТИМИЗАЦИЯ АУДИО ЗАВЕРШЕНА")

        # 6. ВИЗУАЛИЗАЦИЯ СПЕКТРОГРАММЫ
        plot_spectrogram(y, sr, output_image_path.replace(".png", "_spectrogram.png"))

        # 7. ВИЗУАЛИЗАЦИЯ ХРОМОГРАММЫ
        plot_chromagram(y, sr, output_image_path.replace(".png", "_chromagram.png"))
    except Exception as e:
        logging.error(f"ОШИБКА В ОСНОВНОЙ ФУНКЦИИ: {e}", exc_info=True)
        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))