# rama-x1-improved.py
import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel
import numpy as np

# ==================== بخش مدل کوچک شما ====================
class TinyLanguageModel(nn.Module):
    def __init__(self, vocab_size, embedding_dim=128, hidden_dim=256, num_layers=2):
        super(TinyLanguageModel, self).__init__()
        self.vocab_size = vocab_size
        self.embedding_dim = embedding_dim
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.lstm = nn.LSTM(embedding_dim, hidden_dim, num_layers, batch_first=True, dropout=0.1)
        self.fc = nn.Linear(hidden_dim, vocab_size)
        self.dropout = nn.Dropout(0.1)
    
    def forward(self, x, hidden=None):
        batch_size = x.size(0)
        x = self.embedding(x)
        
        if hidden is None:
            hidden = self.init_hidden(batch_size)
            
        lstm_out, hidden = self.lstm(x, hidden)
        lstm_out = self.dropout(lstm_out)
        output = self.fc(lstm_out)
        
        return output, hidden
    
    def init_hidden(self, batch_size):
        device = next(self.parameters()).device
        return (torch.zeros(self.num_layers, batch_size, self.hidden_dim).to(device),
                torch.zeros(self.num_layers, batch_size, self.hidden_dim).to(device))

# ==================== بخش BERT فارسی بهبودیافته ====================
class ImprovedPersianBERT:
    def __init__(self, model_name="HooshvareLab/bert-fa-base-uncased"):
        self.model_name = model_name
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModel.from_pretrained(model_name)
        
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
    
    def get_embeddings(self, text):
        inputs = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
        with torch.no_grad():
            outputs = self.model(**inputs)
        return outputs.last_hidden_state
    
    def generate_text(self, prompt, max_length=50, temperature=0.8):
        inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
        
        generated_text = prompt
        for _ in range(max_length - len(prompt.split())):
            with torch.no_grad():
                outputs = self.model(**inputs)
            
            last_hidden = outputs.last_hidden_state[0, -1, :]
            probabilities = torch.softmax(last_hidden / temperature, dim=0)
            next_token_id = torch.multinomial(probabilities, 1).item()
            next_token = self.tokenizer.decode([next_token_id])
            
            # توقف اگر به توکن پایان رسیدیم
            if next_token in [self.tokenizer.eos_token, self.tokenizer.sep_token, '[PAD]']:
                break
                
            generated_text += next_token
            
            # به روزرسانی inputs برای توکن بعدی
            inputs = self.tokenizer(generated_text, return_tensors="pt", truncation=True, max_length=512)
        
        return generated_text

# ==================== سیستم اصلی ====================
class PersianAISystem:
    def __init__(self):
        print("در حال بارگذاری BERT فارسی...")
        self.bert = ImprovedPersianBERT()
        self.tiny_model = None
        self.char_to_idx = None
        self.idx_to_char = None
        print("BERT بارگذاری شد!")
    
    def load_tiny_model(self, model_path, vocab_size, embedding_dim=128, hidden_dim=256):
        try:
            print("در حال بارگذاری مدل کوچک...")
            self.tiny_model = TinyLanguageModel(vocab_size, embedding_dim, hidden_dim)
            checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
            self.tiny_model.load_state_dict(checkpoint['model_state_dict'])
            self.char_to_idx = checkpoint['char_to_idx']
            self.idx_to_char = checkpoint['idx_to_char']
            print("مدل کوچک بارگذاری شد!")
        except FileNotFoundError:
            print("فایل مدل کوچک یافت نشد. فقط از BERT استفاده می‌شود.")
        except Exception as e:
            print(f"خطا در بارگذاری مدل کوچک: {e}")
    
    def generate_with_bert(self, prompt, length=50, temperature=0.8):
        return self.bert.generate_text(prompt, length, temperature)
    
    def generate_with_tiny_model(self, prompt, length=100, temperature=0.8):
        if self.tiny_model is None:
            return "مدل کوچک بارگذاری نشده است. لطفا ابتدا مدل را آموزش دهید یا فایل مدل را قرار دهید."
        
        # این بخش نیاز به پیاده‌سازی دارد
        return "مدل کوچک فعال است اما تولید متن نیاز به پیاده‌سازی بیشتر دارد."
    
    def interactive_chat(self):
        print("سیستم هوش مصنوعی فارسی فعال شد! برای خروج 'exit' را وارد کنید.")
        print("برای انتخاب مدل: 'bert' یا 'tiny' را وارد کنید")
        
        current_model = "bert"
        
        while True:
            user_input = input("\nشما: ")
            
            if user_input.lower() == 'exit':
                break
            
            if user_input.lower() in ['bert', 'tiny']:
                current_model = user_input.lower()
                print(f"مدل تغییر کرد به: {current_model}")
                continue
            
            try:
                if current_model == "bert":
                    response = self.generate_with_bert(user_input)
                else:
                    response = self.generate_with_tiny_model(user_input)
                
                print(f"هوش مصنوعی ({current_model}): {response}")
                
            except Exception as e:
                print(f"خطا: {e}")

# ==================== اجرای اصلی ====================
def main():
    ai_system = PersianAISystem()
    
    # تست BERT
    print("=" * 50)
    print("تست تولید متن با BERT:")
    try:
        bert_text = ai_system.generate_with_bert("سلام")
        print(bert_text)
    except Exception as e:
        print(f"خطا در تولید متن با BERT: {e}")
    
    # سعی در بارگذاری مدل کوچک
    ai_system.load_tiny_model('tiny_language_model.pth', vocab_size=1000)
    
    # فعال کردن حالت گفتگو
    print("\n" + "=" * 50)
    ai_system.interactive_chat()

if __name__ == "__main__":
    main()
