import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import json
from typing import List, Dict, Tuple
import time

# 주의: 이것은 SEAL 논문의 개념을 보여주는 단순화된 데모입니다.
# 실제 구현에서는 LoRA 파인튜닝과 전체 RL 루프가 필요합니다.

class SEALDemo:
    def __init__(self):
        # 실제로는 Qwen2.5-7B를 사용하지만, 데모에서는 작은 모델 사용
        self.model_name = "gpt2"
        print("Loading model...")
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        self.model = AutoModelForCausalLM.from_pretrained(self.model_name)
        self.tokenizer.pad_token = self.tokenizer.eos_token
        
        # 시뮬레이션을 위한 가상의 self-edit 효과
        self.adaptation_history = []
        
    def generate_self_edit(self, passage: str, edit_type: str = "implications") -> str:
        """주어진 passage에 대한 self-edit 생성"""
        
        prompts = {
            "implications": f"Let's read the following passage and produce a list of implications derived directly or indirectly from the content.\n\nPassage:\n{passage}\n\nImplications:\n",
            "rewrite": f"Let's read the following passage and rewrite it in a few different ways, each one separated by a newline.\n\nPassage:\n{passage}\n\nRewritten passages:\n",
            "qa": f"Let's read the following passage and rewrite it in a question-answer format.\n\nPassage:\n{passage}\n\nQuestion 1:\n"
        }
        
        prompt = prompts.get(edit_type, prompts["implications"])
        
        # 실제 생성 (데모용으로 간단하게)
        inputs = self.tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)
        
        with torch.no_grad():
            outputs = self.model.generate(
                inputs.input_ids,
                max_new_tokens=200,
                temperature=0.8,
                do_sample=True,
                top_p=0.9,
                pad_token_id=self.tokenizer.eos_token_id
            )
        
        generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        self_edit = generated_text[len(prompt):].strip()
        
        # 데모를 위한 예시 self-edit 추가 (실제 모델이 생성하지 못할 경우를 대비)
        if len(self_edit) < 50:
            if edit_type == "implications":
                self_edit = """1. The information in this passage can be used to answer specific questions about the content.
2. Understanding this passage requires identifying key facts and their relationships.
3. The main concepts should be broken down into smaller, learnable components.
4. Each fact in the passage has potential implications for related topics.
5. This content can be restructured to improve learning and retention."""
            elif edit_type == "rewrite":
                self_edit = f"""Version 1: {passage[:100]}...
Version 2: The key information from this passage includes...
Version 3: In summary, this passage discusses..."""
            else:  # qa
                self_edit = """Question 1: What is the main topic of this passage?
Answer: The passage discusses...

Question 2: What are the key facts mentioned?
Answer: The key facts include..."""
        
        return self_edit
    
    def simulate_adaptation(self, passage: str, self_edit: str) -> Dict[str, float]:
        """self-edit를 사용한 적응 과정 시뮬레이션"""
        
        # 실제로는 LoRA 파인튜닝을 수행하지만, 여기서는 시뮬레이션
        adaptation_score = len(self_edit.split('\n')) * 0.1  # 더 많은 implications = 더 나은 적응
        
        # 적응 기록 저장
        self.adaptation_history.append({
            "passage": passage[:100] + "...",
            "self_edit_lines": len(self_edit.split('\n')),
            "adaptation_score": adaptation_score
        })
        
        return {
            "original_performance": 0.33,  # 논문의 baseline
            "adapted_performance": min(0.33 + adaptation_score, 0.47),  # 논문의 SEAL 성능
            "improvement": min(adaptation_score, 0.14)
        }
    
    def answer_question(self, question: str, passage: str, use_adaptation: bool) -> Tuple[str, float]:
        """질문에 대한 답변 생성 (적응 전/후 비교)"""
        
        if use_adaptation:
            # 적응된 모델로 답변 (시뮬레이션)
            prompt = f"Based on the learned information, answer this question concisely:\n\nQuestion: {question}\nAnswer:"
            confidence = 0.8
        else:
            # 원본 모델로 답변
            prompt = f"Answer this question:\n\nQuestion: {question}\nAnswer:"
            confidence = 0.3
        
        inputs = self.tokenizer(prompt, return_tensors="pt", max_length=256, truncation=True)
        
        with torch.no_grad():
            outputs = self.model.generate(
                inputs.input_ids,
                max_new_tokens=50,
                temperature=0.7,
                do_sample=True,
                pad_token_id=self.tokenizer.eos_token_id
            )
        
        answer = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        answer = answer[len(prompt):].strip()
        
        # 답변이 너무 짧으면 기본 답변 제공
        if len(answer) < 10:
            answer = "The model cannot provide a confident answer without adaptation." if not use_adaptation else "Based on the adapted knowledge, the answer is [specific to the passage content]."
        
        return answer, confidence

# 전역 인스턴스
demo = SEALDemo()

def process_seal_adaptation(passage, edit_type, question):
    """SEAL 적응 프로세스 실행"""
    
    if not passage or not question:
        return "Please provide both a passage and a question.", "", {}, "", ""
    
    # 1. Self-edit 생성
    self_edit = demo.generate_self_edit(passage, edit_type)
    
    # 2. 적응 시뮬레이션
    adaptation_results = demo.simulate_adaptation(passage, self_edit)
    
    # 3. 질문에 대한 답변 (적응 전/후)
    answer_before, conf_before = demo.answer_question(question, passage, use_adaptation=False)
    answer_after, conf_after = demo.answer_question(question, passage, use_adaptation=True)
    
    # 결과 포맷팅
    results_text = f"""📊 Adaptation Results:
    
Original Model Performance: {adaptation_results['original_performance']:.1%}
Adapted Model Performance: {adaptation_results['adapted_performance']:.1%}
Improvement: +{adaptation_results['improvement']:.1%}

📈 Answer Confidence:
Before Adaptation: {conf_before:.1%}
After Adaptation: {conf_after:.1%}"""
    
    return self_edit, results_text, adaptation_results, answer_before, answer_after

def visualize_rl_process():
    """RL 훈련 과정 시각화 (논문 Figure 4 재현)"""
    
    # 논문의 RL iteration 결과 시뮬레이션
    iterations = [0, 1, 2]
    performance = [33.5, 43.7, 47.0]  # 논문의 실제 수치
    
    import matplotlib.pyplot as plt
    
    fig, ax = plt.subplots(figsize=(8, 6))
    
    # 다양한 방법들의 성능
    methods = {
        'Base Model': 32.7,
        'Raw Passage': 33.5,
        'Base Model Synthetic': 39.7,
        'GPT-4.1 Synthetic': 46.3,
    }
    
    # 수평선으로 baseline 표시
    for method, score in methods.items():
        ax.axhline(y=score, linestyle='--', alpha=0.5, label=method)
    
    # SEAL 성능 곡선
    ax.plot(iterations, performance, 'ro-', linewidth=2, markersize=8, label='SEAL')
    
    ax.set_xlabel('ReSTEM RL Training Iterations')
    ax.set_ylabel('Knowledge Incorporation (%)')
    ax.set_title('Single-Passage Knowledge Incorporation Performance')
    ax.legend()
    ax.grid(True, alpha=0.3)
    ax.set_ylim(30, 50)
    
    return fig

# Gradio 인터페이스
with gr.Blocks(title="SEAL: Self-Adapting Language Models Demo") as interface:
    gr.Markdown("""
    # 🔄 SEAL: Self-Adapting Language Models Demo
    
    This demo illustrates the key concepts from the paper "Self-Adapting Language Models" (Zweiger et al., 2025).
    SEAL enables LLMs to self-adapt by generating their own finetuning data and update directives.
    
    ## 📚 How it works:
    1. **Input a passage** containing new information
    2. **Model generates self-edits** (implications, rewrites, or Q&A pairs)
    3. **Simulated adaptation** using the self-edits
    4. **Compare performance** before and after adaptation
    """)
    
    with gr.Tab("Knowledge Incorporation"):
        with gr.Row():
            with gr.Column():
                passage_input = gr.Textbox(
                    label="📄 Input Passage",
                    placeholder="Enter a passage with factual information...",
                    lines=8,
                    value="The Apollo program faced significant challenges before achieving its goal. Kennedy's science advisor, Jerome Wiesner, initially opposed manned spacecraft flights. Despite internal disagreements at NASA, the program eventually succeeded in landing humans on the Moon in 1969."
                )
                
                edit_type = gr.Radio(
                    choices=["implications", "rewrite", "qa"],
                    value="implications",
                    label="Self-Edit Type"
                )
                
                question_input = gr.Textbox(
                    label="❓ Test Question",
                    placeholder="Enter a question about the passage...",
                    value="Who was Kennedy's science advisor?"
                )
                
                adapt_btn = gr.Button("🚀 Run SEAL Adaptation", variant="primary")
            
            with gr.Column():
                self_edit_output = gr.Textbox(
                    label="🤖 Generated Self-Edit",
                    lines=8
                )
                
                results_output = gr.Textbox(
                    label="📊 Adaptation Results",
                    lines=6
                )
        
        with gr.Row():
            answer_before = gr.Textbox(
                label="❌ Answer Before Adaptation",
                lines=2
            )
            answer_after = gr.Textbox(
                label="✅ Answer After Adaptation",
                lines=2
            )
        
        adapt_btn.click(
            process_seal_adaptation,
            inputs=[passage_input, edit_type, question_input],
            outputs=[self_edit_output, results_output, gr.State(), answer_before, answer_after]
        )
    
    with gr.Tab("RL Training Visualization"):
        gr.Markdown("""
        ## 📈 Reinforcement Learning Training Process
        
        This visualization shows how SEAL improves over RL iterations, eventually surpassing even GPT-4.1 synthetic data.
        """)
        
        plot_btn = gr.Button("Generate Performance Plot")
        plot_output = gr.Plot()
        
        plot_btn.click(visualize_rl_process, outputs=plot_output)
    
    with gr.Tab("Paper Overview"):
        gr.Markdown("""
        ## 📑 Key Contributions
        
        1. **Self-Edits**: Models generate their own training data in natural language
        2. **RL Training**: Use reinforcement learning to optimize self-edit generation
        3. **Two Applications**:
           - Knowledge Incorporation: 47.0% accuracy (vs 33.5% baseline)
           - Few-shot Learning: 72.5% success rate on ARC tasks
        
        ## 🔧 Technical Details
        
        - **Base Models**: Qwen2.5-7B (knowledge), Llama-3.2-1B (few-shot)
        - **Optimization**: ReSTEM (rejection sampling + SFT)
        - **Adaptation**: LoRA for efficient weight updates
        - **Reward Signal**: Downstream task performance
        
        ## ⚠️ Demo Limitations
        
        This demo simulates the core concepts but doesn't perform actual LoRA finetuning or full RL training due to computational constraints.
        """)

if __name__ == "__main__":
    interface.launch()