"""
Ekalavya - Multi-Lingual Tokenizer
Supports ALL Indian languages + English
"""
import json
import unicodedata
from typing import List, Dict


class IndianLanguageTokenizer:
    """
    Tokenizer for ALL Indian languages + English
    Supports: Hindi, Bengali, Telugu, Tamil, Marathi, Gujarati, Kannada, 
              Malayalam, Odia, Punjabi, Urdu, and more
    """
    
    def __init__(self, vocab_size: int = 150000):
        self.vocab_size = vocab_size
        self.stoi: Dict[str, int] = {}
        self.itos: Dict[int, str] = {}
        self._build_base_vocab()
    
    def _build_base_vocab(self):
        """Build base vocabulary with all Indian scripts"""
        idx = 0
        
        # Special tokens
        special_tokens = ['<pad>', '<unk>', '<s>', '</s>', '<mask>', '<think>', '</think>']
        for token in special_tokens:
            self.stoi[token] = idx
            self.itos[idx] = token
            idx += 1
        
        # English characters + common words
        for i in range(32, 127):  # ASCII printable
            char = chr(i)
            self.stoi[char] = idx
            self.itos[idx] = char
            idx += 1
        
        # Devanagari (Hindi, Marathi, Sanskrit)
        for i in range(0x0900, 0x097F):
            char = chr(i)
            self.stoi[char] = idx
            self.itos[idx] = char
            idx += 1
        
        # Bengali
        for i in range(0x0980, 0x09FF):
            char = chr(i)
            self.stoi[char] = idx
            self.itos[idx] = char
            idx += 1
        
        # Gurmukhi (Punjabi)
        for i in range(0x0A00, 0x0A7F):
            char = chr(i)
            self.stoi[char] = idx
            self.itos[idx] = char
            idx += 1
        
        # Gujarati
        for i in range(0x0A80, 0x0AFF):
            char = chr(i)
            self.stoi[char] = idx
            self.itos[idx] = char
            idx += 1
        
        # Oriya (Odia)
        for i in range(0x0B00, 0x0B7F):
            char = chr(i)
            self.stoi[char] = idx
            self.itos[idx] = char
            idx += 1
        
        # Tamil
        for i in range(0x0B80, 0x0BFF):
            char = chr(i)
            self.stoi[char] = idx
            self.itos[idx] = char
            idx += 1
        
        # Telugu
        for i in range(0x0C00, 0x0C7F):
            char = chr(i)
            self.stoi[char] = idx
            self.itos[idx] = char
            idx += 1
        
        # Kannada
        for i in range(0x0C80, 0x0CFF):
            char = chr(i)
            self.stoi[char] = idx
            self.itos[idx] = char
            idx += 1
        
        # Malayalam
        for i in range(0x0D00, 0x0D7F):
            char = chr(i)
            self.stoi[char] = idx
            self.itos[idx] = char
            idx += 1
        
        # Sinhala
        for i in range(0x0D80, 0x0DFF):
            char = chr(i)
            self.stoi[char] = idx
            self.itos[idx] = char
            idx += 1
        
        # Urdu/Arabic (for Urdu language)
        for i in range(0x0600, 0x06FF):
            char = chr(i)
            self.stoi[char] = idx
            self.itos[idx] = char
            idx += 1
        
        self.vocab_size = len(self.stoi)
    
    def encode(self, text: str) -> List[int]:
        """Encode text to token IDs"""
        tokens = []
        for char in text:
            if char in self.stoi:
                tokens.append(self.stoi[char])
            else:
                tokens.append(self.stoi['<unk>'])
        return tokens
    
    def decode(self, tokens: List[int]) -> str:
        """Decode token IDs to text"""
        chars = []
        for token in tokens:
            if token in self.itos:
                char = self.itos[token]
                if char not in ['<pad>', '<unk>', '<s>', '</s>', '<mask>', '<think>', '</think>']:
                    chars.append(char)
        return ''.join(chars)
    
    def save(self, path: str):
        """Save tokenizer"""
        data = {
            'stoi': self.stoi,
            'itos': {str(k): v for k, v in self.itos.items()},
            'vocab_size': self.vocab_size
        }
        with open(path, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
    
    @classmethod
    def load(cls, path: str):
        """Load tokenizer"""
        with open(path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        tok = cls()
        tok.stoi = data['stoi']
        tok.itos = {int(k): v for k, v in data['itos'].items()}
        tok.vocab_size = data['vocab_size']
        return tok


# Sample multi-lingual training data
MULTILINGUAL_DATA = """
# English
Hello, how are you? I am fine, thank you.
The weather is beautiful today. Let's go for a walk.

# Hindi (हिंदी)
नमस्ते, आप कैसे हैं? मैं ठीक हूँ, धन्यवाद।
आज मौसम बहुत सुंदर है। चलिए सैर पर चलते हैं।

# Bengali (বাংলা)
নমস্কার, আপনি কেমন আছেন? আমি ভালো আছি, ধন্যবাদ।
আজ আবহাওয়া খুব সুন্দর। চলুন হাঁটতে যাই।

# Telugu (తెలుగు)
నమస్కారం, మీరు ఎలా ఉన్నారు? నేను బాగున్నాను, ధన్యవాదాలు.
ఈ రోజు వాతావరణం చాలా అందంగా ఉంది. నడుద్దాం.

# Tamil (தமிழ்)
வணக்கம், நீங்கள் எப்படி இருக்கிறீர்கள்? நான் நலமாக இருக்கிறேன், நன்றி.
இன்று வானிலை மிகவும் அழகாக உள்ளது. நடக்க செல்வோம்.

# Marathi (मराठी)
नमस्कार, तुम्ही कसे आहात? मी बरा आहे, धन्यवाद.
आज हवामान खूप सुंदर आहे. चला फिरायला जाऊ.

# Gujarati (ગુજરાતી)
નમસ્તે, તમે કેમ છો? હું બરાબર છું, આભાર.
આજે હવામાન ખૂબ સુંદર છે. ચાલો ફરવા જઈએ.

# Kannada (ಕನ್ನಡ)
ನಮಸ್ಕಾರ, ನೀವು ಹೇಗಿದ್ದೀರಿ? ನಾನು ಚೆನ್ನಾಗಿದ್ದೇನೆ, ಧನ್ಯವಾದಗಳು.
ಇಂದು ಹವಾಮಾನ ತುಂಬಾ ಸುಂದರವಾಗಿದೆ. ನಡೆಯಲು ಹೋಗೋಣ.

# Malayalam (മലയാളം)
നമസ്കാരം, സുഖമാണോ? ഞാൻ സുഖമായിരിക്കുന്നു, നന്ദി.
ഇന്ന് കാലാവസ്ഥ വളരെ മനോഹരമാണ്. നടക്കാൻ പോകാം.

# Punjabi (ਪੰਜਾਬੀ)
ਸਤਿ ਸ੍ਰੀ ਅਕਾਲ, ਤੁਸੀਂ ਕਿਵੇਂ ਹੋ? ਮੈਂ ਠੀਕ ਹਾਂ, ਧੰਨਵਾਦ।
ਅੱਜ ਮੌਸਮ ਬਹੁਤ ਸੁੰਦਰ ਹੈ। ਚੱਲੋ ਸੈਰ ਕਰਨ ਚੱਲੀਏ।

# Odia (ଓଡ଼ିଆ)
ନମସ୍କାର, ଆପଣ କେମିତି ଅଛନ୍ତି? ମୁଁ ଭଲ ଅଛି, ଧନ୍ୟବାଦ।
ଆଜି ପାଗ ବହୁତ ସୁନ୍ଦର ଅଛି। ଚଲନ୍ତୁ ବୁଲିବାକୁ ଯିବା।

# Urdu (اردو)
السلام علیکم، آپ کیسے ہیں؟ میں ٹھیک ہوں، شکریہ۔
آج موسم بہت خوبصورت ہے۔ چلیں سیر کرتے ہیں۔

# Mathematics
2 + 2 = 4
The area of a circle is πr²
Euler's formula: e^(iπ) + 1 = 0

# Science
Photosynthesis: 6CO₂ + 6H₂O → C₆H₁₂O₆ + 6O₂
Newton's second law: F = ma
Speed of light: c = 299,792,458 m/s

# Common phrases across languages
Good morning / सुप्रभात / सुप्रভাত / ಶುಭೋದಯ / സുപ്രഭാതം / காலை வணக்கம்
Thank you / धन्यवाद / ধন্যবাদ / ధన్యవాదాలు / நன்றி / धन्यवाद / ਧੰਨਵਾਦ
"""


if __name__ == '__main__':
    print("="*70)
    print("EKALAVYA MYTHOS - Multi-Lingual Tokenizer")
    print("="*70)
    
    # Create tokenizer
    tok = IndianLanguageTokenizer()
    print(f"\n✅ Tokenizer created: {tok.vocab_size} tokens")
    
    # Test with different languages
    test_texts = [
        ("English", "Hello, how are you?"),
        ("Hindi", "नमस्ते, आप कैसे हैं?"),
        ("Bengali", "নমস্কার, আপনি কেমন আছেন?"),
        ("Telugu", "నమస్కారం, మీరు ఎలా ఉన్నారు?"),
        ("Tamil", "வணக்கம், நீங்கள் எப்படி இருக்கிறீர்கள்?"),
        ("Marathi", "नमस्कार, तुम्ही कसे आहात?"),
        ("Gujarati", "નમસ્તે, તમે કેમ છો?"),
        ("Kannada", "ನಮಸ್ಕಾರ, ನೀವು ಹೇಗಿದ್ದೀರಿ?"),
        ("Malayalam", "നമസ്കാരം, സുഖമാണോ?"),
        ("Punjabi", "ਸਤਿ ਸ੍ਰੀ ਅਕਾਲ, ਤੁਸੀਂ ਕਿਵੇਂ ਹੋ?"),
    ]
    
    print("\n🌍 Testing Multi-Lingual Support:")
    for lang, text in test_texts:
        encoded = tok.encode(text)
        decoded = tok.decode(encoded)
        match = "✓" if decoded == text else "✗"
        print(f"  {match} {lang:12s}: {text}")
    
    # Train on sample data
    print("\n📚 Training tokenizer on multi-lingual data...")
    data_tokens = tok.encode(MULTILINGUAL_DATA)
    print(f"   Sample data: {len(MULTILINGUAL_DATA)} chars → {len(data_tokens)} tokens")
    
    print("\n" + "="*70)
    print("✅ Tokenizer ready for ALL Indian languages!")
    print("="*70)