import json
import torch
import random
import numpy as np
from sklearn.model_selection import train_test_split
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    TrainingArguments,
    Trainer
)
from collections import Counter



random.seed(42)
np.random.seed(42)
torch.manual_seed(42)
torch.cuda.manual_seed_all(42)



with open("intents_augmented.json", encoding="utf-8") as f:
    data = json.load(f)

sentences = []
labels = []
label2id = {}
id2label = {}

for i, intent in enumerate(data["intents"]):
    tag = intent["tag"]
    label2id[tag] = i
    id2label[i] = tag
    for pattern in intent["patterns"]:
        sentences.append(pattern)
        labels.append(i)



counts = Counter(labels)
max_count = min(max(counts.values()), 300)

balanced_sentences = []
balanced_labels = []

for label in set(labels):
    label_sentences = [s for s, l in zip(sentences, labels) if l == label]
    if len(label_sentences) > max_count:
        
        sampled = random.sample(label_sentences, max_count)
    else:
        
        sampled = label_sentences * (max_count // len(label_sentences)) + \
                  random.sample(label_sentences, max_count % len(label_sentences))
    balanced_sentences.extend(sampled)
    balanced_labels.extend([label] * max_count)

sentences, labels = balanced_sentences, balanced_labels


train_texts, val_texts, train_labels, val_labels = train_test_split(
    sentences, labels, test_size=0.2, random_state=42
)



MODEL_NAME = "aubmindlab/bert-base-arabertv02"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

train_encodings = tokenizer(train_texts, truncation=True, padding=True, max_length=128)
val_encodings = tokenizer(val_texts, truncation=True, padding=True, max_length=128)

class IntentDataset(torch.utils.data.Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels
    def __len__(self):
        return len(self.labels)
    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
        item["labels"] = torch.tensor(self.labels[idx])
        return item

train_dataset = IntentDataset(train_encodings, train_labels)
val_dataset = IntentDataset(val_encodings, val_labels)



model = AutoModelForSequenceClassification.from_pretrained(
    MODEL_NAME,
    num_labels=len(label2id),
    id2label=id2label,
    label2id=label2id
)



training_args = TrainingArguments(
    output_dir="./results",
    learning_rate=5e-5,                  
    per_device_train_batch_size=16,     
    per_device_eval_batch_size=16,
    num_train_epochs=4,                 
    weight_decay=0.01,                   
    save_strategy="epoch",
    logging_dir="./logs",
    save_total_limit=2,
)



trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
    tokenizer=tokenizer
)

trainer.train()



model.save_pretrained("./intent_model")
tokenizer.save_pretrained("./intent_model")

print(" تم حفظ النموذج في مجلد intent_model")