import numpy as np
from sklearn.datasets import fetch_20newsgroups
from sklearn.model_selection import train_test_split
from tensorflow.keras.preprocessing.text import Tokenizer
from tensorflow.keras.preprocessing.sequence import pad_sequences
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Embedding, Conv1D, GlobalMaxPooling1D, Dense, Dropout
from tensorflow.keras.utils import to_categorical

# 1. Load Dataset
data = fetch_20newsgroups(subset='all')
X = data.data
y = data.target

num_classes = len(set(y))

# 2. Split Data
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

# 3. Tokenization
vocab_size = 10000
max_len = 200

tokenizer = Tokenizer(num_words=vocab_size)
tokenizer.fit_on_texts(X_train)

X_train_seq = tokenizer.texts_to_sequences(X_train)
X_test_seq = tokenizer.texts_to_sequences(X_test)

# 4. Padding
X_train_pad = pad_sequences(X_train_seq, maxlen=max_len)
X_test_pad = pad_sequences(X_test_seq, maxlen=max_len)

# 5. One-hot Encoding
y_train_cat = to_categorical(y_train, num_classes)
y_test_cat = to_categorical(y_test, num_classes)

# 6. Build CNN Model
model = Sequential()
model.add(Embedding(vocab_size, 128, input_length=max_len))
model.add(Conv1D(128, 5, activation='relu'))
model.add(GlobalMaxPooling1D())
model.add(Dense(64, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(num_classes, activation='softmax'))

# 7. Compile
model.compile(loss='categorical_crossentropy',
              optimizer='adam',
              metrics=['accuracy'])

# 8. Train
model.fit(X_train_pad, y_train_cat,
          epochs=5,
          batch_size=64,
          validation_split=0.1)

# 9. Evaluate
loss, acc = model.evaluate(X_test_pad, y_test_cat)
print("Test Accuracy:", acc)

# 🔟 Prediction Function
def predict_category(text):
    seq = tokenizer.texts_to_sequences([text])
    pad = pad_sequences(seq, maxlen=max_len)
   
    pred = model.predict(pad)
    class_index = np.argmax(pred)
   
    print("\nInput Sentence:", text)
    print("Predicted Category:", data.target_names[class_index])

# 🔥 Test with your own sentences
predict_category("The government is planning new policies for the economy")
predict_category("The car engine performance is very fast and efficient")
predict_category("The team played a great match and won the tournament")