import numpy as np
import random
import time
import re
import os
import ast
import operator

# --- 基礎數學與安全的 AST 數學運算 ---
def safe_math_eval(expr):
    """安全地解析並計算數學表達式"""
    allowed_operators = {
        ast.Add: operator.add, ast.Sub: operator.sub,
        ast.Mult: operator.mul, ast.Div: operator.truediv,
        ast.USub: operator.neg, ast.UAdd: operator.pos
    }
    def evaluate(node):
        if isinstance(node, ast.Constant):
            return node.value
        elif isinstance(node, ast.BinOp):
            return allowed_operators[type(node.op)](evaluate(node.left), evaluate(node.right))
        elif isinstance(node, ast.UnaryOp):
            return allowed_operators[type(node.op)](evaluate(node.operand))
        else:
            raise TypeError("Unsupported operation")
    try:
        tree = ast.parse(expr, mode='eval')
        return evaluate(tree.body)
    except Exception:
        return None

def sigmoid(x):
    return 1 / (1 + np.exp(-np.clip(x, -20, 20)))

def tanh(x):
    return np.tanh(np.clip(x, -20, 20))

# --- LanAI 1.0 Seq2Seq 架構 (引入 Encoder-Decoder) ---
class LanAISeq2Seq:
    # 腦容量提升至 384，詞向量提升至 128，賦予更強的特徵提取能力
    def __init__(self, vocab_size, embedding_dim=128, hidden_size=384): 
        self.vocab_size = vocab_size
        self.embedding_dim = embedding_dim
        self.hidden_size = hidden_size
        
        limit = np.sqrt(6 / (vocab_size + hidden_size))
        
        self.E = np.random.uniform(-limit, limit, (vocab_size, embedding_dim))
        
        # GRU 權重 (Encoder 與 Decoder 共用以節省記憶體並共享語境)
        self.W_r = np.random.uniform(-limit, limit, (embedding_dim, hidden_size))
        self.U_r = np.random.uniform(-limit, limit, (hidden_size, hidden_size))
        self.b_r = np.zeros(hidden_size)
        
        self.W_z = np.random.uniform(-limit, limit, (embedding_dim, hidden_size))
        self.U_z = np.random.uniform(-limit, limit, (hidden_size, hidden_size))
        self.b_z = np.zeros(hidden_size)
        
        self.W_h = np.random.uniform(-limit, limit, (embedding_dim, hidden_size))
        self.U_h = np.random.uniform(-limit, limit, (hidden_size, hidden_size))
        self.b_h = np.zeros(hidden_size)
        
        self.W_y = np.random.uniform(-limit, limit, (hidden_size, vocab_size))
        self.b_y = np.zeros(vocab_size)

    def expand_weights(self, new_vocab_size):
        diff = new_vocab_size - self.vocab_size
        if diff <= 0: return
        
        limit = np.sqrt(6 / (new_vocab_size + self.hidden_size))
        
        new_E = np.random.uniform(-limit, limit, (diff, self.embedding_dim))
        self.E = np.vstack((self.E, new_E))
            
        new_Wy = np.random.uniform(-limit, limit, (self.hidden_size, diff))
        self.W_y = np.hstack((self.W_y, new_Wy))
        self.b_y = np.concatenate((self.b_y, np.zeros(diff)))
        
        self.vocab_size = new_vocab_size

    def clip_gradient(self, grad, threshold=3.0):
        norm = np.linalg.norm(grad)
        if norm > threshold:
            return grad * (threshold / norm)
        return grad

    def train_sequence(self, input_indices, target_indices, lr=0.05):
        """真正的 Seq2Seq 訓練法：先編碼(Encode)整句話，再解碼(Decode)答案"""
        
        # === 1. Encoder 階段 (讀取輸入，產生思想向量) ===
        h_state = np.zeros(self.hidden_size)
        for idx in input_indices:
            idx = idx if idx < self.vocab_size else 0
            x = self.E[idx]
            r = sigmoid(np.dot(x, self.W_r) + np.dot(h_state, self.U_r) + self.b_r)
            z = sigmoid(np.dot(x, self.W_z) + np.dot(h_state, self.U_z) + self.b_z)
            h_tilde = tanh(np.dot(x, self.W_h) + np.dot(r * h_state, self.U_h) + self.b_h)
            h_state = (1 - z) * h_state + z * h_tilde
            
        # 經過 Encoder，h_state 現在包含了整句問題的「語意精華」
        
        # === 2. Decoder 階段 (Teacher Forcing 反向傳播) ===
        total_loss = 0
        current_x_idx = input_indices[-1] if len(input_indices) > 0 else 0
        
        for target_word_idx in target_indices:
            x = self.E[current_x_idx]
            h_prev = h_state.copy()

            # Decoder 前向傳播
            r = sigmoid(np.dot(x, self.W_r) + np.dot(h_prev, self.U_r) + self.b_r)
            z = sigmoid(np.dot(x, self.W_z) + np.dot(h_prev, self.U_z) + self.b_z)
            h_tilde = tanh(np.dot(x, self.W_h) + np.dot(r * h_prev, self.U_h) + self.b_h)
            h_state = (1 - z) * h_prev + z * h_tilde
            y = np.dot(h_state, self.W_y) + self.b_y
            
            # Softmax 與 Loss
            y_shifted = y - np.max(y)
            exp_y = np.exp(y_shifted)
            probs = exp_y / (np.sum(exp_y) + 1e-9)
            total_loss -= np.log(max(probs[target_word_idx], 1e-10))

            # 反向傳播計算梯度
            dy = probs.copy()
            dy[target_word_idx] -= 1.0

            dW_y = np.outer(h_state, dy)
            db_y = dy
            dh = np.dot(self.W_y, dy)

            dh_tilde = dh * z * (1 - h_tilde**2)
            dz = dh * (h_tilde - h_prev) * z * (1 - z)
            dr = np.dot(self.U_h, dh_tilde) * h_prev * r * (1 - r)

            dW_h = np.outer(x, dh_tilde); dU_h = np.outer(r * h_prev, dh_tilde); db_h = dh_tilde
            dW_z = np.outer(x, dz);       dU_z = np.outer(h_prev, dz);           db_z = dz
            dW_r = np.outer(x, dr);       dU_r = np.outer(h_prev, dr);           db_r = dr
            dx = np.dot(self.W_h, dh_tilde) + np.dot(self.W_z, dz) + np.dot(self.W_r, dr)

            # 更新權重
            self.W_y -= lr * self.clip_gradient(dW_y); self.b_y -= lr * self.clip_gradient(db_y)
            self.W_h -= lr * self.clip_gradient(dW_h); self.U_h -= lr * self.clip_gradient(dU_h); self.b_h -= lr * self.clip_gradient(db_h)
            self.W_z -= lr * self.clip_gradient(dW_z); self.U_z -= lr * self.clip_gradient(dU_z); self.b_z -= lr * self.clip_gradient(db_z)
            self.W_r -= lr * self.clip_gradient(dW_r); self.U_r -= lr * self.clip_gradient(dU_r); self.b_r -= lr * self.clip_gradient(db_r)
            self.E[current_x_idx] -= lr * self.clip_gradient(dx)

            # Teacher Forcing: 下一步的輸入強制使用正確答案，加速模型收斂
            current_x_idx = target_word_idx
            
        return total_loss

    def generate(self, input_indices, word_to_idx, top_k=3, temperature=0.5):
        """引入 Top-K 採樣的生成器，產生像 LLM 一樣穩定且具創造力的回答"""
        # === 1. Encoder 階段 ===
        h_state = np.zeros(self.hidden_size)
        for idx in input_indices:
            idx = idx if idx < self.vocab_size else 0
            x = self.E[idx]
            r = sigmoid(np.dot(x, self.W_r) + np.dot(h_state, self.U_r) + self.b_r)
            z = sigmoid(np.dot(x, self.W_z) + np.dot(h_state, self.U_z) + self.b_z)
            h_tilde = tanh(np.dot(x, self.W_h) + np.dot(r * h_state, self.U_h) + self.b_h)
            h_state = (1 - z) * h_state + z * h_tilde
            
        # === 2. Decoder 階段 ===
        res_indices = []
        current_x_idx = input_indices[-1] if len(input_indices) > 0 else word_to_idx.get("<PAD>", 0)
        
        for _ in range(25): # 允許生成更長的句子
            x = self.E[current_x_idx]
            
            r = sigmoid(np.dot(x, self.W_r) + np.dot(h_state, self.U_r) + self.b_r)
            z = sigmoid(np.dot(x, self.W_z) + np.dot(h_state, self.U_z) + self.b_z)
            h_tilde = tanh(np.dot(x, self.W_h) + np.dot(r * h_state, self.U_h) + self.b_h)
            h_state = (1 - z) * h_state + z * h_tilde
            y = np.dot(h_state, self.W_y) + self.b_y
            
            # 溫度縮放
            logits = y / temperature
            logits = logits - np.max(logits)
            exp_l = np.exp(logits)
            probs = exp_l / (np.sum(exp_l) + 1e-9)
            
            # [1.0 核心升級] Top-K 採樣：過濾掉機率極低的荒謬單字
            if top_k > 0:
                indices_to_remove = probs < np.sort(probs)[-top_k]
                probs[indices_to_remove] = 0
                probs /= np.sum(probs) # 重新歸一化
            
            best_idx = np.random.choice(len(probs), p=probs)
            
            # 防呆機制：避免第一句話就句號
            if (best_idx == word_to_idx["<EOS>"] or best_idx == 0) and len(res_indices) == 0:
                probs[best_idx] = 0
                probs /= np.sum(probs)
                best_idx = np.random.choice(len(probs), p=probs)
                
            if best_idx == word_to_idx["<EOS>"] or best_idx == 0:
                break
                
            res_indices.append(best_idx)
            current_x_idx = best_idx
            
        return res_indices, h_state

# --- LanAI 1.0 主系統 ---
class LanAI:
    def __init__(self, model_path="lanai_v1_memory.npz"): # 為了 Seq2Seq 架構，使用新的記憶檔名
        self.model_path = model_path
        self.context_h = None 
        
        self.base_training_data = [
            (["hello"], ["hi", "there", "how", "can", "i", "help", "you", "today"]),
            (["hi"], ["hello", "how", "are", "you", "doing"]),
            (["hi", "there"], ["hello", "nice", "to", "meet", "you"]),
            (["how", "are", "you"], ["i", "am", "doing", "great", "thank", "you", "for", "asking"]),
            (["i", "am", "good"], ["that", "is", "great", "to", "hear"]),
            (["what", "is", "your", "name"], ["my", "name", "is", "lanai", "version", "1.0"]),
            (["hi", "what", "is", "your", "name"], ["hello", "my", "name", "is", "lanai"]),
            (["who", "are", "you"], ["i", "am", "lanai", "an", "artificial", "intelligence"]),
            (["who", "made", "you"], ["i", "was", "created", "by", "a", "human", "developer"]),
            (["what", "can", "you", "do"], ["i", "can", "chat", "calculate", "math", "and", "learn", "new", "things"]),
            (["ok"], ["okay", "then"]),
            (["yes"], ["i", "agree", "with", "you"]),
            (["thanks"], ["you", "are", "very", "welcome"]),
            (["bye"], ["goodbye", "have", "a", "nice", "day"]),
            (["have", "you", "eaten", "yet"], ["yes", "i", "have", "already", "eaten"]),
            (["what", "is", "passive", "voice"], ["the", "action", "is", "done", "to", "the", "subject"]),
            (["the", "book", "was", "written"], ["by", "a", "very", "famous", "writer"]),
            (["do", "you", "know", "the", "boy"], ["who", "is", "playing", "basketball"]),
            (["this", "is", "the", "house"], ["where", "i", "lived", "before", "with", "my", "family"])
        ]
        self.training_data = list(self.base_training_data)
        
        self.default_vocab = [
            "<PAD>", "<EOS>", "hello", "hi", "there", "how", "can", "i", "help", "you", "today", "are", "doing", 
            "am", "great", "thank", "for", "asking", "good", "that", "is", "to", "hear", "fine", "glad",
            "what", "your", "name", "my", "lanai", "version", "1.0", "who", "an", "artificial", "intelligence",
            "made", "was", "created", "by", "a", "human", "developer", "old", "do", "not", "have", "age", "ai",
            "calculate", "math", "and", "learn", "new", "things", "ok", "okay", "then", "yes", "agree", "with",
            "no", "understand", "thanks", "very", "welcome", "bye", "goodbye", "nice", "day", "eaten", "yet",
            "already", "done", "learned", "recently", "she", "has", "been", "there", "since", "last", "year",
            "passive", "voice", "action", "subject", "book", "written", "famous", "writer", "know", "boy",
            "playing", "basketball", "this", "house", "where", "lived", "before", "family", "meet"
        ]

        if not self.load_model():
            print("[系統] 正在初始化 1.0 Beta 核心 (Seq2Seq Transformer-Lite Architecture)...")
            self.vocab = list(self.default_vocab)
            self.word_to_idx = {word: i for i, word in enumerate(self.vocab)}
            self.idx_to_word = {i: word for i, word in enumerate(self.vocab)}
            self.model = LanAISeq2Seq(len(self.vocab))
            self.is_new_model = True
        else:
            self.is_new_model = False
            self.merge_new_knowledge() 

    def merge_new_knowledge(self):
        new_words_count = 0
        for word in self.default_vocab:
            if word not in self.word_to_idx:
                self.vocab.append(word)
                self.word_to_idx[word] = len(self.vocab) - 1
                self.idx_to_word[len(self.vocab) - 1] = word
                new_words_count += 1
                
        if new_words_count > 0:
            self.model.expand_weights(len(self.vocab))
            self.training_data.extend(self.base_training_data)
            self.save_model() 

    def save_model(self):
        print("[系統] 正在寫入 Seq2Seq 記憶狀態 (.npz)...")
        np.savez(self.model_path,
                 vocab=np.array(self.vocab),
                 E=self.model.E,
                 W_r=self.model.W_r, U_r=self.model.U_r, b_r=self.model.b_r,
                 W_z=self.model.W_z, U_z=self.model.U_z, b_z=self.model.b_z,
                 W_h=self.model.W_h, U_h=self.model.U_h, b_h=self.model.b_h,
                 W_y=self.model.W_y, b_y=self.model.b_y)
        print("[系統] 記憶儲存成功！")

    def load_model(self):
        if os.path.exists(self.model_path):
            try:
                print(f"[系統] 偵測到 V1.0 記憶檔 {self.model_path}，正在喚醒...")
                data = np.load(self.model_path, allow_pickle=True)
                self.vocab = data["vocab"].tolist()
                self.word_to_idx = {word: i for i, word in enumerate(self.vocab)}
                self.idx_to_word = {i: word for i, word in enumerate(self.vocab)}
                
                loaded_hidden_size = data["W_y"].shape[0]
                loaded_embedding_dim = data["E"].shape[1]
                
                self.model = LanAISeq2Seq(len(self.vocab), embedding_dim=loaded_embedding_dim, hidden_size=loaded_hidden_size)
                
                self.model.E = data["E"]
                self.model.W_r = data["W_r"]; self.model.U_r = data["U_r"]; self.model.b_r = data["b_r"]
                self.model.W_z = data["W_z"]; self.model.U_z = data["U_z"]; self.model.b_z = data["b_z"]
                self.model.W_h = data["W_h"]; self.model.U_h = data["U_h"]; self.model.b_h = data["b_h"]
                self.model.W_y = data["W_y"]; self.model.b_y = data["b_y"]
                return True
            except Exception as e:
                print(f"[警告] 記憶載入失敗 ({e})。")
                return False
        return False

    def add_new_word(self, word):
        if word not in self.word_to_idx:
            self.vocab.append(word)
            self.word_to_idx[word] = len(self.vocab) - 1
            self.idx_to_word[len(self.vocab) - 1] = word
            self.model.expand_weights(len(self.vocab))
            return True
        return False

    def solve_math(self, text):
        num_map = {"one": "1", "two": "2", "three": "3", "four": "4", "five": "5"}
        for k, v in num_map.items():
            text = re.sub(rf'\b{k}\b', v, text)
            
        clean_text = text.replace("plus", "+").replace("minus", "-").replace("times", "*").replace("divided by", "/")
        expr = "".join(re.findall(r'[\d\+\-\*\/\(\)\.]', clean_text))
        
        if expr and any(op in expr for op in "+-*/"):
            res = safe_math_eval(expr)
            if res is not None:
                res = int(res) if res == int(res) else round(res, 4)
                return f"The answer is {res}."
        return None

    def train_epochs(self, target_epochs, initial_lr=0.03): # Seq2Seq 需要較小的學習率以防震盪
        for epoch in range(target_epochs):
            random.shuffle(self.training_data)
            loss = 0
            current_lr = initial_lr * (1 - (epoch / (target_epochs * 1.2))) 
            
            for words, targets in self.training_data:
                indices = [self.word_to_idx.get(w, 0) for w in words]
                t_indices = [self.word_to_idx.get(w, 0) for w in targets] + [self.word_to_idx["<EOS>"]]
                loss += self.model.train_sequence(indices, t_indices, lr=current_lr)
            if epoch > 0 and epoch % 20 == 0:
                print(f"訓練進度 - Epoch: {epoch} | Loss: {loss/len(self.training_data):.4f} | LR: {current_lr:.4f}  ", end="\r")

    def thinking_layer(self, user_input, words):
        math_res = self.solve_math(user_input)
        
        logic_keywords = ["why", "how", "can", "what", "math", "solve"]
        grammar_keywords = ["grammar", "tense", "passive", "voice"]
        
        if any(w in words for w in grammar_keywords): intent = "grammar"
        elif any(w in words for w in logic_keywords): intent = "logic"
        else: intent = "chat"
            
        return math_res, intent

    def chat(self):
        print("="*65)
        print("  LanAI 1.0 Beta (Seq2Seq | Top-K Nucleus | Context Encoder)")
        print("="*65)
        
        if self.is_new_model:
            print("[系統] 全新 Seq2Seq 大腦建構中，執行深度語義壓縮 (800 Epochs)...")
            self.train_epochs(800, initial_lr=0.04)
            print("\n[系統] 基礎思維網路建立完成。")
            self.save_model()
        else:
            print("\n[系統] 記憶載入完畢，準備就緒。")

        while True:
            try:
                user_input = input("\nUser: ")
                if not user_input.strip(): continue
                low_input = user_input.lower().strip()
                if low_input in ['exit', 'quit']:
                    self.save_model()
                    break
                
                words = low_input.replace("?", "").replace("!", "").replace(",", "").split()
                math_res, intent = self.thinking_layer(low_input, words)

                if math_res:
                    print(f"AI (Logic Mode): {math_res}")
                    continue
                
                new_words_learned = False
                for w in words:
                    if self.add_new_word(w):
                        print(f"[系統意識] 擴充詞彙空間: '{w}'")
                        new_words_learned = True

                if new_words_learned and len(words) >= 3:
                    print("[動態微調] 正在將新句型寫入 Seq2Seq 記憶體...", end="\r")
                    self.training_data.append((words, ["i", "understand"] + words))
                    self.train_epochs(40, initial_lr=0.02)
                else:
                    self.train_epochs(3, initial_lr=0.005)
                
                print(" " * 60, end="\r") 
                
                indices = [self.word_to_idx.get(w, 0) for w in words]
                
                # Top-K 設定：設定為 3 強制收斂語義，Temperature 稍微調低以追求精確
                res_indices, self.context_h = self.model.generate(
                    indices, self.word_to_idx, top_k=3, temperature=0.4
                )
                
                print("AI: ", end="", flush=True)
                if not res_indices:
                    print("I am not sure how to respond to that.")
                else:
                    for idx in res_indices:
                        word = self.idx_to_word.get(idx, "...")
                        for char in word:
                            print(char, end="", flush=True); time.sleep(0.01)
                        print(" ", end="", flush=True); time.sleep(0.02)
                    print()
                    
            except KeyboardInterrupt: 
                print("\n[系統] 收到中斷訊號。")
                self.save_model()
                break

if __name__ == "__main__":
    lanai = LanAI()
    lanai.chat()