import os
import gc
import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
import spaces  # <--- 1. Import spaces for ZeroGPU

CONFIG = {
    "UNEXPECTED": {"repo_id": "JPQ24/Contrarian-GGUF", "filename": "Meta-Llama-3.1-8B-Instruct-Contrarian-Q4_K_M.gguf"},
    "LOGIC":      {"repo_id": "JPQ24/Logic-4-GGUF",      "filename": "Meta-Llama-3.1-8B-Instruct-logic-v4-Q4_K_M.gguf"},
    "ORGANIC":    {"repo_id": "JPQ24/Llama-3.1-8b-Natural-Synthesis-merged-GGUF",    "filename": "source_model.Q4_K_M.gguf"}
}

def run_gguf_inference(repo_id: str, filename: str, prompt: str, max_tokens: int = 768, temperature: float = 0.7) -> str:
    model_path = hf_hub_download(repo_id=repo_id, filename=filename)

    # 2. Enable full GPU offloading (n_gpu_layers=-1)
    llm = Llama(
        model_path=model_path,
        n_ctx=4096,
        n_gpu_layers=-1,  # <--- Offloads ALL layers to the A100 GPU
        verbose=False
    )

    response = llm(
        prompt=prompt,
        max_tokens=max_tokens,
        temperature=temperature,
        stop=["<|endoftext|>", "<|eot_id|>", "### END"]
    )
    
    output_text = response["choices"][0]["text"].strip()

    del llm
    gc.collect()
    return output_text


# 3. Add @spaces.GPU decorator (allocates the A100 for up to 90 seconds)
@spaces.GPU(duration=90)
def execute_triad_pipeline(context: str, problem: str, progress=gr.Progress()):
    # ... (Keep the rest of your pipeline logic identical) ...
    if not context.strip() or not problem.strip():
        raise gr.Error("Both Context Documents and Problem Statement are required.")

    # --------------------------------------------------------------------------
    # STAGE 1: UNEXPECTED (Grounded Lateral Inversion)
    # --------------------------------------------------------------------------
    progress(0.1, desc="⚡ Stage 1/4: Unexpected Engine (Extracting Grounded Inversions)...")
    
    stage1_prompt = f"""<|system|>
You are an adversarial lateral strategist and red-teamer.
STRICT CLOSED-WORLD GROUNDING RULES:
1. You MUST NOT invent any physical assets, infrastructure, or technology not explicitly stated in the CONTEXT.
2. Innovation MUST come strictly from inverting the FUNCTIONAL RELATIONSHIPS of existing context entities.
3. Explicitly list the Context Entities used before generating hypotheses.
<|user|>
[CONTEXT]
{context}

[PROBLEM]
{problem}

### TASK:
1. Identify the fatal traps and rebound failure modes of past attempts.
2. List available context entities.
3. Provide 3 grounded, high-leverage unconventional hypotheses.
<|assistant|>
"""
    stage1_output = run_gguf_inference(
        CONFIG["UNEXPECTED"]["repo_id"],
        CONFIG["UNEXPECTED"]["filename"],
        stage1_prompt,
        max_tokens=650,
        temperature=0.75
    )

    # --------------------------------------------------------------------------
    # STAGE 2: LOGIC BOUNDS (Formal Axioms & Measurable Criteria)
    # --------------------------------------------------------------------------
    progress(0.35, desc="📐 Stage 2/4: Logic Engine (Formulating Constraint Matrix)...")

    stage2_prompt = f"""<|system|>
You are a formal deductive logic engine. Convert qualitative requirements and proposed hypotheses into hard mathematical bounds, formal axioms, and verifiable success criteria (C1..Cn).
<|user|>
[CONTEXT]
{context}

[LATERAL HYPOTHESES]
{stage1_output}

### TASK:
1. Define Formal Axioms (P -> Q).
2. Set Feasibility & Grounding Boundaries (reject any non-context entities).
3. Establish Mandatory Success Criteria (C1..Cn).
<|assistant|>
"""
    stage2_output = run_gguf_inference(
        CONFIG["LOGIC"]["repo_id"],
        CONFIG["LOGIC"]["filename"],
        stage2_prompt,
        max_tokens=500,
        temperature=0.1
    )

    # --------------------------------------------------------------------------
    # STAGE 3: ORGANIC SYNTHESIS (Biomimetic Contrastive Pruning)
    # --------------------------------------------------------------------------
    progress(0.65, desc="🌿 Stage 3/4: Organic Engine (Biomimetic CoT & Pruning)...")

    stage3_prompt = f"""<|system|>
You are a biomimetic systems architect. Reason using the 5-phase organic cognitive trace:
1. Seed (Conceptual Core)
2. Germination (Divergent Loops)
3. Selective Nourishment (Explicitly NOURISH viable paths, explicitly WITHER weak/ungrounded paths)
4. Canopy Formation (Convergent System Design)
5. Homeostatic Review (Shock testing against heatwaves, floods, economic stress)
<|user|>
[CONTEXT]
{context}

[CRITERIA & BOUNDS]
{stage2_output}

[HYPOTHESES]
{stage1_output}

### TASK:
Produce the <thought> trace using the 5-phase method, then formulate the unified Multi-Component Mechanism.
<|assistant|>
"""
    stage3_output = run_gguf_inference(
        CONFIG["ORGANIC"]["repo_id"],
        CONFIG["ORGANIC"]["filename"],
        stage3_prompt,
        max_tokens=900,
        temperature=0.35
    )

    # --------------------------------------------------------------------------
    # STAGE 4: LOGICAL AUDIT & FINAL VERDICT
    # --------------------------------------------------------------------------
    progress(0.9, desc="⚖️ Stage 4/4: Logic Auditor (Verifying C1..Cn Compliance)...")

    stage4_prompt = f"""<|system|>
You are a ruthless logical verification auditor. Check the synthesized plan against each criterion (C1..Cn) defined in the Criteria list. Check for contradictions and ungrounded hallucinations.
<|user|>
[CRITERIA (C1..Cn)]
{stage2_output}

[PROPOSED SYSTEM PLAN]
{stage3_output}

### TASK:
1. Audit each criterion: [SATISFIED] or [FAILED] with justification.
2. Contradiction & Grounding Scan.
3. Logical Verdict (VALID & SOUND or INVALID).
<|assistant|>
"""
    stage4_output = run_gguf_inference(
        CONFIG["LOGIC"]["repo_id"],
        CONFIG["LOGIC"]["filename"],
        stage4_prompt,
        max_tokens=400,
        temperature=0.05
    )

    progress(1.0, desc="Execution Complete.")

    final_dashboard = f"""==============================================================================
FINAL AUDITED RESULT
==============================================================================
{stage4_output}

==============================================================================
FINAL SYSTEMIC INTERVENTION PLAN
==============================================================================
{stage3_output}
"""
    return stage1_output, stage2_output, stage3_output, final_dashboard


# ==============================================================================
# GRADIO UI SETUP
# ==============================================================================

sample_context = """[DOC 1: Hydrogeology & Climate Risk Report - District 9]
- The district faces severe aquifer depletion: water extraction currently exceeds recharge by 35% annually.
- Climate modeling projects a 40% increase in 5-day extreme heatwave events and a 20% risk of sudden flash floods in the low-lying southern valley.
- Emergency shock capacity: Current municipal water storage lasts only 48 hours in the event of an electrical grid blackout.

[DOC 2: Socio-Economic & Vulnerability Matrix]
- 70% of the population are smallholder farmers (under 2 hectares) living on less than $2.50/day.
- Smallholders rely on high-interest informal credit (40-60% annual rates) to buy diesel fuel for irrigation pumps. When diesel prices spike or pumps fail, default rates surge, causing land foreclosures.
- 30% of agricultural land is owned by 3 commercial agribusinesses that hold prioritized legal water rights and private solar arrays.

[DOC 3: Technology & Policy Pilot Evaluations]
- Pilot A (Subsidized Solar Pumps): Replaced diesel pumps with free solar pumps. Result: Farmers expanded pumping hours due to zero marginal energy cost, increasing aquifer depletion by an additional 18%.
- Pilot B (Automated Micro-Drip with Smart Metering): Reduced water consumption by 50% and doubled crop yields, but requires a continuous 4G/IoT connection and $1,200/hectare upfront capital cost.
- Pilot C (Community Water Quota & Insurance Co-op): Successfully prevented over-extraction, but collapsed within 8 months due to lack of transparent enforcement and elite capture by commercial agribusiness owners."""

sample_problem = """District 9 suffers from a compound crisis: extreme poverty among smallholders, escalating climate disaster risks (heatwaves/floods), and acute resource collapse (aquifer exhaustion). 

Design a systemic intervention that simultaneously:
1. Eliminates the smallholder debt-poverty trap caused by energy/pumping costs.
2. Halts aquifer over-extraction without reducing smallholder caloric/income baseline.
3. Establishes a localized crisis buffer for water/energy during sudden grid outages or climate shocks.

Constraints:
- You must ground your solution in the provided documents.
- You must explicitly avoid the failure modes identified in Pilots A, B, and C."""

theme = gr.themes.Monochrome(
    primary_hue="emerald",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("JetBrains Mono"), "monospace"]
)

with gr.Blocks(title="Triad: Cognitive SLM Pipeline", theme=theme) as demo:
    gr.Markdown("# 🧠 Triad: SLM Multi-Agent Cognitive Architecture")
    gr.Markdown(
        "Sequential orchestration of three specialized 8B parameter models (*Unexpected*, *Logic*, and *Organic*) "
        "designed to solve multi-constraint systemic problems without monolithic parameter bloat."
    )

    with gr.Row():
        with gr.Column(scale=5):
            context_input = gr.Textbox(
                label="Knowledge Base / Context Documents",
                lines=10,
                placeholder="Paste reference documents here..."
            )
            problem_input = gr.Textbox(
                label="Problem Statement & Constraints",
                lines=5,
                placeholder="Define the problem, objectives, and constraints..."
            )
            run_btn = gr.Button("🚀 Run Triad Pipeline", variant="primary", size="lg")

        with gr.Column(scale=6):
            with gr.Accordion("⚡ Stage 1: Contrarian Traps & Inversions (Unexpected Engine)", open=False):
                s1_out = gr.Markdown()
            
            with gr.Accordion("📐 Stage 2: Formal Axioms & Success Bounds (Logic Engine)", open=False):
                s2_out = gr.Markdown()

            with gr.Accordion("🌿 Stage 3: Biomimetic CoT & Synthesis (Organic Engine)", open=False):
                s3_out = gr.Markdown()

            gr.Markdown("### 📋 Final Audited System Plan")
            final_out = gr.Markdown()

    # Preset Example Click
    gr.Examples(
        examples=[[sample_context, sample_problem]],
        inputs=[context_input, problem_input],
        label="Load Preset Case Study (District 9 Crisis)"
    )

    run_btn.click(
        fn=execute_triad_pipeline,
        inputs=[context_input, problem_input],
        outputs=[s1_out, s2_out, s3_out, final_out]
    )

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