import gradio as gr
import torch
from transformers import WhisperForConditionalGeneration, WhisperProcessor
from pydub import AudioSegment
import re

# 업로드한 모델 로드
repo_name = "ireneminhee/speech-to-depression"
model = WhisperForConditionalGeneration.from_pretrained(repo_name)
processor = WhisperProcessor.from_pretrained(repo_name)

# 음성을 텍스트로 변환하는 함수
def transcribe(audio):
    inputs = processor(audio, return_tensors="pt", sampling_rate=16000)
    generated_ids = model.generate(inputs.input_features)
    transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return transcription

# 우울증 예측 모델 로드
def load_model_from_safetensors(model_name, safetensors_path):
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(model_name, config=model_name)
    state_dict = torch.load(safetensors_path)  # safetensors를 모델로 로드
    model.load_state_dict(state_dict)
    model.eval()
    return model, tokenizer

# 예측 함수
def predict_depression(sentences, model, tokenizer):
    results = []
    for sentence in sentences:
        inputs = tokenizer(sentence, return_tensors="pt", truncation=True, padding=True)
        with torch.no_grad():
            outputs = model(**inputs)
            logits = outputs.logits
            prediction = torch.argmax(logits, dim=-1).item()
            results.append((sentence, prediction))
    return results

# 전체 프로세스를 실행하는 함수
def process_audio_and_predict(audio):
    # 1. Whisper 모델로 음성을 텍스트로 변환
    text = transcribe_audio(audio)

    # 2. 텍스트를 문장 단위로 나누기
    #sentences = split_sentences_using_gpt(text)

    # 3. 모델 로드 (모델 경로에 맞게 수정)
    # 모델과 tokenizer 경로 (사용자 환경에 맞게 설정)
    safetensors_path = "./model/model.safetensors"  # SafeTensors 모델 파일 경로
    model_name = "klue/bert-base"  # 모델 이름 또는 Hugging Face 경로
    model, tokenizer = load_model_from_safetensors(model_name, safetensors_path)

    # 4. 문장별로 우울 증상 예측
    results = predict_depression(text, model, tokenizer)

    # 5. 결과를 반환
    df_result = pd.DataFrame(results, columns=["Sentence", "Depression_Prediction"])
    average_probability = df_result["Depression_Prediction"].mean()

    return f"Average Depression Probability: {average_probability:.2f}"

# Gradio 인터페이스로 연결할 함수
def gradio_process_audio(audio_data):
    # 사용자가 마이크로 입력한 음성을 임시 파일로 저장
    temp_audio_path = "temp_audio.wav"
    with open(temp_audio_path, "wb") as f:
        f.write(audio_data)

    # 오디오 처리 및 예측
    average_probability, df_result = process_audio_and_detect_depression(temp_audio_path, safetensors_path, model_name)

    # 결과 출력
    return f"Average Depression Probability: {average_probability:.2f}", df_result

# Gradio 인터페이스 정의
interface = gr.Interface(
    fn=gradio_process_audio,           # Gradio에서 호출할 함수
    inputs=gr.Audio(type="numpy"),  # 사용자 음성 입력 (마이크)
    outputs=[
        gr.Textbox(label="Depression Probability"),  # 평균 확률
        gr.Dataframe(label="Sentence-wise Analysis")  # 상세 분석 결과
    ],
    title="Depression Detection from Audio",
    description="Record your voice, and the model will analyze the text for depression likelihood."
)

# Gradio 실행
interface.launch(share=True)