import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# 🔃 モデル・トークナイザーの読み込み（slow tokenizer指定）
@st.cache_resource
def load_model():
    tokenizer = AutoTokenizer.from_pretrained(
        "rinna/japanese-gpt-neox-3.6b", use_fast=False
    )
    model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt-neox-3.6b")
    return tokenizer, model

tokenizer, model = load_model()

# 📋 サンプルデータ（4頭分）
cases = [
    {
        "user_input": """1歳時馬体重：430kg
体高：155.0cm
胸囲：180.0cm
管囲：20.8cm
生年月日：2023-03-15
測尺日：2024-06-10""",
        "result_summary": """予測馬体重：480kg、ベストな馬体重
標準管囲との乖離：太い
過去の名馬：トーホウジャッカル 菊花賞G1芝3000m、オーソリティ アルゼンチン共和国杯G2芝2500m"""
    },
    {
        "user_input": """1歳時馬体重：500kg
体高：152.0cm
胸囲：175.0cm
管囲：18.0cm
生年月日：2023-02-20
測尺日：2024-06-20""",
        "result_summary": """予測馬体重：540kg、やや重め
標準管囲との乖離：細い
過去の名馬：サクセスブロッケン フェブラリーS G1ダ1600m、スマートファルコン 東京大賞典G1ダ2000m"""
    },
    {
        "user_input": """1歳時馬体重：380kg
体高：158.0cm
胸囲：172.0cm
管囲：20.2cm
生年月日：2023-04-10
測尺日：2024-06-01""",
        "result_summary": """予測馬体重：455kg、成長余地あり
標準管囲との乖離：標準
過去の名馬：ダノンザキッド ホープフルS G1芝2000m、ジャスタウェイ 安田記念G1芝1600m"""
    },
    {
        "user_input": """1歳時馬体重：470kg
体高：164.5cm
胸囲：179.0cm
管囲：21.5cm
生年月日：2023-03-05
測尺日：2024-07-01""",
        "result_summary": """予測馬体重：510kg、完成度高め
標準管囲との乖離：非常に太い
過去の名馬：キズナ 東京優駿G1芝2400m、ハーツクライ 有馬記念G1芝2500m"""
    }
]

# ✨ コメント生成関数
def generate_comment(case):
    prompt = f"""以下は、一口馬主の募集馬についてのAI総括コメントです。

ユーザーが入力した項目：
{case["user_input"]}

予測した結果：
{case["result_summary"]}

この馬について得られた結果を総括したうえで、ポジティブな言葉で締めくくってください。
コメント："""

    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    output = model.generate(input_ids, max_new_tokens=120, do_sample=True, temperature=0.7)
    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
    return decoded.split("コメント：")[-1].strip()

# 🖼️ UI部分
st.title("🐴 AIによる募集馬コメント生成デモ")

if st.button("🎯 コメントを一括生成"):
    for i, case in enumerate(cases, 1):
        with st.spinner(f"Case {i} を生成中..."):
            comment = generate_comment(case)
            st.markdown(f"### 🐎 Case {i}")
            st.markdown(f"**📥 入力**\n```\n{case['user_input']}\n```\n")
            st.markdown(f"**📊 予測結果**\n```\n{case['result_summary']}\n```\n")
            st.markdown(f"**📝 コメント**\n> {comment}")
            st.markdown("---")