import numpy as np
import random
import time
from tinygrad import Tensor, nn
from tinygrad.nn.state import safe_load, load_state_dict
import librosa
import sounddevice as sd
from model import gen
from data import spec_to_audio

SAMPLE_RATE = 22050

def load_model(filepath="model.safetensors"):
    """Loads the model structure and weights."""
    model = gen()
    state_dict = safe_load(filepath)
    load_state_dict(model, state_dict)
    return model

def load_data(filepath="data.npz"):
    """Loads the pre-processed spectrogram data."""
    print(f"Loading data from {filepath}...")
    data = np.load(filepath)
    x = data["arr_0"]
    return x

def play_spec(spec,i):
    """Converts a spectrogram numpy array to audio and plays it."""
    audio = spec_to_audio(spec)
    sd.wait()
    print(f"chunk:{i}")
    sd.play(audio, samplerate=SAMPLE_RATE)

def run_prediction_loop(model, data_x):
    current_spect = data_x[0:1]
    for i in range(10):
        play_spec(current_spect[0][0],i)
        current_spect = model(Tensor(current_spect)).numpy()

if __name__ == "__main__":
    model = load_model()
    data_x = load_data()
    run_prediction_loop(model, data_x)