import os
import sys
import json
import time

# Ensure user packages are in path for onnxruntime
sys.path.append("/home/hypr4/.local/lib/python3.12/site-packages")

import numpy as np
import onnxruntime as ort

class ONNXTransliterator:
    def __init__(self, model_dir):
        # Load vocabularies
        with open(os.path.join(model_dir, "input_vocab.json"), "r", encoding="utf-8") as f:
            self.src_vocab = json.load(f)
        with open(os.path.join(model_dir, "target_vocab.json"), "r", encoding="utf-8") as f:
            self.tgt_vocab = json.load(f)
            
        self.src_idx2char = {v: k for k, v in self.src_vocab.items()}
        self.tgt_idx2char = {v: k for k, v in self.tgt_vocab.items()}
        
        # Load ONNX (force 1-thread CPU mode)
        opts = ort.SessionOptions()
        opts.intra_op_num_threads = 1
        opts.inter_op_num_threads = 1
        opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        
        self.encoder_sess = ort.InferenceSession(os.path.join(model_dir, "encoder.onnx"), sess_options=opts)
        self.decoder_sess = ort.InferenceSession(os.path.join(model_dir, "decoder.onnx"), sess_options=opts)

    def transliterate_word(self, word):
        src_ids = [self.src_vocab["<s>"]]
        for char in word:
            src_ids.append(self.src_vocab.get(char, self.src_vocab["<unk>"]))
        src_ids.append(self.src_vocab["</s>"])
        
        input_ids = np.array([src_ids], dtype=np.int64)
        
        enc_outputs, enc_h, enc_c = self.encoder_sess.run(
            ["encoder_outputs", "h_states", "c_states"],
            {"input_ids": input_ids}
        )
        
        num_layers = 2
        hidden_dim = 256
        dec_h = np.zeros((num_layers, 1, hidden_dim), dtype=np.float32)
        dec_c = np.zeros((num_layers, 1, hidden_dim), dtype=np.float32)
        
        for i in range(num_layers):
            dec_h[i] = (enc_h[2*i] + enc_h[2*i+1]) / 2.0
            dec_c[i] = (enc_c[2*i] + enc_c[2*i+1]) / 2.0
            
        dec_input = np.array([self.tgt_vocab["<s>"]], dtype=np.int64)
        output_chars = []
        
        for step in range(32):
            logits, dec_h, dec_c, _ = self.decoder_sess.run(
                ["logits", "h", "c", "attn_weights"],
                {
                    "input_char": dec_input,
                    "prev_h": dec_h,
                    "prev_c": dec_c,
                    "encoder_outputs": enc_outputs
                }
            )
            
            next_char_idx = int(np.argmax(logits[0]))
            if next_char_idx == self.tgt_vocab["</s>"] or next_char_idx == self.tgt_vocab["<pad>"]:
                break
                
            output_chars.append(self.tgt_idx2char.get(next_char_idx, ""))
            dec_input = np.array([next_char_idx], dtype=np.int64)
            
        return "".join(output_chars)

    def transliterate_sentence(self, sentence):
        words = sentence.split()
        translated_words = []
        for word in words:
            clean_word = "".join(c for c in word if '\u0900' <= c <= '\u097f')
            if not clean_word:
                translated_words.append(word)
                continue
                
            translated = self.transliterate_word(clean_word)
            
            prefix = ""
            for c in word:
                if not ('\u0900' <= c <= '\u097f'):
                    prefix += c
                else:
                    break
                    
            suffix = ""
            for c in reversed(word):
                if not ('\u0900' <= c <= '\u097f'):
                    suffix = c + suffix
                else:
                    break
                    
            translated_words.append(prefix + translated + suffix)
            
        return " ".join(translated_words)

def main():
    print("=" * 80)
    print("EVALUATING UPDATED TEXT-ALIGNED ONNX MODEL ON REAL TRANSCRIPTS")
    print("=" * 80)
    
    base_dir = os.path.dirname(os.path.abspath(__file__))
    onnx_dir = os.path.join(base_dir, "../models")
    transcript_file = os.path.join(base_dir, "./eval_transcript.json")
    
    # 1. Load ONNX model
    onnx_rnn = ONNXTransliterator(onnx_dir)
    print("✓ ONNX RNN model loaded successfully!")
    
    # 2. Load transcript segments
    with open(transcript_file, "r", encoding="utf-8") as f:
        data = json.load(f)
    segments = data.get("segments", [])
    
    test_sentences = []
    for seg in segments:
        text = seg.get("text", "").strip()
        if any('\u0900' <= char <= '\u097f' for char in text):
            test_sentences.append(text)
            if len(test_sentences) >= 30:
                break
                
    # 3. Transliterate and print side-by-side
    print("\n" + "=" * 90)
    print(f"{'Devanagari Transcript Segment':<45} | {'Texting-Aligned Hinglish Output':<45}")
    print("=" * 90)
    
    latencies = []
    for idx, sentence in enumerate(test_sentences):
        t_start = time.perf_counter()
        output = onnx_rnn.transliterate_sentence(sentence)
        latency = (time.perf_counter() - t_start) * 1000  # ms
        latencies.append(latency)
        
        print(f"{idx+1:2d}. {sentence}")
        print(f"    -> '{output}' ({latency:.1f} ms)")
        print("-" * 90)
        
    print(f"\nAverage Sentence Transliteration Latency: {sum(latencies)/len(latencies):.2f} ms")
    print("=" * 90)

if __name__ == "__main__":
    main()