import torch
import torch.nn as nn
import os
import requests
import json
from torch.optim.lr_scheduler import ReduceLROnPlateau

class JowaMAI(nn.Module):
    def __init__(self, vocab_size, embedding_dim=256, hidden_dim=512, num_layers=2, dropout_rate=0.4):
        super(JowaMAI, self).__init__()
        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=dropout_rate)
        self.dropout = nn.Dropout(dropout_rate)
        self.fc = nn.Linear(hidden_dim, vocab_size)

    def forward(self, x, hidden=None):
        embedded = self.embedding(x)
        
        if hidden is None:
            h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).to(x.device)
            c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).to(x.device)
            hidden = (h0, c0)
            
        out, hidden = self.lstm(embedded, hidden)
        out = self.dropout(out)
        out = self.fc(out)
        return out, hidden

data_tokens = [
    "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", 
    "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z",
    "A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", 
    "N", "O", "P", "Q", "R", "S", "T", "U", "V", "W", "X", "Y", "Z",
    "0", "1", "2", "3", "4", "5", "6", "7", "8", "9",
    "!", "\"", "#", "$", "%", "&", "\"", "(", ")", "*", "+", ",", "-", 
    ".", "/", ":", ";", "<", "=", ">", "?", "@", "[", "\\", "]", "^", 
    "_", "`", "{", "|", "}", "~", "¿", " ", "\n"
]

vocab_size = len(data_tokens)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = JowaMAI(vocab_size).to(device)
model_path = "jowa-mAi-1.6.pt"
db_url = "https://raw.githubusercontent.com/AhmedXV-V12/jowamAi-db/main/database.json"

if os.path.exists(model_path):
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.eval()

def fetch_remote_data(url):
    try:
        response = requests.get(url)
        response.raise_for_status()
        content = response.json()
        
        if isinstance(content, dict):
            return content
        else:
            print("Warning: Remote data is not a dictionary. Training might be affected.")
            return {}
    except Exception as e:
        print(f"Error fetching data: {e}")
        return {}

def train_model(model, db_data, data_tokens, epochs=100, patience=10):
    if not db_data:
        print("No data available for training.")
        return

    model.train()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-5)
    criterion = nn.CrossEntropyLoss()
    scheduler = ReduceLROnPlateau(optimizer, 'min', patience=5, factor=0.5, verbose='True')
    
    qa_pairs = []
    for q, a in db_data.items():
        qa_pairs.append(f"Q: {q}\nA: {a}\n")

    if not qa_pairs:
        print("No formatted Q&A pairs for training.")
        return

    best_loss = float('inf')
    epochs_no_improve = 0

    for epoch in range(epochs):
        total_loss = 0
        for text in qa_pairs:
            indices = [data_tokens.index(c) for c in text if c in data_tokens]
            if len(indices) < 2: continue
            
            inputs = torch.tensor(indices[:-1]).unsqueeze(0).to(device)
            targets = torch.tensor(indices[1:]).unsqueeze(0).to(device)
            
            optimizer.zero_grad()
            outputs, _ = model(inputs)
            
            loss = criterion(outputs.view(-1, vocab_size), targets.view(-1))
            
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        
        avg_loss = total_loss / max(1, len(qa_pairs))
        print(f"Epoch {epoch+1}/{epochs} | Loss: {avg_loss:.4f}")

        scheduler.step(avg_loss)

        if avg_loss < best_loss:
            best_loss = avg_loss
            epochs_no_improve = 0
            torch.save(model.state_dict(), model_path)
        else:
            epochs_no_improve += 1
            if epochs_no_improve == patience:
                print(f"Early stopping triggered after {epoch+1} epochs. No improvement for {patience} epochs.")
                break
    
    print("Training complete. Weights saved.")

def generate_response(model, start_str, max_len=100, temperature=0.7):
    model.eval()
    
    prompt_text = f"Q: {start_str}\nA: "
    
    filtered_start = "".join([c for c in prompt_text if c in data_tokens])
    if not filtered_start: return "Invalid input characters or empty prompt."
    
    indices = [data_tokens.index(c) for c in filtered_start]
    input_seq = torch.tensor(indices).unsqueeze(0).to(device)
    
    generated_text = "" 
    hidden = None
    
    with torch.no_grad():
        for _ in range(max_len):
            output, hidden = model(input_seq, hidden)
            
            logits = output[0, -1, :] / temperature
            probabilities = torch.softmax(logits, dim=-1)
            predicted_index = int(torch.multinomial(probabilities, 1).item())
            
            char = data_tokens[predicted_index]
            
            if char == '\n' and generated_text.strip().endswith(('.', '?', '!')):
                break
            
            generated_text += char
            input_seq = torch.tensor([[predicted_index]]).to(device)
            
            if char in ["?", ".", "!"] and len(generated_text.strip()) > 5:
                break
            
    start_answer_idx = generated_text.find('\nA: ')
    if start_answer_idx != -1:
        generated_text = generated_text[start_answer_idx + len('\nA: '):]
    
    return generated_text.strip()

if __name__ == "__main__":
    remote_data = fetch_remote_data(db_url)
    
    file_exists = os.path.exists(model_path)
    if file_exists:
        print(f"Status: Existing model weights found ({model_path}).")
        model.load_state_dict(torch.load(model_path, map_location=device))
        model.eval()
    else:
        print("Status: No existing weights found. Training is recommended.")
    
    inp2 = input("Do you want to train the model with remote data? [y/N]: ").lower().strip()
    if inp2 == "y":
        train_model(model, remote_data, data_tokens, epochs=100)
    else:
        pass
    print("_".center(50,"_"))
    print("jowa-mAi-1.6".center(50,"_"))
    print("_".center(50,"_"))

    print("""
jowa-mAi-1.6 
url: %s
to exit the app write :" shutdown now the app "
or click : ctrl + c
or  ctrl + z           
          """%db_url)
        
    while True:
        prompt = input("You: ")
        if prompt == "shutdown now the app":
            print("exit successfully!")
            break
        else:
            print(f"AI: {generate_response(model, prompt)}")
