import multiprocessing as mp
from functools import partial
import os
import numpy as np
import torch
import json
import pandas as pd
from tqdm import tqdm
from suno_utils.audio import Audio
import matplotlib.pyplot as plt

def trim_sample_audio(path):
    audio = Audio.from_file(path)
    sample_rate = audio.sample_rate
    wav = torch.from_numpy(audio.array_float)
    sample_length = len(wav) / sample_rate
    start_time = np.random.randint(0, int(sample_length)-2)
    trimmed_wav = wav[start_time*sample_rate:start_time*sample_rate + sample_rate]
    trimmed_audio = Audio.from_array_float(trimmed_wav, sample_rate)
    out_path = os.path.join(output_path, os.path.split(path)[1])
    trimmed_audio.write_wav(out_path)
    return start_time, out_path

def process_sample(sample_path):
    return trim_sample_audio(sample_path)

def read_json(filepath):
    full_song_json_data = []
    with open(filepath, 'rb') as f:
        for line in tqdm(f):
            json_line = line.decode('utf-8').strip()
            try:
                full_song_json_data.append({**json.loads(json_line)})
            except json.JSONDecodeError as e:
                print(f'Error parsing JSON in line: {json_line}, Error: {str(e)}')
    return full_song_json_data

if __name__ == '__main__':
    train_samps_df = pd.DataFrame(read_json('suno_seal_tr.jsonl'))
    num_cores = mp.cpu_count()  # Get the number of CPU cores
    total_samples = len(train_samps_df['path'])

    output_path = '/app/suno/christian_c/suno_seal_one_sec_48kHz_data'
    
    # Create a pool of worker processes
    with mp.Pool(processes=num_cores) as pool:
        # Use pool.map to apply the function to all sample paths in parallel
        results = list(tqdm(pool.imap(process_sample, train_samps_df['path']), 
                            total=total_samples, 
                            desc="Trimming audio samples"))
    
    # Unpack the results
    start_times, out_paths = zip(*results)
    trimmed_audio_df = pd.DataFrame({'path': out_paths, 'start_time': start_times})
    trimmed_audio_df.to_csv('suno_seal_one_sec_tr.csv')