import math
import random
import time
import re
import json
import os

# --- 基礎數學與激活函數 ---
def sigmoid(x):
    return 1 / (1 + math.exp(-max(min(x, 20), -20)))

def tanh(x):
    return math.tanh(max(min(x, 20), -20))

# --- LanAI 0.4.0 核心架構 (GRU) ---
class LanAIGRU:
    def __init__(self, vocab_size, hidden_size=160): 
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        
        limit = math.sqrt(6 / (vocab_size + hidden_size))
        def init_matrix(r, c): return [[random.uniform(-limit, limit) for _ in range(c)] for _ in range(r)]
        
        self.W_r = init_matrix(vocab_size, hidden_size)
        self.U_r = init_matrix(hidden_size, hidden_size)
        self.b_r = [0.0] * hidden_size
        
        self.W_z = init_matrix(vocab_size, hidden_size)
        self.U_z = init_matrix(hidden_size, hidden_size)
        self.b_z = [0.0] * hidden_size
        
        self.W_h = init_matrix(vocab_size, hidden_size)
        self.U_h = init_matrix(hidden_size, hidden_size)
        self.b_h = [0.0] * hidden_size
        
        # W_y 維度: [hidden_size, vocab_size]
        self.W_y = init_matrix(hidden_size, vocab_size)
        self.b_y = [0.0] * vocab_size

    def expand_weights(self, new_vocab_size):
        """權重動態擴展機制：當詞彙庫增加時，動態擴充網路矩陣維度"""
        diff = new_vocab_size - self.vocab_size
        if diff <= 0: return
        
        limit = math.sqrt(6 / (new_vocab_size + self.hidden_size))
        
        # 擴充輸入層的詞彙列 (W_r, W_z, W_h) -> 增加 diff 列
        for _ in range(diff):
            self.W_r.append([random.uniform(-limit, limit) for _ in range(self.hidden_size)])
            self.W_z.append([random.uniform(-limit, limit) for _ in range(self.hidden_size)])
            self.W_h.append([random.uniform(-limit, limit) for _ in range(self.hidden_size)])
            
        # 擴充輸出層的詞彙行 (W_y, b_y)
        for i in range(self.hidden_size):
            self.W_y[i].extend([random.uniform(-limit, limit) for _ in range(diff)])
        self.b_y.extend([0.0] * diff)
        
        self.vocab_size = new_vocab_size

    def forward(self, inputs, h_prev=None):
        h = h_prev[:] if h_prev else [0.0] * self.hidden_size
        for idx in inputs:
            idx = idx if idx < self.vocab_size else 0
            wr_row = self.W_r[idx]; wz_row = self.W_z[idx]; wh_row = self.W_h[idx]
            r = [sigmoid(wr_row[j] + sum(h[k] * self.U_r[k][j] for k in range(self.hidden_size)) + self.b_r[j]) for j in range(self.hidden_size)]
            z = [sigmoid(wz_row[j] + sum(h[k] * self.U_z[k][j] for k in range(self.hidden_size)) + self.b_z[j]) for j in range(self.hidden_size)]
            h_tilde = [tanh(wh_row[j] + sum((r[k] * h[k]) * self.U_h[k][j] for k in range(self.hidden_size)) + self.b_h[j]) for j in range(self.hidden_size)]
            h = [(1 - z[j]) * h[j] + z[j] * h_tilde[j] for j in range(self.hidden_size)]
        y = [sum(h[k] * self.W_y[k][j] for k in range(self.hidden_size)) + self.b_y[j] for j in range(self.vocab_size)]
        return y, h

    def train_sequence(self, input_indices, target_indices, lr=0.08):
        h_state = None
        current_input = input_indices
        total_loss = 0
        for target_word_idx in target_indices:
            y, h_state = self.forward(current_input, h_state)
            max_y = max(y)
            exp_y = [math.exp(v - max_y) for v in y]
            sum_exp = sum(exp_y) + 1e-9
            probs = [e / sum_exp for e in exp_y]
            dy = probs[:]
            dy[target_word_idx] -= 1.0
            for i in range(self.hidden_size):
                h_val = h_state[i]
                for j in range(self.vocab_size):
                    self.W_y[i][j] -= lr * dy[j] * h_val
            total_loss -= math.log(max(probs[target_word_idx], 1e-10))
            current_input = [target_word_idx]
        return total_loss

# --- LanAI 0.4.0 主系統 ---
class LanAI:
    def __init__(self, model_path="lanai_memory.json"):
        self.model_path = model_path
        self.context_h = None 
        
        # [更新 0.4.0] 擴充國三 9 年級基礎訓練資料集 (包含文法結構)
        self.base_training_data = [
            (["hello"], ["hi", "how", "are", "you"]),
            (["what", "is", "your", "name"], ["i", "am", "lanai", "version", "0.4.0"]),
            (["what", "can", "you", "do"], ["i", "can", "calculate", "and", "learn", "new", "grammar"]),
            (["thanks"], ["you", "are", "welcome"]),
            (["bye"], ["goodbye", "see", "you"]),
            # --- 國三文法: 現在完成式 ---
            (["have", "you", "eaten", "yet"], ["yes", "i", "have", "already", "eaten"]),
            (["what", "have", "you", "done"], ["i", "have", "learned", "new", "things", "recently"]),
            (["she", "has", "been", "there"], ["since", "last", "year"]),
            # --- 國三文法: 從屬連接詞 & 條件句 ---
            (["although", "it", "is", "hard"], ["i", "will", "never", "give", "up"]),
            (["because", "it", "is", "raining"], ["we", "cannot", "go", "out", "to", "play"]),
            (["if", "it", "rains", "tomorrow"], ["we", "will", "stay", "at", "home"]),
            (["unless", "you", "study", "hard"], ["you", "will", "not", "pass", "the", "test"]),
            # --- 國三文法: too...to / so...that ---
            (["he", "is", "too", "tired"], ["to", "walk", "any", "farther"]),
            (["it", "is", "so", "beautiful"], ["that", "everyone", "likes", "it"]),
            # --- 國三文法: 被動語態 ---
            (["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"]),
            (["the", "apple", "which", "is", "red"], ["looks", "very", "delicious"])
        ]
        self.training_data = list(self.base_training_data)
        
        # [更新 0.4.0] 擴充國三核心單字庫
        self.default_vocab = [
            "<PAD>", "<EOS>", "hello", "hi", "how", "are", "you", "i", "am", "is", "it", "they", "we",
            "fine", "good", "great", "well", "happy", "sad", "fast", "slow", "smart", "cool", "tired",
            "what", "your", "name", "lanai", "ai", "who", "creator", "human", "thanks", "welcome",
            "bye", "see", "later", "can", "do", "help", "think", "understand", "talk", "learn",
            "math", "solve", "calculate", "plus", "minus", "times", "divided", "by", "equals",
            "a", "the", "an", "this", "that", "my", "working", "system", "version", "0.4.0",
            # --- 國三新增單字: 完成式動詞與助動詞 ---
            "have", "has", "had", "been", "done", "finished", "written", "eaten", "seen", "gone", "played",
            # --- 國三新增單字: 連接詞與副詞 ---
            "because", "although", "though", "however", "therefore", "if", "unless",
            "already", "yet", "since", "recently", "still", "never",
            # --- 國三新增單字: 關係代名詞 ---
            "whom", "whose", "which", "where", "when", "why",
            # --- 國三新增單字: 程度副詞 ---
            "too", "to", "so", "such", "enough", "very",
            # --- 國三新增單字: 核心名詞與形容詞 ---
            "experience", "knowledge", "foreign", "language", "successful", "environment", 
            "pollution", "protect", "serious", "consider", "depend", "famous", "writer",
            "careful", "dangerous", "exciting", "interesting", "surprised", "boring", "bored",
            "grammar", "tense", "passive", "voice", "relative", "clause", "subject", "object",
            "boy", "basketball", "house", "lived", "before", "with", "family", "apple", "red", "looks", "delicious",
            "raining", "cannot", "go", "out", "play", "tomorrow", "stay", "at", "home", "study", "hard", "not", "pass", "test",
            "walk", "any", "farther", "beautiful", "everyone", "likes", "action", "new", "things", "she", "there", "last", "year", "give", "up"
        ]

        # 嘗試載入記憶，若無則初始化
        if not self.load_model():
            print("[系統] 找不到既有記憶，初始化 0.4.0 國三擴充模型...")
            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 = LanAIGRU(len(self.vocab))
            self.is_new_model = True
        else:
            self.is_new_model = False
            self.merge_new_knowledge() # 自動將新版文法單字寫入舊記憶

    def merge_new_knowledge(self):
        """[新增 0.4.0] 智慧合併：將新版的國三詞彙與文法自動擴充到使用者的舊記憶中"""
        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:
            print(f"[系統升級] 偵測到舊版記憶，已成功注入 {new_words_count} 個國三程度新詞彙與文法結構！")
            self.model.expand_weights(len(self.vocab))
            # 將新版基礎訓練資料加進去進行微調準備
            self.training_data.extend(self.base_training_data)
            self.save_model() # 儲存擴充後的記憶

    def save_model(self):
        """持久化記憶：將模型權重與詞彙表存入硬碟"""
        print("[系統] 正在儲存記憶到磁碟...")
        data = {
            "vocab": self.vocab,
            "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
        }
        with open(self.model_path, 'w', encoding='utf-8') as f:
            json.dump(data, f)
        print("[系統] 記憶儲存成功！")

    def load_model(self):
        """持久化記憶：從硬碟讀取模型"""
        if os.path.exists(self.model_path):
            try:
                print("[系統] 偵測到既有記憶，正在喚醒...")
                with open(self.model_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                self.vocab = data["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 = LanAIGRU(len(self.vocab))
                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",
                   "six": "6", "seven": "7", "eight": "8", "nine": "9", "ten": "10"}
        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 "+-*/"):
            try:
                res = eval(expr, {"__builtins__": None}, {})
                return f"The answer is {res}."
            except:
                return None
        return None

    def train_epochs(self, target_epochs, initial_lr=0.05):
        """引入 Learning Rate Decay"""
        for epoch in range(target_epochs):
            random.shuffle(self.training_data)
            loss = 0
            current_lr = initial_lr * (1 - (epoch / (target_epochs * 1.5))) 
            
            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 % 100 == 0:
                print(f"訓練進度 - Epoch: {epoch} | Loss: {loss/len(self.training_data):.4f} | LR: {current_lr:.4f}", end="\r")

    def generate_response(self, input_indices, temperature=0.6):
        """輸出層取樣策略：引入 Temperature 係數"""
        res_indices = []
        h = self.context_h
        curr_input = input_indices
        for _ in range(15):
            y, h = self.model.forward(curr_input, h)
            
            logits = [v / temperature for v in y]
            max_l = max(logits)
            exp_l = [math.exp(l - max_l) for l in logits]
            sum_exp = sum(exp_l)
            probs = [e / sum_exp for e in exp_l]
            
            r = random.random()
            cum = 0
            best_idx = 0
            for i, p in enumerate(probs):
                cum += p
                if r <= cum:
                    best_idx = i
                    break
            
            if best_idx == self.word_to_idx["<EOS>"] or best_idx == 0: 
                break
            res_indices.append(best_idx)
            curr_input = [best_idx]
            
        self.context_h = h
        return res_indices

    def thinking_layer(self, user_input, words):
        """思考預處理層：判斷意圖與提取數學/文法邏輯"""
        print("» [系統思考中] 分析語境與意圖...", end="\r")
        time.sleep(0.5) 
        
        math_res = self.solve_math(user_input)
        
        logic_keywords = ["why", "how", "can", "what", "math", "solve", "calculate", "think", "who"]
        grammar_keywords = ["grammar", "tense", "passive", "active", "voice", "sentence", "clause", "relative"]
        
        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"
            
        print(" " * 40, end="\r") 
        if math_res:
            print("» [思考結果] 偵測到明確數學算式。")
        elif intent == "grammar":
            print("» [思考結果] 偵測到文法探討，啟用結構性回答模式...")
        elif intent == "logic":
            print("» [思考結果] 偵測到邏輯需求，準備強化推論與訓練...")
        else:
            print("» [思考結果] 一般閒聊模式。")
            
        return math_res, intent

    def chat(self):
        print("="*60)
        print("    LanAI 0.4.0 (Grammar & Vocab Extended | Dynamic GRU)    ")
        print("="*60)
        
        if self.is_new_model:
            print("[系統] 模型為初始狀態，執行基礎認知與文法訓練 (1200 Epochs)...")
            # 由於加了大量文法，初始訓練加強到 1200 Epochs
            self.train_epochs(1200, initial_lr=0.08)
            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 intent == "grammar":
                    print("[動態微調] 文法結構強化訓練中...", end="\r")
                    if new_words_learned:
                        self.training_data.append((words, ["this", "is", "a", "grammar", "rule"]))
                    self.train_epochs(80, initial_lr=0.03)
                elif intent == "logic":
                    print("[動態微調] 邏輯強化深度訓練中...", end="\r")
                    if new_words_learned:
                        self.training_data.append((words, ["i", "think", "and", "learn"] + words))
                    self.train_epochs(100, initial_lr=0.04)
                elif new_words_learned:
                    print("[動態微調] 吸收新知識中...", end="\r")
                    self.training_data.append((words, ["i", "learned", "about"] + words))
                    self.train_epochs(50, initial_lr=0.03)
                else:
                    self.train_epochs(5, initial_lr=0.01)
                
                indices = [self.word_to_idx.get(w, 0) for w in words]
                
                # 文法模式降低 Temperature 確保回答精準度
                temp = 0.2 if intent in ["logic", "grammar"] else 0.7
                res_indices = self.generate_response(indices, temperature=temp)
                
                print("AI: ", end="", flush=True)
                if not res_indices:
                    print("I am processing new information. Tell me more.")
                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()