#!/usr/bin/env python3
"""
Eigenverse Theorem Prover
AI-assisted research tool for machine-verified theorems.
"""

import json
import os
import numpy as np
import gradio as gr
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')
from huggingface_hub import InferenceClient

# HuggingFace Inference API setup
HF_TOKEN = os.environ.get("HF_TOKEN", None)  # Set as Space secret
client = InferenceClient(
    model="mistralai/Mistral-7B-Instruct-v0.3",
    token=HF_TOKEN
)

# Load knowledge base
with open('theorems_kb.json', 'r') as f:
    KB = json.load(f)

# Prepare TF-IDF search
corpus = [f"{item['name']} {item['statement']} {item['docstring']}" for item in KB]
vectorizer = TfidfVectorizer(stop_words='english', max_features=500)
tfidf_matrix = vectorizer.fit_transform(corpus)

# Coherence scales
COHERENCE_SCALES = {
    'kernel': 1.0,
    'silver': 0.707,
    'golden': 0.667
}

def coherence_function(r):
    """C(r) = 2r/(1+r²)"""
    return (2 * r) / (1 + r**2)

def find_relevant_theorems(hypothesis, top_k=5):
    """Find top-k most relevant theorems using TF-IDF."""
    if not hypothesis.strip():
        return []
    
    query_vec = vectorizer.transform([hypothesis])
    similarities = cosine_similarity(query_vec, tfidf_matrix).flatten()
    top_indices = np.argsort(similarities)[::-1][:top_k]
    
    results = []
    for idx in top_indices:
        if similarities[idx] > 0:
            results.append({
                'theorem': KB[idx],
                'score': float(similarities[idx])
            })
    
    return results

def generate_template_scaffold(hypothesis, relevant_theorems):
    """Generate basic Lean 4 proof scaffold template (fallback)."""
    
    scaffold = f"""-- Hypothesis: {hypothesis}
-- Upstream dependencies:
"""
    
    for i, item in enumerate(relevant_theorems, 1):
        thm = item['theorem']
        scaffold += f"-- {i}. {thm['name']} from {thm['module']}\n"
    
    scaffold += f"""
theorem my_hypothesis : <replace_with_formal_statement> := by
"""
    
    for i, item in enumerate(relevant_theorems, 1):
        thm = item['theorem']
        scaffold += f"  have h{i} := {thm['name']}\n"
    
    scaffold += """  -- Continue proof using h1, h2, ...
  sorry -- replace with actual proof
"""
    
    return scaffold

def generate_llm_hypothesis(user_hypothesis: str, relevant_theorems: list) -> str:
    """Generate LLM-powered hypothesis analysis and Lean 4 draft."""
    
    theorem_context = "\n".join([
        f"- {item['theorem']['name']} ({item['theorem']['module']}): {item['theorem']['statement'][:200]}"
        for item in relevant_theorems[:5]
    ])
    
    prompt = f"""You are a Lean 4 theorem proving assistant for the Eigenverse mathematical framework.

The Eigenverse framework is built around the critical eigenvalue μ = exp(i·3π/4) at the 135° phase, with coherence function C(r) = 2r/(1+r²). It has 688 machine-verified theorems across physics, chemistry, and cosmology.

User hypothesis: {user_hypothesis}

Relevant existing theorems found:
{theorem_context}

Your task:
1. Analyze how the hypothesis relates to these theorems
2. Identify what mathematical structures are needed
3. Draft a Lean 4 theorem statement and proof skeleton

Output format:
## Analysis
<2-3 sentences connecting hypothesis to existing theorems>

## Lean 4 Draft
```lean4
-- Hypothesis: {user_hypothesis}
-- Dependencies: list the relevant theorems

theorem hypothesis_name : <formal_statement> := by
  -- Step 1: <what to do>
  have h1 := relevant_theorem_1
  -- Step 2: <what to do>  
  sorry -- replace with proof
```

## Next Steps
<what needs to be proven to complete this>"""

    try:
        response = client.text_generation(
            prompt,
            max_new_tokens=600,
            temperature=0.3,
            do_sample=True,
        )
        return response
    except Exception as e:
        # Fallback to template if LLM fails
        fallback = f"⚠️ LLM generation failed ({str(e)}). Showing template scaffold.\n\n"
        fallback += generate_template_scaffold(user_hypothesis, relevant_theorems)
        return fallback

def hypothesis_prover_tab(hypothesis):
    """Tab 1: Hypothesis Prover - now with LLM-powered analysis."""
    if not hypothesis.strip():
        return "Please enter a hypothesis.", ""
    
    relevant = find_relevant_theorems(hypothesis, top_k=5)
    
    if not relevant:
        return "No relevant theorems found. Try different keywords.", ""
    
    # Format relevant theorems
    theorems_output = f"## Relevant Theorems (Top {len(relevant)})\n\n"
    
    for i, item in enumerate(relevant, 1):
        thm = item['theorem']
        score = item['score']
        theorems_output += f"### {i}. {thm['name']} (relevance: {score:.3f})\n"
        theorems_output += f"**Module:** {thm['module']}\n\n"
        theorems_output += f"**Kind:** {thm['kind']}\n\n"
        if thm['docstring']:
            theorems_output += f"**Description:** {thm['docstring']}\n\n"
        theorems_output += f"```lean\n{thm['statement']}\n```\n\n"
        theorems_output += "---\n\n"
    
    # Generate LLM-powered hypothesis analysis
    llm_output = generate_llm_hypothesis(hypothesis, relevant)
    
    return theorems_output, llm_output

def search_theorems(module_filter, search_query):
    """Tab 2: Theorem Explorer."""
    filtered = KB
    
    # Filter by module
    if module_filter and module_filter != "All Modules":
        filtered = [item for item in filtered if item['module'] == module_filter]
    
    # Search by query
    if search_query.strip():
        query_lower = search_query.lower()
        filtered = [
            item for item in filtered 
            if query_lower in item['name'].lower() 
            or query_lower in item['statement'].lower()
            or query_lower in item['docstring'].lower()
        ]
    
    if not filtered:
        return "No theorems found."
    
    # Format results
    output = f"## Found {len(filtered)} theorem(s)\n\n"
    
    for item in filtered[:50]:  # Limit to 50 for display
        output += f"**{item['name']}** | {item['module']} | {item['kind']}\n"
        if item['docstring']:
            output += f"  _{item['docstring']}_\n"
        output += f"```lean\n{item['statement'][:200]}{'...' if len(item['statement']) > 200 else ''}\n```\n\n"
    
    if len(filtered) > 50:
        output += f"\n_Showing first 50 of {len(filtered)} results._"
    
    return output

def coherence_analyzer_tab(r_value):
    """Tab 3: Coherence Analyzer."""
    try:
        r = float(r_value)
    except:
        return "Invalid input. Please enter a number.", None
    
    c_r = coherence_function(r)
    
    # Find closest scale
    closest_scale = min(COHERENCE_SCALES.items(), key=lambda x: abs(x[1] - c_r))
    
    output = f"## Coherence Analysis\n\n"
    output += f"**Input:** r = {r:.4f}\n\n"
    output += f"**C(r) = 2r/(1+r²) = {c_r:.6f}**\n\n"
    output += f"**Closest coherence scale:** {closest_scale[0]} (C = {closest_scale[1]:.3f})\n\n"
    output += f"**Distance from {closest_scale[0]} scale:** {abs(c_r - closest_scale[1]):.6f}\n\n"
    
    # Find relevant theorems for this scale
    scale_keywords = {
        'kernel': ['kernel', 'critical', 'maximum', 'coherence'],
        'silver': ['silver', 'eta', 'sqrt'],
        'golden': ['golden', 'phi', 'koide']
    }
    
    keywords = scale_keywords.get(closest_scale[0], [])
    relevant_theorems = []
    
    for item in KB:
        if any(kw in item['name'].lower() or kw in item['statement'].lower() for kw in keywords):
            relevant_theorems.append(item)
    
    output += f"### Relevant Theorems ({len(relevant_theorems[:5])} shown)\n\n"
    
    for item in relevant_theorems[:5]:
        output += f"- **{item['name']}** ({item['module']})\n"
    
    # Create bar chart
    fig, ax = plt.subplots(figsize=(8, 5))
    
    scales = list(COHERENCE_SCALES.keys()) + [f'C({r:.2f})']
    values = list(COHERENCE_SCALES.values()) + [c_r]
    colors = ['#4A90E2', '#50E3C2', '#F5A623', '#E74C3C']
    
    bars = ax.bar(scales, values, color=colors)
    ax.set_ylabel('Coherence C(r)', fontsize=12)
    ax.set_title('Coherence Scale Comparison', fontsize=14, fontweight='bold')
    ax.set_ylim(0, 1.1)
    ax.grid(axis='y', alpha=0.3)
    
    # Add value labels on bars
    for bar, val in zip(bars, values):
        height = bar.get_height()
        ax.text(bar.get_x() + bar.get_width()/2., height + 0.02,
                f'{val:.3f}', ha='center', va='bottom', fontsize=10)
    
    plt.tight_layout()
    
    return output, fig

def about_tab():
    """Tab 4: About."""
    return f"""# Eigenverse Theorem Prover

## Overview

AI-assisted research tool for the **Eigenverse** framework — a mathematical foundation proving reality emerges uniquely from three pre-physical axioms.

**688 machine-verified theorems. Zero sorry.**

## Architecture

### Funneling vs Tunneling

- **Funneling:** Converge many possibilities to one solution (μ-state uniqueness)
- **Tunneling:** Traverse barriers via coherence-preserving paths

The Eigenverse demonstrates both:
- Reality funnels to μ = e^(i3π/4) as the unique observer state
- Quantum tunneling emerges from coherence function C(r) = 2r/(1+r²)

## Coherence Scales

Three critical thresholds:

1. **Kernel (C = 1.0):** Maximum coherence, critical eigenvalue
2. **Silver (C = 0.707):** η = 1/√2, balanced μ-state amplitude  
3. **Golden (C = 0.667):** φ² ratio, Koide mass formula

## Statistics

- **Total Items:** {len(KB)}
- **Theorems:** {len([x for x in KB if x['kind'] == 'theorem'])}
- **Lemmas:** {len([x for x in KB if x['kind'] == 'lemma'])}
- **Definitions:** {len([x for x in KB if x['kind'] == 'def'])}
- **Modules:** {len(set(x['module'] for x in KB))}

## Key Theorems

- `reality_unique` — μ is the only solution in Q2 with balanced sectors
- `observer_fixed_point_unique` — C(1 + 1/η) = η has unique solution
- `fine_structure_constant_approx` — α ≈ 1/137 from μ¹³⁷
- `mu_eighth_power_unity` — μ⁸ = 1 (closure)

## Links

- **GitHub:** [beanapologist/Eigenverse](https://github.com/beanapologist/Eigenverse)
- **Formal Proofs:** 27 Lean 4 modules
- **Paper:** [Beach.Science](https://beach.science)

---

Built with Gradio • Powered by sklearn TF-IDF • Deployed on HuggingFace Spaces
"""

# Get unique modules for dropdown
modules = sorted(set(item['module'] for item in KB))

# Build Gradio interface
with gr.Blocks(theme=gr.themes.Soft(), title="Eigenverse Theorem Prover") as app:
    
    gr.Markdown("# 🧬 Eigenverse Theorem Prover")
    gr.Markdown("*Machine-verified mathematical framework. 688 theorems. Zero sorry.*")
    
    with gr.Tabs():
        
        # Tab 1: Hypothesis Prover
        with gr.Tab("🔍 Hypothesis Prover"):
            gr.Markdown("### Describe your hypothesis and get relevant theorems + LLM-powered proof draft")
            gr.Markdown("*Powered by Mistral-7B via HuggingFace Inference API. Set `HF_TOKEN` as a Space secret for better rate limits.*")
            
            hypothesis_input = gr.Textbox(
                label="Hypothesis",
                placeholder="e.g., The μ-state is stable under perturbations...",
                lines=3
            )
            
            prove_button = gr.Button("Find Theorems & Generate Proof Draft", variant="primary")
            
            with gr.Row():
                with gr.Column():
                    relevant_output = gr.Markdown(label="Relevant Theorems")
                with gr.Column():
                    scaffold_output = gr.Markdown(label="LLM Proof Draft")
            
            prove_button.click(
                hypothesis_prover_tab,
                inputs=[hypothesis_input],
                outputs=[relevant_output, scaffold_output]
            )
        
        # Tab 2: Theorem Explorer
        with gr.Tab("📚 Theorem Explorer"):
            gr.Markdown("### Browse and search the theorem database")
            
            with gr.Row():
                module_dropdown = gr.Dropdown(
                    choices=["All Modules"] + modules,
                    value="All Modules",
                    label="Filter by Module"
                )
                search_box = gr.Textbox(
                    label="Search",
                    placeholder="Enter keywords...",
                )
            
            search_button = gr.Button("Search", variant="primary")
            explorer_output = gr.Markdown()
            
            search_button.click(
                search_theorems,
                inputs=[module_dropdown, search_box],
                outputs=[explorer_output]
            )
        
        # Tab 3: Coherence Analyzer
        with gr.Tab("📊 Coherence Analyzer"):
            gr.Markdown("### Analyze coherence function C(r) = 2r/(1+r²)")
            
            r_input = gr.Number(
                label="Input value r",
                value=1.0
            )
            
            analyze_button = gr.Button("Analyze", variant="primary")
            
            coherence_output = gr.Markdown()
            coherence_plot = gr.Plot()
            
            analyze_button.click(
                coherence_analyzer_tab,
                inputs=[r_input],
                outputs=[coherence_output, coherence_plot]
            )
        
        # Tab 4: About
        with gr.Tab("ℹ️ About"):
            gr.Markdown(about_tab())

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