"""
zmierz-model.py - ローカルモデルの速度、GPUメモリ、電気代を測定（記事URL）
作者：Piotr Olszewski。スクリプトのライセンス：MIT。2026年9月26日確認：llama.cpp b11200、RTX 4090。

OpenAI API互換サーバーで動作します：llama-server（ポート8080）、LM Studio（1234）、Ollama（11434）。
元のポーランド語テスト用プロンプトを5件送り、トークン/秒、GPUメモリの最大使用量、平均消費電力（nvidia-smi）を測定し、
トークン100万個あたりの電気代を計算します。Python 3.9以降の標準ライブラリのみを使用します。

実行例：
  python zmierz-model.py                                   （127.0.0.1:8080のllama-server）
  python zmierz-model.py --url http://127.0.0.1:1234 --model gemma-4-12b      （LM Studio）
  python zmierz-model.py --url http://127.0.0.1:11434 --model gemma4:12b      （Ollama）
  python zmierz-model.py --cena 0.27 --reszta-w 100        （電気料金：ドル/kWh、PC本体分の加算電力）
"""
import argparse, json, statistics, subprocess, sys, threading, time, urllib.error, urllib.request

# 以下は、実際のテストに使ったポーランド語プロンプトの日本語訳です。
# 日本語で再テストしたプロンプトではありません。
POLECENIA = [
    "【ポーランド語で行ったテストのプロンプトの訳】荷物の到着が5日遅れている顧客からの苦情に、丁寧な返信を書いてください。120語以内。",
    "【ポーランド語で行ったテストのプロンプトの訳】12歳の子どもに、インフレーションとは何かを説明してください。100～130語で、例を1つ挙げてください。",
    "【ポーランド語で行ったテストのプロンプトの訳】次の内容を5つの箇条書きで要約してください：経営陣は、アプリのリリースを1か月延期し、テスト予算を20%増やし、テスターを2人採用することを決定した。",
    "【ポーランド語で行ったテストのプロンプトの訳】集合住宅での節電について、記事のタイトルを10個提案してください。それぞれ60文字以内。",
    "【ポーランド語で行ったテストのプロンプトの訳】次の英文を、直訳調を避けてポーランド語に翻訳してください：『新製品のリリースを発表できて、とてもうれしく思います。実現まで長い時間がかかりました。』",
]


class Czujnik(threading.Thread):
    def __init__(self):
        super().__init__(daemon=True); self.stop = False; self.vram = []; self.moc = []
    def run(self):
        while not self.stop:
            try:
                o = subprocess.check_output(["nvidia-smi", "--query-gpu=memory.used,power.draw", "--format=csv,noheader,nounits"], text=True, timeout=5)
                v, w = [float(x) for x in o.strip().splitlines()[0].split(",")]
                self.vram.append(v); self.moc.append(w)
            except Exception:
                return
            time.sleep(0.5)


def main():
    a = argparse.ArgumentParser(description="ローカルモデルの測定：トークン/秒、GPUメモリ、電気代")
    a.add_argument("--url", default="http://127.0.0.1:8080", help="OpenAI API互換サーバーのURL")
    a.add_argument("--model", default="local", help="モデル名（LM StudioとOllamaでは指定が必要）")
    a.add_argument("--cena", type=float, default=1.05 / 3.857, help="1kWhあたりの電気料金（ドル、デフォルトは記事の換算値）")
    a.add_argument("--reszta-w", type=float, default=0, help="PC本体分として加算する電力（例：100W）")
    a.add_argument("--max-tokens", type=int, default=800)
    x = a.parse_args()
    try:
        o = subprocess.check_output(["nvidia-smi", "--query-gpu=name,memory.used", "--format=csv,noheader,nounits"], text=True, timeout=5)
        karta, tlo = o.strip().splitlines()[0].rsplit(",", 1); tlo = float(tlo)
        print(f"GPU：{karta.strip()}、テスト開始前の使用メモリ：{tlo:.0f}MB")
    except Exception:
        karta, tlo = None, 0
        print("nvidia-smiが見つかりません。速度のみ測定し、メモリと電力は測定しません。")
    # llama-serverはモデルの読み込み中に503を返すため、最大5分待ちます。
    for proba in range(150):
        try:
            urllib.request.urlopen(x.url.rstrip("/") + "/v1/models", timeout=5); break
        except urllib.error.HTTPError as e:
            if e.code != 503: break
        except Exception:
            pass
        if proba == 0: print("サーバーがモデルを読み込むまで待機しています...")
        time.sleep(2)
    wyniki = []
    for i, p in enumerate(POLECENIA, 1):
        body = {"model": x.model, "messages": [{"role": "user", "content": p}], "max_tokens": x.max_tokens, "temperature": 0.7}
        cz = Czujnik(); cz.start(); t0 = time.time()
        try:
            req = urllib.request.Request(x.url.rstrip("/") + "/v1/chat/completions", data=json.dumps(body).encode(), headers={"Content-Type": "application/json"})
            j = json.load(urllib.request.urlopen(req, timeout=900))
        except Exception as e:
            sys.exit(f"サーバー{x.url}に接続できません（{e}）。モデルが起動しているか確認してください。")
        czas = time.time() - t0; cz.stop = True; cz.join()
        tok = j.get("usage", {}).get("completion_tokens") or 0
        tps = j.get("timings", {}).get("predicted_per_second") or (tok / czas if czas else 0)
        moc = statistics.mean(cz.moc) if cz.moc else 0
        wyniki.append({"tok": tok, "czas": czas, "tps": tps, "vram": max(cz.vram or [0]), "moc": moc})
        tekst = (j["choices"][0]["message"].get("content") or "").strip().replace("\n", " ")
        print(f"[{i}/5] {tok}トークン、{czas:.1f}秒、{tps:.1f}トークン/秒、{moc:.0f}W | {tekst[:70]}...")
    tok = sum(w["tok"] for w in wyniki); czas = sum(w["czas"] for w in wyniki)
    print("\n=== 集計 ===")
    print(f"平均生成速度：{statistics.mean(w['tps'] for w in wyniki):.1f}トークン/秒")
    if karta:
        wh = sum(w["czas"] * (w["moc"] + x.reszta_w) for w in wyniki) / 3600
        print(f"GPUメモリ：最大{max(w['vram'] for w in wyniki):.0f}MB、モデル使用分は約{max(w['vram'] for w in wyniki) - tlo:.0f}MB")
        print(f"エネルギー：{tok}トークンあたり{wh:.2f}Wh")
        if tok:
            print(f"電気代：テスト全体で{wh / 1000 * x.cena:.6f}ドル、トークン100万個あたり{wh / 1000 * x.cena / tok * 1e6:.4f}ドル "
                  f"（料金：{x.cena}ドル/kWh{', PC本体分として+' + format(x.reszta_w, '.0f') + 'W' if x.reszta_w else '、GPUのみ'}）")


if __name__ == "__main__":
    main()