# 模擬詞表
vocab = {
    "<PAD>": 0,  # 填充位 (佔位符)
    "<UNK>": 1,  # 未知詞 (遇到沒見過的詞就用它代替)
    "I": 2,
    "love": 3,
    "AI": 4,
    "hate": 5,
    ...
}

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from collections import Counter
import re

# 模擬一些數據 (真實場景中你會從文件讀取)
raw_data = [
    ("This movie is amazing I love it", 1), # 1 代表正面
    ("What a waste of time terrible acting", 0), # 0 代表負面
    ("Great story and fantastic visual effects", 1),
    ("I fell asleep halfway through boring", 0),
    ("The plot is confusing but the music is good", 1),
    ("Worst movie ever do not watch", 0)
] * 100 # 複製多一點以便跑得起來

# --- 1. 簡易分詞器 ---
def tokenizer(text):
    # 轉小寫，去掉非字母字符，按空格切分
    text = re.sub(r'[^a-zA-Z\s]', '', text.lower())
    return text.split()

# --- 2. 構建詞表 ---
def build_vocab(data, min_freq=1):
    all_tokens = []
    for text, _ in data:
        all_tokens.extend(tokenizer(text))
    
    # 統計詞頻
    word_counts = Counter(all_tokens)
    
    # 定義特殊字符
    vocab = {"<PAD>": 0, "<UNK>": 1}
    idx = 2
    for word, count in word_counts.items():
        if count >= min_freq:
            vocab[word] = idx
            idx += 1
    return vocab

# 構建詞表
vocab = build_vocab(raw_data)
print(f"詞表大小: {len(vocab)}")
print(f"Token示例: 'movie' -> {vocab.get('movie')}")

# --- 3. 自定義 Dataset ---
class IMDBDataset(Dataset):
    def __init__(self, data, vocab, max_len=10):
        self.data = data
        self.vocab = vocab
        self.max_len = max_len

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        text, label = self.data[idx]
        tokens = tokenizer(text)
        
        # 文本轉數字索引
        token_ids = [self.vocab.get(token, self.vocab["<UNK>"]) for token in tokens]
        
        # 填充或截斷 (Padding / Truncation)
        if len(token_ids) < self.max_len:
            # 短了就補 0
            token_ids += [self.vocab["<PAD>"]] * (self.max_len - len(token_ids))
        else:
            # 長了就截斷
            token_ids = token_ids[:self.max_len]
            
        return torch.tensor(token_ids), torch.tensor(label, dtype=torch.float32)

# 實例化 DataLoader
dataset = IMDBDataset(raw_data, vocab, max_len=10)
dataloader = DataLoader(dataset, batch_size=4, shuffle=True)

# 測試一下管道
inputs, labels = next(iter(dataloader))
print(f"Input Batch Shape: {inputs.shape}") # Should be [4, 10]

class SentimentLSTM(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim, output_dim):
        super(SentimentLSTM, self).__init__()
        
        # 1. Embedding 層: 把整數索引變成向量
        # padding_idx=0 告訴模型: 索引為0的是填充物，不要計算它的梯度，也沒意義
        self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)
        
        # 2. LSTM 層
        # batch_first=True 讓輸入格式變成 (batch, seq_len, feature)
        self.lstm = nn.LSTM(embedding_dim, hidden_dim, batch_first=True)
        
        # 3. 全連接層 (分類器)
        self.fc = nn.Linear(hidden_dim, output_dim)
        
        # 4. Sigmoid 激活 (因為是二分類，輸出 0~1 的概率)
        self.sigmoid = nn.Sigmoid()
        
    def forward(self, text):
        # text shape: [batch_size, seq_len] -> [4, 10]
        
        # embedded shape: [batch_size, seq_len, emb_dim] -> [4, 10, 32]
        embedded = self.embedding(text)
        
        # lstm 輸出: output (所有時刻的狀態), (hidden, cell) (最後時刻的狀態)
        # 我們只需要最後時刻的狀態來代表整句話的意思
        _, (hidden, _) = self.lstm(embedded)
        
        # hidden shape: [1, batch_size, hidden_dim] -> Squeeze -> [batch_size, hidden_dim]
        # 取最後一層的 hidden state
        last_hidden = hidden[-1]
        
        # 全連接 + 激活
        return self.sigmoid(self.fc(last_hidden))

# 初始化模型
vocab_size = len(vocab)
embedding_dim = 32
hidden_dim = 64
output_dim = 1

model = SentimentLSTM(vocab_size, embedding_dim, hidden_dim, output_dim)
print(model)

import torch.optim as optim

criterion = nn.BCELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# --- 訓練循環 ---
epochs = 5
for epoch in range(epochs):
    total_loss = 0
    for texts, labels in dataloader:
        optimizer.zero_grad()
        
        # 前向傳播
        predictions = model(texts).squeeze(1) # 把 [4, 1] 變成 [4]
        
        loss = criterion(predictions, labels)
        
        # 反向傳播
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
    
    print(f"Epoch {epoch+1}, Loss: {total_loss/len(dataloader):.4f}")

# --- 預測新句子 ---
def predict_sentiment(sentence):
    model.eval（)
    # 1. 預處理
    tokens = tokenizer(sentence)
    idx = [vocab.get(t, vocab["<UNK>"]) for t in tokens]
    # 簡單的 padding 處理 (生產環境需要更嚴謹)
    if len(idx) < 10:
        idx += [0] * (10 - len(idx))
    else:
        idx = idx[:10]
        
    tensor = torch.LongTensor(idx).unsqueeze(0) # 加 batch 維度
    
    # 2. 推理
    with torch.no_grad():
        prediction = model(tensor).item()
        
    status = "正面 😄" if prediction > 0.5 else "負面 😡"
    print(f"評論: '{sentence}' -> {status} (概率: {prediction:.2f})")

print("\n--- 測試模型 ---")
predict_sentiment("This movie is amazing")
predict_sentiment("I hate this boring movie")