# -----------------------------
# 1. Required Libraries
# -----------------------------

import nltk
import numpy as np
from nltk.tokenize import sent_tokenize

from tensorflow.keras.preprocessing.text import Tokenizer
from tensorflow.keras.preprocessing.sequence import pad_sequences
from tensorflow.keras.models import Model
from tensorflow.keras.layers import (
    Input, Embedding, Conv1D,
    GlobalMaxPooling1D, Dense, Concatenate
)

nltk.download('punkt_tab')

# -----------------------------
# 2. Input Document
# -----------------------------

document = """
Machine Learning is changing many industries.
It allows computers to learn from data.
ML models help businesses make better decisions.
Many companies use machine learning for predictions.
"""

# -----------------------------
# 3. Sentence Tokenization
# -----------------------------
# Split document into individual sentences
sentences = sent_tokenize(document)

# Each sentence will be treated as one training instance
print("Sentences:")
for i, s in enumerate(sentences, 1):
    print(f"{i}. {s}")

# -----------------------------
# 4. Text → Integer Encoding
# -----------------------------

# Tokenizer maps each word to a unique integer ID
tokenizer = Tokenizer()
tokenizer.fit_on_texts(sentences)

# Convert sentences into sequences of integers
sequences = tokenizer.texts_to_sequences(sentences)

# Pad all sentences to the same length
# This is required because CNN expects fixed-size input
MAX_LEN = 20
X = pad_sequences(sequences, maxlen=MAX_LEN, padding='post')

# Vocabulary size (+1 for padding index 0)
VOCAB_SIZE = len(tokenizer.word_index) + 1

# -----------------------------
# 5. CNN Model Architecture
# -----------------------------

# Input layer
# Shape = (MAX_LEN,)
# Each input is a sentence represented as word indices
inputs = Input(shape=(MAX_LEN,))

# -----------------------------
# Embedding Layer
# -----------------------------
# Converts word indices → dense vectors
# Output shape = (batch_size, MAX_LEN, EMBED_DIM)

EMBED_DIM = 100

embedding = Embedding(
    input_dim=VOCAB_SIZE,   # size of vocabulary
    output_dim=EMBED_DIM,   # vector size for each word
    input_length=MAX_LEN
)(inputs)

# -----------------------------
# CNN Layers (n-gram feature extraction)
# -----------------------------
# Conv1D slides over word embeddings
# Kernel size decides n-gram length


# 3-gram features
conv3 = Conv1D(
    filters=128,           # number of feature maps
    kernel_size=3,         # trigram
    activation='relu'
)(embedding)               # <-- applied ON embedding output
# Formula : output_length = input_length - kernel_size + 1
#: 20-3+1 = 18
# 4-gram features
conv4 = Conv1D(
    filters=128,
    kernel_size=4,         # four-word phrases
    activation='relu'
)(embedding)

# 5-gram features
conv5 = Conv1D(
    filters=128,
    kernel_size=5,         # five-word phrases
    activation='relu'
)(embedding)

# -----------------------------
# Pooling Layers
# -----------------------------
# GlobalMaxPooling picks the strongest feature
# This makes the model position-invariant

pool3 = GlobalMaxPooling1D()(conv3)
pool4 = GlobalMaxPooling1D()(conv4)
pool5 = GlobalMaxPooling1D()(conv5)

# -----------------------------
# Feature Concatenation
# -----------------------------
# Combine features from different n-gram filters
concat = Concatenate()([pool3, pool4, pool5])

# -----------------------------
# Output Layer
# -----------------------------
# Single neuron → sentence importance score
# Sigmoid gives value between 0 and 1

output = Dense(1, activation='sigmoid')(concat)

# -----------------------------
# Model Creation
# -----------------------------
model = Model(inputs, output)

model.compile(
    optimizer='adam',
    loss='binary_crossentropy',
    metrics=['accuracy']
)

model.summary()

# -----------------------------
# 6. Training (Demo labels)
# -----------------------------
# 1 → important sentence
# 0 → not important sentence

y = np.array([1, 1, 0, 0])

model.fit(X, y, epochs=10, verbose=1)

# -----------------------------
# 7. Sentence Scoring & Ranking
# -----------------------------

# Predict importance score for each sentence
scores = model.predict(X).flatten()

# Pair scores with sentences and sort
ranked_sentences = sorted(
    zip(scores, sentences),
    reverse=True
)

print("\nRanked Sentences:")
for rank, (score, sent) in enumerate(ranked_sentences, 1):
    print(f"{rank}. Score: {score:.4f} | {sent}")

# -----------------------------
# 8. Final Extractive Summary
# -----------------------------

TOP_K = 2
summary = [sent for _, sent in ranked_sentences[:TOP_K]]

print("\n=== CNN Extractive Summary ===")
for s in summary:
    print(s)