app.py

import os
import io
import streamlit as st
import google.generativeai as genai
import pandas as pd
from streamlit_extras.colored_header import colored_header
from streamlit_extras.add_vertical_space import add_vertical_space

# --- 1. 기본 설정 및 AI 구성 ---
st.set_page_config(layout="wide", page_title="AI 학습 데이터 만들기")

# Gemini API 키 설정
try:
    genai.configure(api_key=st.secrets["GEMINI_API_KEY"])
except (KeyError, AttributeError):
    st.error("🚨 GEMINI_API_KEY를 설정해주세요! (Streamlit secrets에 추가)")
    st.stop()

# AI 모델 설정
generation_config = {
    "temperature": 0.3, # 약간의 창의성은 유지하되, 규칙을 더 잘 따르도록 온도를 살짝 낮춤
    "top_p": 0.95,
    "top_k": 40,
    "max_output_tokens": 4096,
    "response_mime_type": "text/plain",
}
model = genai.GenerativeModel(
    model_name="gemini-2.5-flash",
    generation_config=generation_config,
)

# --- 2. AI 프롬프트 (현실적인 데이터 생성을 위해 대폭 수정) ---
REALISTIC_DATA_PROMPT = """
당신은 현실적인 교육 데이터를 시뮬레이션하는 데이터 과학자 AI입니다.
당신의 임무는 사용자가 제공한 변수들의 현실적인 관계를 고려하여, 머신러닝 회귀 분석 학습에 적합한 데이터를 **Markdown 테이블 형식**으로 생성하는 것입니다.

**사용자 입력:**
*   원인 (X 변수) 이름: "{x_name}"
*   결과 (Y 변수) 이름: "{y_name}"
*   생성할 데이터 개수: {num_rows}

**수행할 작업 (매우 중요):**
1.  **현실적 관계 모델링:** 두 변수 간에 양의 상관관계가 있도록 데이터를 생성합니다. (X가 증가하면 Y도 대체로 증가)
2.  **현실적 변동성 추가:** 관계가 완벽한 직선이 아닌, 현실 데이터처럼 보이도록 약간의 무작위 변동성을 추가합니다.
3.  **현실적 제약 조건 적용:** 변수의 의미를 고려하여 상한선(Maximum)과 하한선(Minimum)을 자연스럽게 적용합니다.
    *   예를 들어, '{y_name}'이 '시험 점수'나 '만족도'와 관련 있다면 **결과값이 100을 넘지 않고, 0 미만이 되지 않도록** 데이터를 생성해야 합니다.
    *   '{x_name}'이 '공부 시간'이나 '운동 시간'이라면, 시간이 늘어날수록 '{y_name}'의 상승폭이 점차 둔화되는 **'수확 체감(diminishing returns)' 현상**을 반영하여 데이터를 더욱 현실적으로 만듭니다. (예: 1시간 공부해서 20점 오르지만, 9시간 공부해서는 5점 오르는 식)
4.  **출력 형식 준수:** 결과는 **오직 Markdown 테이블 형식**으로만 출력합니다.

**출력 형식 (절대 변경 금지):**
*   첫 줄은 헤더 `| {x_name} | {y_name} |` 입니다.
*   두 번째 줄은 구분선 `|---|---|` 입니다.
*   그 이후로는 `| 값1 | 값2 |` 형식의 데이터 행을 {num_rows}개 만큼 생성합니다.
*   설명, 코드, ``` 등 다른 어떤 텍스트도 포함하지 마세요.

**현실적인 데이터 출력 예시 (규칙을 잘 따르는 예시):**
| 공부 시간 | 시험 점수 |
|---|---|
| 0.5 | 45.8 |
| 1.2 | 62.1 |
| 2.0 | 78.2 |
| 3.5 | 89.5 |
| 5.0 | 94.7 |
| 6.5 | 98.1 |
| 1.8 | 75.3 |
| 4.2 | 92.0 |

이제 아래 정보를 바탕으로 현실적인 제약 조건을 따른 Markdown 테이블을 생성해주세요.
"""

# --- 3. 핵심 기능 함수 ---

def generate_markdown_data(x_name, y_name, num_rows):
    """AI를 호출하여 Markdown 테이블 형식의 데이터를 받아오는 함수"""
    prompt = REALISTIC_DATA_PROMPT.format( # 수정된 프롬프트 사용
        x_name=x_name,
        y_name=y_name,
        num_rows=num_rows
    )
    try:
        response = model.generate_content([prompt])
        return response.text.strip()
    except Exception as e:
        st.error(f"AI 호출 중 오류가 발생했습니다: {e}")
        return None

def parse_markdown_to_df(markdown_text):
    """Markdown 테이블 텍스트를 Pandas DataFrame으로 변환하는 함수"""
    try:
        lines = markdown_text.strip().split('\n')
        cleaned_lines = [line.strip().strip('|').strip() for line in lines if '---' not in line]
        
        header = [h.strip() for h in cleaned_lines[0].split('|')]
        data = [
            [cell.strip() for cell in row.split('|')]
            for row in cleaned_lines[1:]
        ]
        
        df = pd.DataFrame(data, columns=header)
        
        for col in df.columns:
            df[col] = pd.to_numeric(df[col], errors='coerce')
        
        df.dropna(inplace=True)
        return df

    except Exception as e:
        st.error(f"생성된 데이터를 테이블로 변환하는 중 오류가 발생했습니다: {e}")
        st.info("AI가 생성한 원본 데이터:")
        st.text(markdown_text)
        return None


# --- 4. Streamlit UI 구성 ---
colored_header(
    label="🤖 나만의 AI 학습 데이터 만들기 (회귀용)",
    description="AI가 여러분이 정한 규칙에 따라 '진짜 같은' 학습 데이터를 만들어줘요!",
    color_name="blue-70",
)
add_vertical_space(1)

# 세션 상태 초기화
if 'generated_df' not in st.session_state:
    st.session_state.generated_df = None
if 'markdown_output' not in st.session_state:
    st.session_state.markdown_output = ""

# 입력 영역
st.subheader("1. 데이터 규칙 정하기")
col1, col2, col3 = st.columns(3)

with col1:
    x_input_name = st.text_input(
        "**원인(X) 이름은?**",
        value="공부 시간",
        help="데이터의 원인이 되는 값의 이름을 정해주세요. (예: 공부 시간, 운동량)"
    )

with col2:
    y_input_name = st.text_input(
        "**결과(Y) 이름은?**",
        value="시험 점수",
        help="원인(X)에 따라 변하는 결과 값의 이름을 정해주세요. (예: 시험 점수, 만족도)"
    )

with col3:
    num_data_rows = st.number_input(
        "**데이터는 몇 개 만들까요?**",
        min_value=10,
        max_value=200,
        value=30,
        step=10,
        help="AI를 학습시키려면 데이터가 충분해야 해요. 10개 이상을 추천해요!"
    )

add_vertical_space(1)
generate_button = st.button("🚀 데이터 생성 시작!", type="primary", use_container_width=True)
add_vertical_space(2)


# 실행 및 결과 출력 영역
if generate_button:
    if not x_input_name or not y_input_name:
        st.warning("⚠️ '원인(X) 이름'과 '결과(Y) 이름'을 모두 입력해주세요!")
    else:
        with st.spinner("똑똑한 AI가 현실 세계의 규칙을 생각하며 데이터를 만들고 있어요... 🤖"):
            markdown_data = generate_markdown_data(x_input_name, y_input_name, num_data_rows)

            if markdown_data:
                df = parse_markdown_to_df(markdown_data)

                if df is not None:
                    st.session_state.generated_df = df
                    st.session_state.markdown_output = markdown_data
                    st.success("🎉 데이터 생성 완료! 아래에서 확인하고 다운로드하세요.")
                else:
                    st.error("데이터 생성에 실패했어요. AI가 만든 데이터를 분석할 수 없는 것 같아요.")
                    st.session_state.generated_df = None
            else:
                st.error("AI가 데이터를 생성하지 못했어요. 잠시 후 다시 시도해주세요.")
                st.session_state.generated_df = None

# 세션 상태에 저장된 데이터가 있으면 화면에 표시
if st.session_state.generated_df is not None:
    st.subheader("2. 생성된 데이터 미리보기 (Markdown)")
    st.markdown(st.session_state.markdown_output)

    # 데이터의 최대/최소값을 확인하여 현실성 검증
    try:
        y_col_name = st.session_state.generated_df.columns[1]
        max_val = st.session_state.generated_df[y_col_name].max()
        min_val = st.session_state.generated_df[y_col_name].min()
        st.info(f"**결과값('{y_col_name}')의 범위:** 최소 {min_val:.1f} ~ 최대 {max_val:.1f}")
    except (IndexError, TypeError):
        # DataFrame이 비어있거나 숫자형이 아닐 경우를 대비한 예외 처리
        pass


    st.subheader("3. CSV 파일로 다운로드하기")
    st.markdown("이 버튼을 눌러 위에 보이는 데이터를 CSV 파일로 컴퓨터에 저장하세요.")

    csv = st.session_state.generated_df.to_csv(index=False).encode('utf-8-sig')

    st.download_button(
        label="📥 CSV 파일 다운로드",
        data=csv,
        file_name=f"{x_input_name.replace(' ', '_')}_{y_input_name.replace(' ', '_')}_data.csv",
        mime="text/csv",
        use_container_width=True
    )

    with st.expander("👀 AI가 생성한 원본 Markdown 텍스트가 궁금하다면?"):
        st.text(st.session_state.markdown_output)