"""
04_autograd_from_scratch.py
================================
手刻一個「自動微分（autograd）」引擎，再用它訓練一個小神經網路解 XOR。
這補上 docs/01 講的「反向傳播 backprop」——PyTorch 的 .backward() 底層就是這個想法。

核心觀念：
  - 每個數值都用一個 Value 物件包起來，記住「自己是怎麼被算出來的」（建立計算圖）。
  - forward：照算式往前算出答案。
  - backward：用鏈鎖律，從輸出往回，把梯度傳回每個參數（這就是反向傳播）。
  - 有了梯度，就能做梯度下降 θ ← θ − η·∇L 來訓練。

執行：python3 04_autograd_from_scratch.py   （純 Python，不需 numpy、不需 GPU）
參考：Andrej Karpathy 的 micrograd。
"""

import math
import random

random.seed(1337)


class Value:
    """一個純量，外加它的梯度與「怎麼算出來的」紀錄。"""

    def __init__(self, data, _children=(), _op=""):
        self.data = data
        self.grad = 0.0
        self._backward = lambda: None      # 這個節點把梯度往回傳的方法
        self._prev = set(_children)
        self._op = _op

    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other), "+")
        def _backward():               # d(out)/d(self)=1, d(out)/d(other)=1
            self.grad += out.grad
            other.grad += out.grad
        out._backward = _backward
        return out

    def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data * other.data, (self, other), "*")
        def _backward():               # 乘法的微分：互相是對方的係數
            self.grad += other.data * out.grad
            other.grad += self.data * out.grad
        out._backward = _backward
        return out

    def __pow__(self, p):              # 只支援常數次方
        assert isinstance(p, (int, float))
        out = Value(self.data ** p, (self,), f"**{p}")
        def _backward():
            self.grad += (p * self.data ** (p - 1)) * out.grad
        out._backward = _backward
        return out

    def tanh(self):                    # 非線性激活函數
        t = math.tanh(self.data)
        out = Value(t, (self,), "tanh")
        def _backward():
            self.grad += (1 - t * t) * out.grad   # tanh 的導數 = 1 - tanh²
        out._backward = _backward
        return out

    # 一些方便的運算
    def __neg__(self): return self * -1
    def __sub__(self, other): return self + (-other)
    def __radd__(self, other): return self + other
    def __rmul__(self, other): return self * other
    def __truediv__(self, other): return self * other ** -1

    def backward(self):
        """反向傳播：先把節點做拓撲排序，再從輸出往回呼叫每個 _backward。"""
        topo, visited = [], set()
        def build(v):
            if v not in visited:
                visited.add(v)
                for child in v._prev:
                    build(child)
                topo.append(v)
        build(self)
        self.grad = 1.0                # 輸出對自己的梯度 = 1
        for v in reversed(topo):
            v._backward()


# ---------- 用上面的引擎組一個小神經網路 ----------
class Neuron:
    def __init__(self, n_in):
        self.w = [Value(random.uniform(-1, 1)) for _ in range(n_in)]
        self.b = Value(0.0)

    def __call__(self, x):
        act = sum((wi * xi for wi, xi in zip(self.w, x)), self.b)
        return act.tanh()

    def parameters(self):
        return self.w + [self.b]


class Layer:
    def __init__(self, n_in, n_out):
        self.neurons = [Neuron(n_in) for _ in range(n_out)]

    def __call__(self, x):
        outs = [n(x) for n in self.neurons]
        return outs[0] if len(outs) == 1 else outs

    def parameters(self):
        return [p for n in self.neurons for p in n.parameters()]


class MLP:
    def __init__(self, n_in, sizes):
        dims = [n_in] + sizes
        self.layers = [Layer(dims[i], dims[i + 1]) for i in range(len(sizes))]

    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

    def parameters(self):
        return [p for layer in self.layers for p in layer.parameters()]


def main():
    # XOR 問題：輸入兩個 0/1，輸出它們是否「不同」。用 +1/-1 當標籤配合 tanh。
    X = [[0, 0], [0, 1], [1, 0], [1, 1]]
    Y = [-1.0, 1.0, 1.0, -1.0]         # XOR：相同→-1，不同→+1

    model = MLP(2, [4, 4, 1])          # 2 輸入 → 兩層各 4 神經元 → 1 輸出
    n_params = len(model.parameters())
    print(f"模型參數量：{n_params}（每個都是一個會記梯度的 Value）\n開始訓練 XOR ...\n")

    lr = 0.1
    for epoch in range(200):
        # forward：算預測與損失（均方誤差）
        preds = [model(x) for x in X]
        loss = sum((p - y) ** 2 for p, y in zip(preds, Y))

        # backward：先把梯度歸零，再反向傳播
        for p in model.parameters():
            p.grad = 0.0
        loss.backward()

        # 梯度下降：沿梯度反方向更新參數
        for p in model.parameters():
            p.data -= lr * p.grad

        if epoch % 40 == 0 or epoch == 199:
            print(f"epoch {epoch:3d} | loss = {loss.data:.4f}")

    print("\n訓練後預測（target：相同→-1，不同→+1）：")
    for x, y in zip(X, Y):
        pred = model(x).data
        ok = "✓" if (pred > 0) == (y > 0) else "✗"
        print(f"  輸入 {x} → 預測 {pred:+.2f}（應為 {y:+.0f}）{ok}")
    print("\n重點：上面那條 loss.backward() 就是反向傳播；PyTorch 幫你把這套做到大規模、跑在 GPU 上。")


if __name__ == "__main__":
    main()

"""
07_lora_finetune_reference.py   ★ 參考腳本（需要 GPU，請在你自己的機器跑）★
======================================================================
docs/00 第 4 站的核心實戰：用 LoRA 微調一個開源 LLM。
這支「不會」在本沙盒跑（沒有 GPU、沒有網路、沒裝套件），是給你照著在自己環境執行的範本。

為什麼用 LoRA：完整微調一個大模型要很多顯存；LoRA 只在模型裡插入少量可訓練的小矩陣，
凍結原權重，就能用一張消費級顯卡微調出自己的模型。再加 4-bit 量化（QLoRA）可微調更大的模型。

安裝（在你自己的環境，建議 Python 3.10+、有 NVIDIA GPU）：
    pip install "transformers>=4.44" "trl>=0.9" "peft>=0.12" "datasets" "accelerate" "bitsandbytes"

執行：
    python3 07_lora_finetune_reference.py

⚠ 注意：transformers / trl / peft 的 API 改版頻繁，若報錯請對照當下官方文件微調參數。
   建議先用很小的模型（如 0.5B）在小資料上跑通，再換大模型 / 大資料。
"""

import torch
from datasets import Dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig
from trl import SFTConfig, SFTTrainer

# ---------------- 1. 選一個小的開源基礎模型 ----------------
# 建議先用小模型把流程跑通（你做教材/股票分析的資料正好可拿來微調）。
MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"   # 也可換 TinyLlama、Llama-3.2-1B 等

# ---------------- 2. 準備資料 ----------------
# 真實情況可從 JSONL 或 Hugging Face datasets 載入。這裡用幾筆示意，
# 格式是「指令 → 回應」的對話樣本（SFT 監督式微調）。
raw = [
    {"instruction": "用一句話解釋什麼是注意力機制。",
     "response": "注意力讓模型在處理每個字時，依相關度去加權參考其他字的資訊。"},
    {"instruction": "什麼是過擬合？",
     "response": "模型把訓練資料背得太死，在沒看過的資料上反而表現變差。"},
    # ... 你的真實資料放這裡（越多越好，幾百~幾千筆起跳）
]

def to_text(example):
    # 用模型的對話模板把「指令+回應」組成一段訓練文字
    messages = [
        {"role": "user", "content": example["instruction"]},
        {"role": "assistant", "content": example["response"]},
    ]
    return {"text": tokenizer.apply_chat_template(messages, tokenize=False)}

# ---------------- 3. 載入 tokenizer 與模型（可選 4-bit 量化 = QLoRA）----------------
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

USE_4BIT = True   # 顯存不夠就設 True（QLoRA）；顯存充足可設 False
quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
) if USE_4BIT else None

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    quantization_config=quant_config,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

dataset = Dataset.from_list(raw).map(to_text)

# ---------------- 4. 設定 LoRA ----------------
# r 越大可學的東西越多但也越吃資源；target_modules 是要插入 LoRA 的層。
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],  # 注意力的投影層
)

# ---------------- 5. 訓練（TRL 的 SFTTrainer 會自動套用 LoRA）----------------
sft_config = SFTConfig(
    output_dir="./lora-out",
    num_train_epochs=3,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    logging_steps=10,
    save_strategy="epoch",
    bf16=True,
    max_seq_length=512,
    dataset_text_field="text",
)

trainer = SFTTrainer(
    model=model,
    args=sft_config,
    train_dataset=dataset,
    peft_config=lora_config,
    processing_class=tokenizer,   # 舊版本可能叫 tokenizer=，依你的版本調整
)

if __name__ == "__main__":
    trainer.train()
    trainer.save_model("./lora-out")        # 只會存下很小的 LoRA adapter
    print("✅ 訓練完成，LoRA adapter 已存到 ./lora-out")

    # ---------------- 6. 用微調後的模型推論 ----------------
    prompt = "用一句話解釋什麼是過擬合。"
    messages = [{"role": "user", "content": prompt}]
    inputs = tokenizer.apply_chat_template(messages, return_tensors="pt",
                                           add_generation_prompt=True).to(model.device)
    out = model.generate(inputs, max_new_tokens=80)
    print(tokenizer.decode(out[0], skip_special_tokens=True))

# ----------------------------------------------------------------------
# 進階方向：
#   - 把它部署成 API：用 FastAPI 包起來，或用 vLLM 加速推論。
#   - 評估（evals）：準備一組測試題，比較微調前後的回答品質。
#   - 合併權重：peft 的 merge_and_unload() 可把 LoRA 併回基礎模型成為獨立模型。
# 這些做完，你就走完了 docs/00 的第 4 站「資料 → 微調 → 評估 → 部署」。
# ----------------------------------------------------------------------