import streamlit as st
import torch
from peft import PeftConfig, PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
import time
import re
from typing import Tuple
from datasets import load_dataset
import psutil

st.set_page_config(page_title="LG AIMERS EXAONE 4.0 GSM8K", layout="wide")


@st.cache_resource
def load_model():
    st.info("🚀 최신 LoRA 로딩 중...")
    
    REPO_ID = "paddacoco/exaone4_gsm8k2"
    SUBFOLDER = "exaone4_gsm8k/final_exaone4_gsm8k"
    
    try:
        # 1. LoRA 설정 로드
        config = PeftConfig.from_pretrained(REPO_ID, subfolder=SUBFOLDER)
        st.write(f"✅ LoRA r={config.r}, target_modules={config.target_modules}")
        
        # 2. EXAONE 4.0 베이스
        base_model = AutoModelForCausalLM.from_pretrained(
            "LGAI-EXAONE/EXAONE-4.0-1.2B",
            torch_dtype=torch.float32,
            device_map="cpu",
            trust_remote_code=True
        )
        
        # 3. LoRA 어댑터 결합
        model = PeftModel.from_pretrained(base_model, REPO_ID, subfolder=SUBFOLDER)
        
        # 4. 토크나이저 (LoRA에서)
        tokenizer = AutoTokenizer.from_pretrained(REPO_ID, subfolder=SUBFOLDER, trust_remote_code=True)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
            
        st.success("🎉 LG AIMERS LoRA 완전 로드!")
        st.metric("LoRA 크기", "43.7MB")
        st.metric("베이스", "EXAONE 4.0 1.2B")
        return model, tokenizer
        
    except Exception as e:
        st.error(f"❌ {e}")
        return None, None


def extract_math_answer(text: str) -> str:
    """수학 답변 추출"""
    if "####" in text:
        answer_part = text.split("####")[-1].strip()
        numbers = re.findall(r"\d+", answer_part)
        return numbers[0] if numbers else answer_part[:50]
    numbers = re.findall(r"\d+", text)
    return numbers[-1] if numbers else "답 없음"


def generate_response(model, tokenizer, prompt: str, max_length: int = 256, 
                     temperature: float = 0.7, top_p: float = 0.9) -> Tuple[str, float]:
    """모델 응답 생성"""
    if model is None or tokenizer is None:
        return "모델 로드 실패", 0.0
    
    start_time = time.time()
    
    try:
        # 수학 모드 최적화 프롬프트
        if "수학" in st.session_state.get("mode", ""):
            system_prompt = """수학 문제를 단계별로 정확히 풀어주세요.
단위와 계산 과정을 명확히 설명하세요.
최종 답은 #### 뒤에 숫자만 적으세요."""
            full_prompt = f"{system_prompt}\n\n문제: {prompt}"
        else:
            full_prompt = prompt
        
        chat = [{'role': 'user', 'content': full_prompt}]
        formatted = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
        
        inputs = tokenizer(formatted, return_tensors="pt", truncation=True, max_length=1024)
        
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=max_length,
                temperature=temperature,
                top_p=top_p,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
                repetition_penalty=1.1
            )
        
        response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
        return response.strip(), time.time() - start_time
    
    except Exception as e:
        return f"생성 에러: {str(e)}", 0.0


@st.cache_data
def evaluate_gsm8k(n_samples=100):
    """GSM8K 벤치마크 평가 (해커톤용)"""
    try:
        st.info(f"📊 GSM8K {n_samples}개 샘플 평가 중... (시간 걸림)")
        dataset = load_dataset("openai/gsm8k", "main")['test'].select(range(min(n_samples, 100)))
        
        correct = 0
        total_time = 0
        results = []
        
        progress_bar = st.progress(0)
        status_text = st.empty()
        
        for i, example in enumerate(dataset):
            prompt = example['question']
            start = time.time()
            response, _ = generate_response(st.session_state.model, st.session_state.tokenizer, prompt, max_length=128)
            elapsed = time.time() - start
            total_time += elapsed
            
            # 간단 숫자 추출
            pred = extract_math_answer(response)
            ground_truth = example['answer'].split('####')[-1].strip()
            
            is_correct = pred == ground_truth
            if is_correct:
                correct += 1
            
            results.append({
                'question': prompt[:60],
                'pred': pred,
                'gt': ground_truth,
                'correct': '✅' if is_correct else '❌'
            })
            
            progress_bar.progress((i + 1) / len(dataset))
            status_text.text(f"평가 진행 중... {i+1}/{len(dataset)} ({int(correct/(i+1)*100)}%)")
        
        progress_bar.empty()
        status_text.empty()
        
        accuracy = correct / len(dataset) * 100
        avg_time = total_time / len(dataset)
        
        return {
            'accuracy': accuracy,
            'avg_time': avg_time,
            'results': results
        }
    except Exception as e:
        st.error(f"평가 실패: {str(e)}")
        return None


def get_memory_usage():
    """현재 메모리 사용량 (MB)"""
    process = psutil.Process()
    return process.memory_info().rss / 1024**2


# 앱 시작
if "model" not in st.session_state:
    st.session_state.model, st.session_state.tokenizer = load_model()

model = st.session_state.model
tokenizer = st.session_state.tokenizer

if model:
    st.balloons()
else:
    st.stop()

# 메인 UI
st.markdown("""
# 🤖 EXAONE 1.2B LLM 대시보드
### LoRA 경량화 모델 - **LG AIMERS 8th Cohort 해커톤**
""")

# 사이드바 설정
with st.sidebar:
    st.markdown("## ⚙️ 설정")
    st.session_state.mode = st.radio("📌 모드", ["일반 채팅", "수학 풀이"], key="mode_selector")
    
    st.markdown("---")
    st.session_state.temperature = st.slider("🌡️ Temperature", 0.1, 2.0, 0.7, 0.1)
    st.session_state.top_p = st.slider("🎯 Top-P", 0.1, 1.0, 0.9, 0.05)
    st.session_state.max_length = st.slider("📏 길이", 50, 512, 256, 50)
    
    st.markdown("---")
    col1, col2 = st.columns(2)
    with col1:
        st.metric("모델", "EXAONE 1.2B")
    with col2:
        st.metric("최적화", "LoRA 43.7MB")
    
    st.markdown("---")
    eval_mode = st.checkbox("📈 GSM8K 벤치마크 (시간 소요)")
    if eval_mode:
        if st.button("🚀 평가 시작", key="eval_button"):
            benchmark_results = evaluate_gsm8k(n_samples=50)
            if benchmark_results:
                st.success(f"✅ 정확도: {benchmark_results['accuracy']:.1f}%")
                st.info(f"⏱️ 평균 시간: {benchmark_results['avg_time']:.2f}s")
                
                # 샘플 결과 표시
                st.markdown("#### 평가 샘플 (상위 5개)")
                for result in benchmark_results['results'][:5]:
                    st.write(f"{result['correct']} Q: {result['question']}...")
                    st.write(f"  예측: {result['pred']} | 정답: {result['gt']}")

# 탭 UI
tab1, tab2, tab3 = st.tabs(["💬 채팅", "📊 통계", "🏆 해커톤 분석"])

with tab1:
    st.markdown("### 챗봇과 대화하세요")
    
    if "messages" not in st.session_state:
        st.session_state.messages = []
    
    # 이전 메시지 표시
    for msg in st.session_state.messages:
        if msg["role"] == "user":
            st.markdown(f'<div class="message-user"><strong>👤 You:</strong> {msg["content"]}</div>', unsafe_allow_html=True)
        else:
            st.markdown(f'<div class="message-assistant"><strong>🤖 AI:</strong> {msg["content"]}</div>', unsafe_allow_html=True)
    
    st.markdown("---")
    
    # 입력 UI
    col1, col2 = st.columns([4, 1])
    with col1:
        user_input = st.text_input("메시지:", placeholder="질문을 입력하세요", key="user_input")
    with col2:
        send_button = st.button("전송", use_container_width=True)
    
    # 전송 처리
    if send_button and user_input:
        st.session_state.messages.append({"role": "user", "content": user_input})
        
        with st.spinner("🤖 응답 생성 중..."):
            response, processing_time = generate_response(
                model, tokenizer, user_input, 
                st.session_state.max_length, 
                st.session_state.temperature, 
                st.session_state.top_p
            )
            
            st.session_state.messages.append({"role": "assistant", "content": response})
        
        # 메트릭 표시
        col1, col2, col3 = st.columns(3)
        with col1:
            st.metric("길이", f"{len(response)} 자")
        with col2:
            st.metric("시간", f"{processing_time:.2f}초")
        with col3:
            if "수학" in st.session_state.mode:
                answer = extract_math_answer(response)
                st.metric("답", answer)
        
        st.rerun()
    
    if not st.session_state.messages:
        st.info("👋 EXAONE 1.2B에 오신 것을 환영합니다!\n수학 문제나 일반 질문을 해보세요.")

with tab2:
    st.markdown("### 📊 대화 통계")
    
    if st.session_state.messages:
        assistant_msgs = [m for m in st.session_state.messages if m["role"] == "assistant"]
        col1, col2, col3 = st.columns(3)
        with col1:
            st.metric("대화 수", len(st.session_state.messages) // 2)
        with col2:
            avg_length = sum(len(m["content"]) for m in assistant_msgs) / len(assistant_msgs) if assistant_msgs else 0
            st.metric("평균 길이", f"{int(avg_length)} 자")
        with col3:
            total_length = sum(len(m["content"]) for m in assistant_msgs)
            st.metric("총 생성", f"{total_length:,} 자")
    else:
        st.info("아직 대화가 없습니다.")

with tab3:
    st.markdown("### 🏆 LG AIMERS 해커톤 분석")
    
    st.markdown("#### 📊 경량화 효과")
    
    comparison_data = {
        "항목": ["전체 파라미터", "학습 파라미터", "추론 메모리", "학습 VRAM"],
        "Base Model": ["1.2B", "1.2B", "8-10 GB", "20-30 GB ❌"],
        "LoRA (당신)": ["1.2B", "11.4M (0.95%)", "9-11 GB", "T4 16GB ✅"]
    }
    
    st.table(comparison_data)
    
    st.markdown("#### 💾 메모리 효율성")
    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric("현재 메모리", f"{get_memory_usage():.1f} MB")
    with col2:
        st.metric("파라미터 절약", "99.05% ↓")
    with col3:
        st.metric("학습 시간", "~2시간 (T4)")
    
    st.markdown("#### 🎯 수상 포인트")
    st.success("""
    ✅ **99.1% 파라미터 절약** (1.2B → 11.4M 학습)
    ✅ **T4 GPU 호환** (16GB 내 학습/추론 성공)
    ✅ **실시간 Streamlit 데모** (수학 + 채팅 모드)
    ✅ **벤치마크 측정** (GSM8K 정확도 추적)
    ✅ **메모리 vs Accuracy 최적화** 증명
    """)
    
    st.markdown("#### 📈 개선 로드맵")
    st.info("""
    1️⃣ **QLoRA** - 4bit 양자화로 추가 50% 메모리 절약
    2️⃣ **Layer-wise LoRA** - 주요 층만 fine-tuning (5-6M 파라미터)
    3️⃣ **DPO + SFT** - 정확도 추가 3-5% 개선
    """)

# 푸터
st.markdown("---")
st.markdown("""
### 🎓 LG AIMERS 8th Cohort 해커톤 프로젝트
**EXAONE 4.0 1.2B + GSM8K LoRA 최적화**

- 모델: LGAI-EXAONE/EXAONE-4.0-1.2B
- 기법: LoRA (Low-Rank Adaptation)
- 데이터: GSM8K (수학 문제)
- 플랫폼: Streamlit + HuggingFace
- 개발자: LG AIMERS 8기 (군복무 중 학습)
""")