import streamlit as st
import pandas as pd
import os
from src.FisrtModule.module1 import MisconceptionModel
from src.SecondModule.module2 import SimilarQuestionGenerator
from src.ThirdModule.module3 import AnswerVerifier
import logging
from typing import Optional, Tuple
from pylatexenc.latex2text import LatexNodes2Text
import re

logging.basicConfig(level=logging.DEBUG)

# Initialize Misconception Model
@st.cache_resource
def load_misconception_model():
    return MisconceptionModel(
        model_name="minsuas/Misconceptions__1",
        misconception_mapping_path=os.path.join(data_path, 'misconception_mapping.parquet'),
        misconception_embs_paths=[os.path.join(data_path, f'embs_misconception-9-9.npy')]
    )

# Streamlit 페이지 기본 설정
st.set_page_config(
    page_title="MisconcepTutor", 
    layout="wide", 
    initial_sidebar_state="expanded"
)

@st.cache_resource
def load_answer_verifier():
    """답안 검증 모델 로드"""
    from src.ThirdModule.module3 import AnswerVerifier
    return AnswerVerifier()

# 경로 설정
base_path = os.path.dirname(os.path.abspath(__file__))
data_path = os.path.join(base_path, 'Data')
misconception_csv_path = os.path.join(data_path, 'misconception_mapping.csv')

# 로깅 설정
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 세션 상태 초기화 - 가장 먼저 실행되도록 최상단에 배치
if 'initialized' not in st.session_state:
    st.session_state.initialized = True
    st.session_state.wrong_questions = []
    st.session_state.misconceptions = []
    st.session_state.current_question_index = 0
    st.session_state.generated_questions = []
    st.session_state.current_step = 'initial'
    st.session_state.selected_wrong_answer = None
    st.session_state.questions = []
    logger.info("Session state initialized")


# 문제 생성기 초기화
@st.cache_resource
def load_question_generator():
    """문제 생성 모델 로드"""
    if not os.path.exists(misconception_csv_path):
        st.error(f"CSV 파일이 존재하지 않습니다: {misconception_csv_path}")
        raise FileNotFoundError(f"CSV 파일이 존재하지 않습니다: {misconception_csv_path}")
    return SimilarQuestionGenerator(misconception_csv_path=misconception_csv_path)

# CSV 데이터 로드 함수
@st.cache_data
def load_data(data_file='/train.csv', selected_indexes=None):
#def load_data(data_file='/processed_mathqa2.csv', selected_indexes=None):
    try:
        file_path = os.path.join(data_path, data_file.lstrip('/'))
        df = pd.read_csv(file_path)
        logger.info(f"Data loaded successfully from {file_path}")

        if selected_indexes is not None:
            #df = df.loc[selected_indexes]  # 신규 문제
            df = df.loc[df['QuestionId'].isin(selected_indexes)]  # QuestionId 기준 필터링
            logger.info(f"Data filtered to selected indexes: {selected_indexes}")
            
            # 필터링 후 다시 한번 중복 체크
            if df.duplicated(subset=['QuestionText']).any():
                df = df.drop_duplicates(subset=['QuestionText'], keep='first')
                logger.warning("Removed duplicates from selected indexes")
        
        return df
        
    except FileNotFoundError:
        st.error(f"파일을 찾을 수 없습니다: {data_file}")
        logger.error(f"File not found: {data_file}")
        return None

def start_quiz():
    """퀴즈 시작 및 초기화"""

    #selected_indexes = [2519, 3852, 3404, 3896, 7602, 3946, 12977, 1878, 7602, 3589, 9] 
    # 12038 문제 끊김? # 1302 동일 문제인가? # 3473 안나옴? # 3887 답안 틀림 # 9699 수식 이상  # 9752 문제 끊김 

    # train.csv
    selected_indexes = [1866, 1864, 1845, 1862, 1861, 1829, 1827, 1802, 1741, 1725] # 1671]
    # 확정 : 1671, 1725
    # 통과 : 1864, 1866, 1845, 1861, 1862, 1802, 1827, 1829, 1741, 
    # 세모 : 1825, 
    # 문제 안이쁨 : 1868, 1847, 1834, 1841, 1809, 1804, 1672, 1731, 1736, 1746, 1692, 1775, 1781
    # 문제 이상 : 1792, 1804, 1813, 1679, 1711
    
    df = load_data(selected_indexes=selected_indexes)
    if df is None or df.empty:
        st.error("데이터를 불러올 수 없습니다. 데이터셋을 확인해주세요.")
        return

    #st.session_state.questions = df.sample(n=10, random_state=42)
    st.session_state.questions = df
    st.session_state.current_step = 'quiz'
    st.session_state.current_question_index = 0
    st.session_state.wrong_questions = []
    st.session_state.misconceptions = []
    st.session_state.generated_questions = []
    logger.info("Quiz started")

def generate_similar_question(wrong_q, misconception_id, generator):
    """유사 문제 생성"""
    logger.info(f"Generating similar question for misconception_id: {misconception_id}")
    
    # 입력 데이터 유효성 검사
    if not isinstance(wrong_q, dict):
        logger.error(f"Invalid wrong_q type: {type(wrong_q)}")
        st.error("유사 문제 생성에 필요한 데이터 형식이 잘못되었습니다.")
        return None
        
    try:
        # misconception_id가 없거나 NaN인 경우 다른 misconception 사용
        if pd.isna(misconception_id):
            logger.info("Original misconception_id is NaN, trying to find alternative")
            # 현재까지 나온 misconception들 중에서 선택
            available_misconceptions = [m for m in st.session_state.misconceptions if not pd.isna(m)]
            
            if available_misconceptions:
                # 가장 최근에 나온 misconception 선택
                misconception_id = available_misconceptions[-1]
                logger.info(f"Using alternative misconception_id: {misconception_id}")
            else:
                # 기본 misconception ID 사용 (예: 가장 기본적인 misconception)
                misconception_id = 2001  # 적절한 기본값으로 수정 필요
                logger.info(f"Using default misconception_id: {misconception_id}")

        # 데이터 준비 (튜플 변환 방지)
        input_data = {
            'construct_name': str(wrong_q.get('ConstructName', '')),
            'subject_name': str(wrong_q.get('SubjectName', '')),
            'question_text': str(wrong_q.get('QuestionText', '')),
            'correct_answer_text': str(wrong_q.get(f'Answer{wrong_q["CorrectAnswer"]}Text', '')),
            'wrong_answer_text': str(wrong_q.get(f'Answer{st.session_state.selected_wrong_answer}Text', '')),
            'misconception_id': int(misconception_id)
        }
        
        logger.info(f"Prepared input data: {input_data}")

        with st.spinner("📝 유사 문제를 생성하고 있습니다..."):
            # 유사 문제 생성 호출
            generated_q, _ = generator.generate_similar_question_with_text(
                construct_name=input_data['construct_name'],
                subject_name=input_data['subject_name'],
                question_text=input_data['question_text'],
                correct_answer_text=input_data['correct_answer_text'],
                wrong_answer_text=input_data['wrong_answer_text'],
                misconception_id=input_data['misconception_id']
            )

        if generated_q:
            verifier = load_answer_verifier()
            with st.status("🤔 AI가 문제를 검토하고 있습니다..."):
                st.write("답안의 정확성을 검증하고 있습니다...")
                verified_answer = verifier.verify_answer(
                    question=generated_q.question,
                    choices=generated_q.choices
                )
                
                if verified_answer:
                    logger.info(f"Answer verified: {verified_answer}")
                    st.write("✅ 검증 완료!")
                    result = {
                        'question': generated_q.question,
                        'choices': generated_q.choices,
                        'correct': verified_answer,
                        'explanation': generated_q.explanation
                    }
                    st.session_state['current_similar_question_answer'] = verified_answer
                    return result
                else:
                    logger.warning("Answer verification failed, using original answer")
                    st.write("⚠️ 검증에 실패했습니다. 원본 답안을 사용합니다.")
                    result = {
                        'question': generated_q.question,
                        'choices': generated_q.choices,
                        'correct': generated_q.correct_answer,
                        'explanation': generated_q.explanation
                    }
                    st.session_state['current_similar_question_answer'] = generated_q.correct_answer
                    return result
                    
    except Exception as e:
        logger.error(f"Error in generate_similar_question: {str(e)}")
        st.error(f"문제 생성 중 오류가 발생했습니다: {str(e)}")
        return None

    return None

def handle_answer(answer, current_q):
    """답변 처리"""
    if answer != current_q['CorrectAnswer']:
        wrong_q_dict = current_q.to_dict()
        st.session_state.wrong_questions.append(wrong_q_dict)
        st.session_state.selected_wrong_answer = answer
        
        misconception_key = f'Misconception{answer}Id'
        misconception_id = current_q.get(misconception_key)
        st.session_state.misconceptions.append(misconception_id)
    
    st.session_state.current_question_index += 1
    if st.session_state.current_question_index >= len(st.session_state.questions):
        st.session_state.current_step = 'review'
    else:
        st.session_state.current_step = 'quiz'

def display_math_content(content):
    """
    Display mathematical content with proper formatting.
    
    Args:
        content (str): The math content to display
    """
    # Convert LaTeX to plain text for display
    from pylatexenc.latex2text import LatexNodes2Text
    
    # Clean and format the content
    formatted_content = LatexNodes2Text().latex_to_text(content)
    st.markdown(f'<div class="math-container">{formatted_content}</div>', unsafe_allow_html=True)

def add_custom_css():
    st.markdown(
        """
        <style>
            .problem-header {
                color: #FF6B6B;
                font-size: 24px;
                font-weight: bold;
                margin-bottom: 20px;
            }
            .math-container {
                background-color: #f0f8ff;
                padding: 15px 20px;
                border-radius: 5px;
                margin: 5px 0;
            }
            button {
                color: #0066ff;
                font-weight: 500;
            }
        </style>
        """,
        unsafe_allow_html=True
    )

def display_question(question, answers):
    """Display question and options with LaTeX formatting"""
    st.markdown('<div class="problem-header">Problem:</div>', unsafe_allow_html=True)
    display_math_content(question)
    
    # Add custom CSS for options
    st.markdown("""
        <style>
        .option-container {
            background-color: #f0f8ff;
            padding: 10px 20px;
            margin: 5px 0;
            border-radius: 5px;
            cursor: pointer;
            display: flex;
            align-items: center;
            gap: 20px;
        }
        .option-text {
            color: #0066ff;
            font-weight: 500;
            width: 30px;
        }
        </style>
    """, unsafe_allow_html=True)
    
    # Display options
    for opt in ['A', 'B', 'C', 'D']:
        with st.container():
            col1, col2 = st.columns([1, 11])
            with col1:
                if st.button(f"{opt}.", key=f"btn_{opt}", help="Click to select"):
                    handle_answer(opt, st.session_state.questions.iloc[st.session_state.current_question_index])
                    st.rerun()
            with col2:
                display_option_content(answers[opt])

def display_option_content(option_text):
    """Process and display option content with LaTeX formatting"""
    from pylatexenc.latex2text import LatexNodes2Text
    formatted_content = LatexNodes2Text().latex_to_text(option_text)
    st.markdown(f'<div class="math-container">{formatted_content}</div>', unsafe_allow_html=True)

def update_similar_question_display(new_question, i, answered=False):
    """Display similar question and its options"""
    display_math_content(new_question['question'])
    
    # Display options
    for opt in ['A', 'B', 'C', 'D']:
        with st.container():
            col1, col2 = st.columns([1, 11])
            with col1:
                if st.button(f"{opt}.", key=f"sim_btn_{opt}_{i}", help="Click to select"):
                    if not answered:
                        # 선택한 옵션(opt)을 st.session_state에 저장
                        st.session_state[f"similar_question_answered_{i}"] = True
                        st.session_state[f"selected_answer_{i}"] = opt
                        correct_answer = st.session_state.get('current_similar_question_answer')
                        # 정답 여부를 확인
                        st.session_state[f"is_correct_{i}"] = (opt == correct_answer)
                        st.rerun()
            with col2:
                display_option_content(new_question['choices'][opt])


def main():
    """메인 애플리케이션 로직"""
    st.title("MisconcepTutor")

    # Misconception Model 로드
    misconception_model = load_misconception_model()

    # Generator 초기화
    generator = load_question_generator()
    add_custom_css()
    
    # 초기 화면
    if st.session_state.current_step == 'initial':
        st.write("#### 학습을 시작하겠습니다. 10개의 문제를 풀어볼까요?")
        if st.button("학습 시작", key="start_quiz"):
            start_quiz()
            st.rerun()

    # 퀴즈 화면
    elif st.session_state.current_step == 'quiz':
        current_q = st.session_state.questions.iloc[st.session_state.current_question_index]
        
        # 진행 상황 표시
        progress = st.session_state.current_question_index / 10
        st.progress(progress)
        st.write(f"### 문제 {st.session_state.current_question_index + 1}/10")
        
        # 문제 표시
        st.markdown("---")
        question_row = current_q['QuestionText']
        question_text = LatexNodes2Text().latex_to_text(current_q['QuestionText'])
        
        answers ={
            'A': current_q['AnswerAText'],
            'B': current_q['AnswerBText'],
            'C': current_q['AnswerCText'],
            'D': current_q['AnswerDText']
            }
        
        display_question(question_text, answers)
    

    # 복습 화면
    elif st.session_state.current_step == 'review':
        st.write("### 학습 결과")
        
        # 결과 통계
        col1, col2, col3 = st.columns(3)
        col1.metric("총 문제 수", 10)
        col2.metric("맞은 문제", 10 - len(st.session_state.wrong_questions))
        col3.metric("틀린 문제", len(st.session_state.wrong_questions))
        
        # 결과에 따른 메시지 표시
        if len(st.session_state.wrong_questions) == 0:
            st.balloons()  # 축하 효과
            st.success("🎉 축하합니다! 모든 문제를 맞추셨어요!")
            st.markdown("""
            ### 🏆 수학왕이십니다! 
            완벽한 점수를 받으셨네요! 수학적 개념을 정확하게 이해하고 계신 것 같습니다.
            """)
        elif len(st.session_state.wrong_questions) <= 3:
            st.success("잘 하셨어요! 조금만 더 연습하면 완벽할 거예요!")
        else:
            st.info("천천히 개념을 복습해보아요. 연습하다 보면 늘어날 거예요!")
        
        # 네비게이션 버튼
        col1, col2 = st.columns(2)
        with col1:
            if st.button("🔄 새로운 문제 세트 시작하기", use_container_width=True):
                start_quiz()
                st.rerun()
        with col2:
            if st.button("🏠 처음으로 돌아가기", use_container_width=True):
                st.session_state.clear()
                st.rerun()
                
        # 틀린 문제 분석 부분
        if st.session_state.wrong_questions:
            st.write("### ✍️ 틀린 문제 분석")
            tabs = st.tabs([f"📝 틀린 문제 #{i + 1}" for i in range(len(st.session_state.wrong_questions))])
            
            for i, (tab, (wrong_q, misconception_id)) in enumerate(zip(
                tabs,
                zip(st.session_state.wrong_questions, st.session_state.misconceptions)
            )):
                with tab:
                    st.write("**📋 문제:**")
                    display_math_content(wrong_q['QuestionText'])  # 문제 렌더링
                    st.write("**✅ 정답:**")
                    display_option_content(wrong_q[f'Answer{wrong_q["CorrectAnswer"]}Text'])

                    st.write("---")
                    st.write("**🔍 관련된 Misconception:**")
                    if misconception_id and not pd.isna(misconception_id):
                        misconception_text = generator.get_misconception_text(misconception_id)
                        st.info(f"Misconception ID: {int(misconception_id)}\n\n{misconception_text}")
                    else:
                        st.info("Misconception 정보가 없습니다.")
        
                    if st.button(f"📚 유사 문제 풀기", key=f"retry_{i}"):
                        st.session_state[f"show_similar_question_{i}"] = True
                        st.session_state[f"similar_question_answered_{i}"] = False
                        st.rerun()
                    if st.session_state.get(f"show_similar_question_{i}", False):
                        st.divider()
                        new_question = generate_similar_question(wrong_q, misconception_id, generator)
                        if new_question:
                            st.write("### 🎯 유사 문제")
                            #display_math_content(new_question['question'])  # 함수 교체

                            # 답변 상태 확인
                            answered = st.session_state.get(f"similar_question_answered_{i}", False)
                            #update_similar_question_display(new_question, i, answered)
                            # 선택한 옵션을 처리하는 함수를 호출
                            update_similar_question_display(new_question, i)

                            # 답변한 경우 결과 표시
                            if answered:
                                is_correct = st.session_state.get(f"is_correct_{i}", False)
                                correct_answer = st.session_state.get('current_similar_question_answer')
                                if is_correct:
                                    st.success("✅ 정답입니다!")
                                else:
                                    st.error(f"❌ 틀렸습니다. 정답은 {correct_answer}입니다.")
                                
                                # 해설 표시
                                st.write("---")
                                st.write("**📝 해설:**", new_question['explanation'])
                                
                                # 다시 풀기 버튼
                                if st.button("🔄 다시 풀기", key=f"reset_{i}"):
                                    st.session_state[f"similar_question_answered_{i}"] = False
                                    st.session_state[f"selected_answer_{i}"] = None
                                    st.session_state[f"is_correct_{i}"] = None
                                    st.rerun()
                            
                            # 문제 닫기 버튼
                            if st.button("❌ 문제 닫기", key=f"close_{i}"):
                                st.session_state[f"show_similar_question_{i}"] = False
                                st.session_state[f"similar_question_answered_{i}"] = False
                                st.session_state[f"selected_answer_{i}"] = None
                                st.session_state[f"is_correct_{i}"] = None
                                st.rerun()

                            # 화면 아래 여백 추가
                            st.markdown("<br>" * 5, unsafe_allow_html=True)  # 5줄의 빈 줄 추가
                            st.markdown("""
                                <div style="height: 100px;">
                                </div>
                                """, unsafe_allow_html=True)  # 추가 여백
                        else:
                            st.error("유사 문제를 생성할 수 없습니다.")
                            if st.button("❌ 닫기", key=f"close_error_{i}"):
                                st.session_state[f"show_similar_question_{i}"] = False
                                st.rerun()
                                # 화면 아래 여백 추가
                            st.markdown("<br>" * 5, unsafe_allow_html=True)  # 5줄의 빈 줄 추가
                            st.markdown("""
                                <div style="height: 100px;">
                                </div>
                                """, unsafe_allow_html=True)  # 추가 여백

    
    #     # 틀린 문제 분석
    #     if st.session_state.wrong_questions:
    #         st.write("### ✍️ 틀린 문제 분석")
    #         tabs = st.tabs([f"📝 틀린 문제 #{i + 1}" for i in range(len(st.session_state.wrong_questions))])

    #         for i, (tab, (wrong_q, misconception_id)) in enumerate(zip(
    #             tabs,
    #             zip(st.session_state.wrong_questions, st.session_state.misconceptions)
    #         )):
    #             with tab:
    #                 st.write("**📋 문제:**")
    #                 st.write(wrong_q['QuestionText'])
    #                 st.write("**✅ 정답:**", wrong_q['CorrectAnswer'])

    #                 st.write("---")
    #                 st.write("**🔍 관련된 Misconception:**")
    #                 if misconception_id and not pd.isna(misconception_id):
    #                     misconception_text = misconception_model.misconception_names.get(misconception_id, "정보 없음")
    #                     st.info(f"Misconception ID: {int(misconception_id)}\n\n{misconception_text}")
    #                 else:
    #                     st.info("Misconception 정보가 없습니다.")

if __name__ == "__main__":
    main()


# random_state 42에서 정답
    # D C A A C
    # A B B B B