from scipy.io import wavfile as wav
import numpy as np
import sys

print("Portfolio Object 2: Adaptive Tone Control")

# Contants
bands = {"low": (0, 300), "mid": (300, 2000), "high": (2000, 20000)}

file = "test.wav"
if len(sys.argv) > 1:
    file = sys.argv[1]
print("Input file: ", file)
# Init
sample_rate, wav_signal = wav.read(file)
# window_size = 1024
window_size = len(wav_signal)
fft_values = np.fft.fft(wav_signal[:window_size])
fft_freqs = np.fft.fftfreq(window_size, 1 / sample_rate)


def band_energy(band, fft_values):
    idx_band = np.where((fft_freqs >= band[0]) & (fft_freqs <= band[1]))[0]
    return np.sum(np.abs(fft_values[idx_band]) ** 2)


energy_low = band_energy(bands["low"], fft_values)
energy_mid = band_energy(bands["mid"], fft_values)
energy_high = band_energy(bands["high"], fft_values)
avg_energy = (energy_low + energy_mid + energy_high) / 3


print(f"low {energy_low:.2e}")
print(f"mid {energy_mid:.2e}")
print(f"high {energy_high:.2e}")
print(f"avg {avg_energy:.2e}")


def adjust(target_energy, current_energy, fft_values, band):
    idx_band = np.where((fft_freqs >= band[0]) & (fft_freqs <= band[1]))[0]
    gain = np.sqrt(target_energy / (current_energy + 1e-6))
    adjusted_fft_values = np.copy(fft_values)
    adjusted_fft_values[idx_band] *= gain
    return adjusted_fft_values


adjusted_low = adjust(avg_energy, energy_low, fft_values, bands["low"])
adjusted_mid = adjust(avg_energy, energy_mid, fft_values, bands["mid"])
adjusted_high = adjust(avg_energy, energy_high, fft_values, bands["high"])

print(f'adj low {band_energy( bands["low"],adjusted_low):.2e}')
print(f'adj mid {band_energy(bands["mid"],adjusted_mid):.2e}')
print(f'adj high {band_energy( bands["high"],adjusted_high):.2e}')

adjusted_all = adjusted_low + adjusted_mid + adjusted_high
full_spectrum = np.concatenate([adjusted_all])
reconstructed_signal = np.fft.ifft(adjusted_all).real
reconstructed_signal = np.int16(
    reconstructed_signal / np.max(np.abs(reconstructed_signal)) * np.iinfo(np.int16).max
)

output_file = "adj-" + file
wav.write(output_file, sample_rate, reconstructed_signal)
print(f"Adjusted audio written to {output_file}")