import torch
import torchaudio
import jiwer
from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC

# Load the Arabic-specific processor and model
model_name = "omarxadel/wav2vec2-large-xlsr-53-arabic-egyptian"
processor = Wav2Vec2Processor.from_pretrained(model_name)
model = Wav2Vec2ForCTC.from_pretrained(model_name)

# Set the model to evaluation mode
model.eval()

# Define the chunk length in seconds
CHUNK_LENGTH = 30  # 30 seconds
SAMPLE_RATE = 16000  # 16 kHz

def transcribe_chunk(chunk):
    """
    Transcribes a single audio chunk.
    """
    # Process the audio chunk with the processor
    inputs = processor(chunk, sampling_rate=SAMPLE_RATE, return_tensors="pt", padding=True)
    
    # Get logits from the model
    with torch.no_grad():
        logits = model(inputs.input_values).logits
    
    # Decode the predicted ids to text
    predicted_ids = torch.argmax(logits, dim=-1)
    transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
    
    return transcription.strip()

def transcribe_audio_in_chunks(audio_file):
    """
    Loads an audio file, splits it into chunks, transcribes each chunk, and combines the transcriptions.
    """
    # Load the audio file
    print("Loading audio file...")
    waveform, sr = torchaudio.load(audio_file)
    print(f"Original Sample Rate: {sr}, Waveform shape: {waveform.shape}")
    
    # Convert to mono if it's stereo
    if waveform.shape[0] > 1:
        waveform = torch.mean(waveform, dim=0, keepdim=True)
        print(f"Converted to mono. New shape: {waveform.shape}")
    
    # Resample if needed
    if sr != SAMPLE_RATE:
        print(f"Resampling from {sr} to {SAMPLE_RATE}...")
        resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=SAMPLE_RATE)
        waveform = resampler(waveform)
        print(f"Resampled waveform shape: {waveform.shape}")
    
    # Calculate the number of samples per chunk
    chunk_size = CHUNK_LENGTH * SAMPLE_RATE
    total_samples = waveform.shape[1]
    num_chunks = (total_samples + chunk_size - 1) // chunk_size  # Ceiling division
    print(f"Total samples: {total_samples}, Chunk size (in samples): {chunk_size}, Number of chunks: {num_chunks}")
    
    transcriptions = []
    
    for i in range(num_chunks):
        start_sample = i * chunk_size
        end_sample = min(start_sample + chunk_size, total_samples)
        chunk = waveform[:, start_sample:end_sample].squeeze(0)
        print(f"Processing chunk {i+1}/{num_chunks}, Samples {start_sample}:{end_sample}")
        
        # Normalize the audio
        chunk = chunk / torch.abs(chunk).max()
        
        transcription = transcribe_chunk(chunk.numpy())
        transcriptions.append(transcription)
        print(f"Transcription for chunk {i+1}: {transcription}\n")
    
    # Combine all transcriptions
    full_transcription = ' '.join(transcriptions)
    return full_transcription.strip()

# Define reference transcription for WER calculation
reference_transcription = "قال المصنف رحمه الله تعالى ونفعنا الله بعلومه في الدارين امين فصل والربا في الذهب والفضه والمطعومات الربا شكل من اشكال البيوع فيه تبادل مثل الثمن والمثمن فهنا فيه مال بازاء مال ولذلك اختلط الامر على المشركين في مكه فقالوا انما البيع مثل الربا يعني هذا فيه مبادله وهذا فيه مبادله وفي المبادله زياده بين الثمن و المثمن وهنا زياده بين الثمن والمثمن فما الفرق نحن نبيع ونشتري والفرق هو ان الله سبحانه وتعالى امر بحرمة الربا من اجل الحفاظ على نظام الاقتصاد النقدي من داخله فلو اباح الله الربا لكان هناك زياده في وسيط التبادل دون زيادة الانتاج ويحدث بذلك التضخم من داخل النظام فاذا حدث التضخم من ذات النظام ف الذي يضار بذلك محدود الدخل الذي يضار بذلك الفقير فحرم الله الربا من اجل الا يقع ظلم على الفقير واوجب الزكاه حتى يتعادل دخل الفقير لاداء مقتضيات الحياه فتحريم الربا مرتبط في النظر الالهي مع ايجاب الزكاه حرم الله الربا واوجب حرم واوجب الزكاه من اجل حماية انسانية الفقير محدود الدخل فلا يجد الاسعار فوق الطاقه ويستطيع ان يحصل باستهلاكه لما اخذه من الزكاه ما يريد من مقتضيات الحياه فاذا كان الامر كذلك وهذا نظام رباني الهي كان حرا لا يستطيع احد ان يتحكم فيه لا في اجتماع ولا في سياسه ولا في اقتصاد يكمل المنظومه الصلاه فاذا ذهبها الناس وقفوا متجاورين الغني بجواره الفقير والفقير بجواره المسكين والمسكين بجواره الغني لا فرق بينهما والصف الاول مناخ من سبق سبق الى الصف الاول الغني او سبق الى الصف الثاني او سبق الى الصف الاول الفقير لا يستطيع احد ان يجذبه من اوليته حتى يتاخر بسبب غناه او جاهه او سلطانه لان من جاء فوجد الصلاه فليلحق بها حيث ما انتهى به الصف وكذلك في طلب العلم حيث ما انتهى به المجلس من اين تتاتى هذه النفسيه الابيه من دون كبر انما بعزه ورقي وانسانيه بوجوب الزكاه فالزكاة ليست منه ولا عطيه ولا هي يد عليا و بحرمة الربا حتي يحافظ على النظام الاقتصادي من داخله من التضخم ذهبت الايام وجاءت الايام وتسلط علينا ابالسه الجن والانس واخرجوا الذهب والفضه من التعامل واستبدلوا بها حيله سخيفه وهي حيله الورق النقدي جعلوا له اولا غطاء ثم بعد ذلك تناسوا الغطاء فاصبح ورقا لا قيمه له في انتاجه مع قيمته دينار الذهب 4 جرام وربع لو سيحناه يطلع 4 جرام وربع ذهب لو صكيناه يطلع 4 جرام وربع ذهب لو بعناه يطلع 4 جرام وربع الذهب قيمته فيه الورقه ال 200 جنيه اللي في جيبك متكلفه 16 قرش حتة الورقه دي وشوية الاحبار اللي عليها 16 قرش طب الفرق بين 16 قرش و200 جنيه فين سلطة الدوله سلطة الاصدار وسلطة الدوله دي بتتحسب ازاي في نفسها بها يعني جوه مصر كده يبقى بها يعني ال 200 جنيه 200 جنيه طب وفي خارجها يبقى بقوة انتاجها فالدولار يبقى ب 50 جنيه اللي هو كان ب 39 قرش ونص يعني كان الجنيه ب 5 دولار وبعدين لما انهار بقى 2 ونص في الاول في الاربعينات كان الجنيه المصري ب 5 دولار 5 دولار في 50 جنيه يبقى بكام يبقى ب 250 مره يبقى الجنيه ابو 40 سنة 40 بيساوي 250 جنيه من الجنيهات اللي معانا دي تخيل اضرب دخلك في 250 فلو سالنا الشيخ فلان ولا علان بتاخد كام يقوللك عباره عن 6000 جنيه في الشهر طب 6000 جنيه في الشهر اضربها في 250 يبقى 150 قصادها صفر يبقى مليون ونص ها تخيل نفسك بقى ان انت معاك مليون ونص كل شهر هتعرف تسكن هتعرف تشتري لك عربيه من ام 2 مليون دي اقل حاجه حتى العربيه ام 14 مليون ممكن تجيبها بس بالتقسيط الله هو ايه اللي حصل تضخم والتضخم ده معناته ايه معناته ان هذا الوسيط للتبادل لم يعد سببا لاغناء الفقير الكلام ده كله يا اخي لقينا الشافعي فهمه الكلام اللي احنا قلناه ده والامام الغزالي ركن الشافعيه فاهمينه بيكتبوه عادي كده فيقول رضي الله تعالى عنه في الام وعلة الربا تعبديه فالذهب والفضه جزء عله يعني جزء العله يعني اذا خرج عن الذهبيه والفضيه فلا ربا ويقول ابن الحاجب وهل حرم الربا لجوهرهما او لجوهريتهما الذهب والفضه لجوهرهما عشان هم ذهب وفضه ولا لانهم وسيط للتبادل قال بل لجوهرهما لجوهرهما يعني لانهم ذهب وفضه مش لجوهريتهما يعني كونهم وسيط للتبادل بين الناس لو كان وسيط "

# Transcribe the audio file in chunks
audio_file = "sidiali_reba_zahab_feda.wav"
transcription = transcribe_audio_in_chunks(audio_file)

if transcription:
    print("Final Transcription:\n", transcription)
    
    # Calculate Word Error Rate (WER)
    wer = jiwer.wer(reference_transcription, transcription)
    print(f"\nWord Error Rate (WER): {wer}")
    
    # Save the transcription and WER to a file
    with open("evaluation_results.txt", "w", encoding="utf-8") as f:
        f.write(f"Transcription:\n{transcription}\n\n")
        f.write(f"Reference:\n{reference_transcription}\n\n")
        f.write(f"Word Error Rate (WER): {wer}\n")
    
    print("\nEvaluation results saved to evaluation_results.txt")
else:
    print("Transcription failed.")