import os
import numpy as np
from pretty_midi import PrettyMIDI, Instrument, Note
import matplotlib.pyplot as plt
from decimal import Decimal, ROUND_HALF_UP


def precise_round(value, decimals=6):
    """Более точное округление для расчетов времени."""
    return float(Decimal(value).quantize(Decimal(f'1e-{decimals}'), rounding=ROUND_HALF_UP))


def normalize_tempo(midi_file, resolution=0.125):
    """
    Нормализует темп MIDI-файла до целевых временных шагов без дробных значений.
    """
    midi_data = PrettyMIDI(midi_file)
    original_bpm = midi_data.estimate_tempo()
    target_bpm = 60 / resolution  # Рассчитать целевой BPM для заданного разрешения
    time_scaling = target_bpm / original_bpm

    print(f"Оригинальный BPM: {original_bpm:.2f}, Целевой BPM: {target_bpm:.2f}")

    # Масштабируем время событий
    for instrument in midi_data.instruments:
        for note in instrument.notes:
            note.start = precise_round(note.start * time_scaling)
            note.end = precise_round(note.end * time_scaling)
        for control_change in instrument.control_changes:
            control_change.time = precise_round(control_change.time * time_scaling)

    # Масштабируем временные отметки
    midi_data.adjust_times(
        [0, midi_data.get_end_time()],
        [0, precise_round(midi_data.get_end_time() * time_scaling)]
    )

    return midi_data, original_bpm


def midi_to_svd(midi_data, resolution=0.125, num_notes=128):
    """Преобразование MIDI в SVD."""
    max_time = midi_data.get_end_time()
    num_steps = int(np.ceil(max_time / resolution))

    # Изменяем тип данных D на uint16
    S = np.zeros((num_notes, num_steps), dtype=np.uint8)
    D = np.zeros((num_notes, num_steps), dtype=np.uint16)  # Изменено на uint16
    V = np.zeros((num_notes, num_steps), dtype=np.uint8)

    for instrument in midi_data.instruments:
        if instrument.is_drum:
            continue
        for note in instrument.notes:
            pitch = note.pitch
            velocity = note.velocity
            start_idx = int(np.round(note.start / resolution))
            end_idx = int(np.round(note.end / resolution))

            if end_idx <= start_idx:
                end_idx = start_idx + 1

            if start_idx < num_steps:
                S[pitch, start_idx] = 1
                D[pitch, start_idx] = min(end_idx - start_idx, num_steps - start_idx)  # Теперь поддерживает большие значения
                V[pitch, start_idx] = velocity

    return S, D, V


def svd_to_midi(S_dense, D_dense, V_dense, resolution=0.125, program=0):
    """Преобразование SVD обратно в MIDI."""
    midi = PrettyMIDI()
    instrument = Instrument(program=program)

    for pitch in range(S_dense.shape[0]):
        for time in range(S_dense.shape[1]):
            if S_dense[pitch, time] == 1:
                start_time = time * resolution
                duration = D_dense[pitch, time] * resolution
                velocity = int(V_dense[pitch, time])
                end_time = start_time + duration
                note = Note(velocity=velocity, pitch=pitch, start=start_time, end=end_time)
                instrument.notes.append(note)

    midi.instruments.append(instrument)
    return midi


def revert_tempo(midi_data, original_bpm, resolution=0.125):
    """
    Возвращает MIDI обратно к исходному темпу после нормализации.
    :param midi_data: Объект PrettyMIDI
    :param original_bpm: Исходный BPM
    :param resolution: Шаг времени (в секундах)
    :return: PrettyMIDI объект с восстановленным темпом
    """
    target_bpm = 60 / resolution
    time_scaling = target_bpm / original_bpm

    print(f"Возврат к исходному BPM: {original_bpm:.2f} из нормализованного {target_bpm:.2f}")

    # Корректируем времена событий
    for instrument in midi_data.instruments:
        for note in instrument.notes:
            note.start *= time_scaling
            note.end *= time_scaling
        for control_change in instrument.control_changes:
            control_change.time *= time_scaling

    # Корректируем временные отметки
    original_times = [0, midi_data.get_end_time()]
    new_times = [0, midi_data.get_end_time() * time_scaling]
    midi_data.adjust_times(original_times, new_times)

    return midi_data


def visualize_midi_piano_roll(midi, resolution=0.125):
    """Визуализация Piano Roll."""
    max_time = midi.get_end_time()
    num_steps = int(np.ceil(max_time / resolution))
    piano_roll = np.zeros((128, num_steps))

    time_color_map = {}
    color_id = 1

    for instrument in midi.instruments:
        for note in instrument.notes:
            start_idx = int(note.start / resolution)
            end_idx = int(note.end / resolution)

            if start_idx not in time_color_map:
                time_color_map[start_idx] = color_id
                color_id += 1

            piano_roll[note.pitch, start_idx:end_idx] = time_color_map[start_idx]

    cmap = plt.get_cmap("tab20", len(time_color_map) + 1)

    plt.figure(figsize=(12, 6))
    plt.title("Piano Roll (Colored by Start Time)")
    plt.imshow(piano_roll, aspect="auto", cmap=cmap, origin="lower")
    plt.colorbar(label="Start Time Index")
    plt.xlabel("Time Steps")
    plt.ylabel("Pitch")
    plt.show()


if __name__ == "__main__":
    midi_file = "1.mid"
    resolution = 0.0625

    # Нормализация темпа
    midi_data, original_bpm = normalize_tempo(midi_file, resolution=resolution)

    # Преобразование в SVD
    S_dense, D_dense, V_dense = midi_to_svd(midi_data, resolution=resolution)

    # Преобразование обратно в MIDI
    midi_reconstructed = svd_to_midi(S_dense, D_dense, V_dense, resolution=resolution)

    # Возврат к исходному темпу
    midi_reconstructed = revert_tempo(midi_reconstructed, original_bpm, resolution=resolution)

    # Сохранение реконструированного MIDI
    output_dir = os.path.dirname(midi_file)
    output_file = os.path.join(output_dir, "reconstructed_" + os.path.basename(midi_file))
    midi_reconstructed.write(output_file)
    print(f"Реконструированный MIDI сохранен как: {output_file}")

    # Визуализация Piano Roll
    visualize_midi_piano_roll(midi_reconstructed, resolution=resolution)

Ориентируйся на этот код во всем, он работает в данный момент лучше всего. Но в нем все еще наблюдается мизерная погрешность. Я думаю точностью вычислений тут не обойтись и подход неверный, так как time steps больше 70000 получается. Ты как будто решаешь не причину, а следствие. Причина в неверном подборе верного bpm, в связи с чем копится ошибка округления. Да и округление в целом можно поставить под сомнение в плане правильности подхода. Мелодия должна идеально к сетке подходить, как в fl studio и других daw. Там не возникает проблем с оригинальной midi и таких погрешностей тоже нет, хотя все работает очень быстро и вряд ли считается с супер высокой точностью какой то, так как даже она на длинной дистанции вызовет ошибку накопления. Исправь подход и код. Дай исправленный файл полностью