#######################################################################
##    created by cooyou.org
#######################################################################
import os
import sys
import numpy as np
import librosa
import soundfile as sf
from pydub import AudioSegment
import noisereduce as nr

def normalize_audio(y, peak=1.5):
    """音声を正規化して最大振幅をpeakにする"""
    max_amp = np.max(np.abs(y))
    if max_amp == 0:
        return y
    return y / max_amp * peak

def mp3_to_wav(mp3_path, wav_path):
    """MP3をWAVに変換"""
    audio = AudioSegment.from_mp3(mp3_path)
    audio.export(wav_path, format="wav")

def wav_to_mp3(wav_path, mp3_path):
    """WAVをMP3に変換"""
    audio = AudioSegment.from_wav(wav_path)
    audio.export(mp3_path, format="mp3")

def denoise_wav(input_wav, output_clean_wav, output_noise_wav=None, prop_decrease=0.5, peak=1.5):
    """WAVファイルをノイズ除去して保存、希望でノイズだけのファイルも作成"""
    y, sr = librosa.load(input_wav, sr=None)
    reduced = nr.reduce_noise(y=y, sr=sr, prop_decrease=prop_decrease)
    reduced = normalize_audio(reduced, peak=peak)

    # クリーン音声保存
    sf.write(output_clean_wav, reduced, sr)

    # ノイズだけのファイルを作成
    if output_noise_wav is not None:
        noise_only = y - reduced
        noise_only = normalize_audio(noise_only, peak=peak)
        sf.write(output_noise_wav, noise_only, sr)

def main(input_mp3, prop_decrease=0.5, peak=1.5):
    base_dir = os.path.dirname(input_mp3)
    filename = os.path.splitext(os.path.basename(input_mp3))[0]
    temp_wav = os.path.join(base_dir, f"{filename}_temp.wav")
    temp_clean_wav = os.path.join(base_dir, f"{filename}_temp_clean.wav")
    output_clean_mp3 = os.path.join(base_dir, f"{filename}_clean.mp3")
    output_noise_mp3 = os.path.join(base_dir, f"{filename}_noise.mp3")  # ノイズ成分

    print("MP3 → WAV変換中...")
    mp3_to_wav(input_mp3, temp_wav)

    print("ノイズ除去中...")
    denoise_wav(temp_wav, temp_clean_wav, output_noise_wav=temp_wav.replace(".wav", "_noise.wav"),
                prop_decrease=prop_decrease, peak=peak)

    print("WAV → MP3変換中...")
    wav_to_mp3(temp_clean_wav, output_clean_mp3)
    wav_to_mp3(temp_wav.replace(".wav", "_noise.wav"), output_noise_mp3)

    # 一時ファイルの削除
    os.remove(temp_wav)
    os.remove(temp_clean_wav)
    os.remove(temp_wav.replace(".wav", "_noise.wav"))

    print(f"処理完了:")
    print(f"  クリーン音声: {output_clean_mp3}")
    print(f"  ノイズ音声: {output_noise_mp3}")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("使い方: python denoise.py 入力ファイル.mp3 [prop_decrease] [peak]")
        print("  prop_decrease: 0.0〜1.0 (ノイズ除去の強さ, 小さいほど声を残す, default=0.5)")
        print("  peak: 音量正規化の最大値 (default=1.5)")
        sys.exit(1)

    input_mp3 = sys.argv[1]
    prop_decrease = float(sys.argv[2]) if len(sys.argv) > 2 else 0.5
    peak = float(sys.argv[3]) if len(sys.argv) > 3 else 1.5

    main(input_mp3, prop_decrease=prop_decrease, peak=peak)