from pydub import AudioSegment, generators
import numpy as np
import os

# Output folder
os.makedirs("output_tracks", exist_ok=True)

def generate_battle_track(filename, crazy=False, boosted=False):
    # Generate pink noise
    duration_ms = 65000  # 65 seconds
    noise = generators.WhiteNoise().to_audio_segment(duration=duration_ms).apply_gain(-20)

    # Apply pink noise shaping manually
    samples = np.array(noise.get_array_of_samples())
    samples = samples.astype(np.float32)
    samples = samples / np.max(np.abs(samples))  # Normalize

    # Fake pink noise shaping (more low end)
    pink = samples * np.linspace(1.0, 0.3, num=len(samples))

    pink = pink * 32767
    pink = pink.astype(np.int16)
    noise = noise._spawn(pink.tobytes())

    # Add a bass tone at 63 Hz
    tone = generators.Sine(63).to_audio_segment(duration=duration_ms).apply_gain(-10)
    combined = noise.overlay(tone)

    if crazy:
        # Add tremolo effect: simple volume modulation
        mod_samples = np.array(combined.get_array_of_samples())
        t = np.linspace(0, duration_ms/1000, num=len(mod_samples))
        mod = 0.5 * (1.0 + np.sin(2 * np.pi * 5 * t))  # 5 Hz modulation
        modulated = (mod_samples * mod).astype(np.int16)
        combined = combined._spawn(modulated.tobytes())

    if boosted:
        combined = combined.apply_gain(+3)  # Boost 3 dB

    combined.export(f"output_tracks/{filename}.wav", format="wav")
    combined.export(f"output_tracks/{filename}.mp3", format="mp3", bitrate="320k")

# Generate clean, crazy, boosted tracks
generate_battle_track("F2F_Battle_Track_Clean")
generate_battle_track("F2F_Battle_Track_Crazy", crazy=True)
generate_battle_track("F2F_Battle_Track_Boosted", boosted=True)

# For Je Suis Malade simulation
# Just create a soft tone + broadband
voice_sim = generators.Sine(440).to_audio_segment(duration=65000).apply_gain(-25)  # 440 Hz "voice tone"
base = generators.WhiteNoise().to_audio_segment(duration=65000).apply_gain(-30)

combined_voice = base.overlay(voice_sim)
bassline = generators.Sine(63).to_audio_segment(duration=65000).apply_gain(-12)
final_mix = combined_voice.overlay(bassline)

final_mix.export("output_tracks/Je_Suis_Malade_Battle_Remix.wav", format="wav")
final_mix.export("output_tracks/Je_Suis_Malade_Battle_Remix.mp3", format="mp3", bitrate="320k")

print("All battle tracks successfully created in /output_tracks!")