"""
ChatSpatial Engine — Gemma 4 Spatial Omics
Multimodal spatial transcriptomics analysis powered by Gemma 4 31B + LoRA.
"""

import torch
import gradio as gr
from PIL import Image
import os

from visualizations import generate_all_plots, fig_to_pil

# ---------------------------------------------------------------------------
# Model Loading (runs once at startup)
# ---------------------------------------------------------------------------
print("Loading model... This may take a few minutes on first boot.")

from unsloth import FastVisionModel

model, tokenizer = FastVisionModel.from_pretrained(
    "arka2696/gemma-4-spatial-omics-lora-v3",
    load_in_4bit=True,
)
FastVisionModel.for_inference(model)

# Set chat template for multimodal messages
CHAT_TEMPLATE = (
    "{% for message in messages %}"
    "{% if message['role'] == 'user' %}{{ '<start_of_turn>user\\n' }}"
    "{% elif message['role'] == 'model' %}{{ '<start_of_turn>model\\n' }}"
    "{% endif %}"
    "{% if message['content'] is string %}{{ message['content'] }}"
    "{% else %}{% for block in message['content'] %}"
    "{% if block['type'] == 'image' %}{{ '<|image|>' }}"
    "{% elif block['type'] == 'text' %}{{ block['text'] }}"
    "{% endif %}{% endfor %}{% endif %}"
    "{{ '<end_of_turn>\\n' }}"
    "{% endfor %}"
)
tokenizer.chat_template = CHAT_TEMPLATE

SYSTEM_PROMPT = (
    "You are ChatSpatial, an expert spatial transcriptomics analyst. "
    "You analyze H&E-stained tissue microscopy images and predict gene expression "
    "patterns based on cellular morphology, tissue architecture, and spatial context. "
    "Always reason step-by-step about what you observe in the tissue before making predictions. "
    "Use <|think|> tags for your internal reasoning."
)

print("Model loaded successfully!")

# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------

def _load_phoenix():
    """Load Phoenix engine (optional — fails gracefully)."""
    try:
        from phoenix_engine import PhoenixEngine
        from model_loader import download_phoenix_weights
        model_dir = "phoenix_weights"
        if not os.path.exists(model_dir):
            download_phoenix_weights(model_dir)
        engine = PhoenixEngine(model_dir=model_dir, device="cpu", num_samples=3)
        engine._load()
        print("Phoenix engine loaded!", flush=True)
        return engine
    except Exception as e:
        print(f"Phoenix not available: {e}", flush=True)
        return None

phoenix_engine = _load_phoenix()


def analyze(image: Image.Image, question: str, h5ad_file=None):
    """Run inference on an image + question. Returns (analysis_md, bar_img, radar_img, heatmap_img)."""
    if image is None:
        return "Please upload an H&E tissue image to analyze.", None, None, None
    if not question.strip():
        question = "Analyze this tissue image. What cell types and gene expression patterns do you observe?"

    prompt = question
    bar_img, radar_img, heatmap_img = None, None, None

    # Phoenix prediction (if available + image provided)
    if phoenix_engine is not None:
        try:
            phoenix_result = phoenix_engine.predict(image)
            if phoenix_result:
                phoenix_text = phoenix_result.get("summary_text", "")
                prompt += (
                    "\n\nPhoenix quantitative expression prediction results:\n"
                    + phoenix_text
                    + "\n\nIntegrate these quantitative results with your morphological "
                    "observations. Add biological interpretation beyond restating numbers."
                )
                plots = generate_all_plots(phoenix_result)
                if "bar_chart" in plots:
                    bar_img = fig_to_pil(plots["bar_chart"])
                if "radar" in plots:
                    radar_img = fig_to_pil(plots["radar"])
                if "heatmap" in plots:
                    heatmap_img = fig_to_pil(plots["heatmap"])
        except Exception as e:
            print(f"Phoenix error: {e}", flush=True)

    # Build messages with image object directly (Unsloth's expected format)
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": f"{SYSTEM_PROMPT}\n\n{prompt}"},
            ],
        }
    ]

    # Tokenize (Unsloth handles image processing internally)
    inputs = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    )
    inputs = {k: v.to(model.device) for k, v in inputs.items()}

    # Generate
    with torch.inference_mode():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=1024,
            temperature=0.6,
            top_p=0.9,
            do_sample=True,
            use_cache=True,
        )

    # Decode only new tokens
    generated = output_ids[0][inputs["input_ids"].shape[1]:]
    response = tokenizer.decode(generated, skip_special_tokens=True)

    # Format reasoning traces nicely
    response = format_reasoning(response)
    return response, bar_img, radar_img, heatmap_img


def format_reasoning(text: str) -> str:
    """Format <|think|> blocks as collapsible markdown sections."""
    import re

    # Extract think blocks and format as details/summary
    def replace_think(match):
        reasoning = match.group(1).strip()
        return (
            f"\n<details>\n<summary>💭 <b>Reasoning Trace</b></summary>\n\n"
            f"{reasoning}\n\n</details>\n\n"
        )

    text = re.sub(
        r"<\|think\|>(.*?)<\|/think\|>",
        replace_think,
        text,
        flags=re.DOTALL,
    )
    return text.strip()


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

CSS = """
/* ===== Design Tokens ===== */
:root {
    --bg-primary: #0f172a;
    --bg-secondary: #1e293b;
    --bg-card: #1e293b;
    --bg-card-hover: #253349;
    --border-subtle: rgba(255, 255, 255, 0.06);
    --border-active: rgba(59, 130, 246, 0.4);
    --accent-blue: #3b82f6;
    --accent-blue-dim: rgba(59, 130, 246, 0.15);
    --accent-green: #10b981;
    --accent-green-dim: rgba(16, 185, 129, 0.15);
    --accent-amber: #f59e0b;
    --accent-amber-dim: rgba(245, 158, 11, 0.15);
    --text-primary: #f1f5f9;
    --text-secondary: #94a3b8;
    --text-muted: #64748b;
    --font-sans: -apple-system, BlinkMacSystemFont, 'Segoe UI', Inter, Roboto, sans-serif;
    --font-mono: 'JetBrains Mono', 'Fira Code', 'SF Mono', Consolas, monospace;
    --radius-sm: 6px;
    --radius-md: 10px;
    --radius-lg: 14px;
    --shadow-card: 0 1px 3px rgba(0,0,0,0.3), 0 1px 2px rgba(0,0,0,0.2);
    --shadow-elevated: 0 4px 12px rgba(0,0,0,0.4);
}

/* ===== Global ===== */
.gradio-container {
    background: var(--bg-primary) !important;
    font-family: var(--font-sans) !important;
    max-width: 1440px !important;
    color: var(--text-primary) !important;
}

.gradio-container .gr-block,
.gradio-container .gr-box,
.gradio-container .gr-panel {
    background: transparent !important;
}

footer { display: none !important; }

/* ===== Top Navbar ===== */
.navbar {
    display: flex;
    align-items: center;
    justify-content: space-between;
    padding: 12px 24px;
    background: var(--bg-secondary);
    border-bottom: 1px solid var(--border-subtle);
    border-radius: var(--radius-lg);
    margin-bottom: 16px;
}

.navbar-brand {
    display: flex;
    align-items: center;
    gap: 10px;
    font-size: 1.1rem;
    font-weight: 700;
    color: var(--text-primary);
    letter-spacing: -0.02em;
}

.navbar-brand .dot {
    color: var(--accent-green);
    font-size: 1.3rem;
    line-height: 1;
}

.navbar-right {
    display: flex;
    align-items: center;
    gap: 8px;
}

.nav-badge {
    display: inline-flex;
    align-items: center;
    gap: 5px;
    padding: 4px 10px;
    border-radius: 20px;
    font-size: 0.72rem;
    font-weight: 600;
    font-family: var(--font-mono);
    letter-spacing: 0.02em;
    text-transform: uppercase;
}

.nav-badge.blue {
    background: var(--accent-blue-dim);
    color: var(--accent-blue);
    border: 1px solid rgba(59, 130, 246, 0.25);
}

.nav-badge.green {
    background: var(--accent-green-dim);
    color: var(--accent-green);
    border: 1px solid rgba(16, 185, 129, 0.25);
}

.nav-badge.amber {
    background: var(--accent-amber-dim);
    color: var(--accent-amber);
    border: 1px solid rgba(245, 158, 11, 0.25);
}

.status-indicator {
    display: inline-flex;
    align-items: center;
    gap: 6px;
    padding: 4px 12px;
    border-radius: 20px;
    font-size: 0.75rem;
    font-weight: 500;
    background: var(--accent-green-dim);
    color: var(--accent-green);
    border: 1px solid rgba(16, 185, 129, 0.3);
}

.status-indicator .status-dot {
    width: 7px;
    height: 7px;
    background: var(--accent-green);
    border-radius: 50%;
    animation: pulse 2s ease-in-out infinite;
}

@keyframes pulse {
    0%, 100% { opacity: 1; }
    50% { opacity: 0.5; }
}

/* ===== Card Containers ===== */
.ui-card {
    background: var(--bg-card) !important;
    border: 1px solid var(--border-subtle) !important;
    border-radius: var(--radius-lg) !important;
    padding: 20px !important;
    box-shadow: var(--shadow-card) !important;
    transition: border-color 0.2s ease !important;
}

.ui-card:hover {
    border-color: rgba(255, 255, 255, 0.1) !important;
}

.card-header {
    display: flex;
    align-items: center;
    gap: 8px;
    margin-bottom: 12px;
    font-size: 0.82rem;
    font-weight: 600;
    color: var(--text-secondary);
    text-transform: uppercase;
    letter-spacing: 0.05em;
}

.card-header .card-icon {
    font-size: 0.9rem;
}

/* ===== Inputs Styling ===== */
.gradio-container textarea,
.gradio-container input[type="text"] {
    background: var(--bg-primary) !important;
    border: 1px solid var(--border-subtle) !important;
    border-radius: var(--radius-md) !important;
    color: var(--text-primary) !important;
    font-family: var(--font-sans) !important;
    font-size: 0.9rem !important;
    transition: border-color 0.2s ease !important;
}

.gradio-container textarea:focus,
.gradio-container input[type="text"]:focus {
    border-color: var(--accent-blue) !important;
    box-shadow: 0 0 0 3px var(--accent-blue-dim) !important;
}

/* ===== Primary Button ===== */
#run-btn {
    background: linear-gradient(135deg, #3b82f6 0%, #2563eb 100%) !important;
    border: none !important;
    border-radius: var(--radius-md) !important;
    color: #ffffff !important;
    font-weight: 600 !important;
    font-size: 0.92rem !important;
    padding: 12px 28px !important;
    letter-spacing: 0.01em !important;
    transition: all 0.15s ease !important;
    box-shadow: 0 2px 8px rgba(59, 130, 246, 0.25) !important;
}

#run-btn:hover {
    transform: translateY(-1px) !important;
    box-shadow: 0 6px 20px rgba(59, 130, 246, 0.35) !important;
}

#run-btn:active {
    transform: translateY(0) !important;
}

/* ===== Tabs ===== */
.gradio-container .tabs {
    background: transparent !important;
}

.gradio-container .tab-nav {
    border-bottom: 1px solid var(--border-subtle) !important;
    background: transparent !important;
    gap: 0 !important;
}

.gradio-container .tab-nav button {
    background: transparent !important;
    border: none !important;
    border-bottom: 2px solid transparent !important;
    color: var(--text-muted) !important;
    font-weight: 500 !important;
    font-size: 0.85rem !important;
    padding: 10px 18px !important;
    transition: all 0.15s ease !important;
}

.gradio-container .tab-nav button:hover {
    color: var(--text-secondary) !important;
}

.gradio-container .tab-nav button.selected {
    color: var(--accent-blue) !important;
    border-bottom-color: var(--accent-blue) !important;
}

/* ===== Output Markdown ===== */
.output-markdown {
    background: var(--bg-primary) !important;
    border: 1px solid var(--border-subtle) !important;
    border-radius: var(--radius-md) !important;
    padding: 24px !important;
    color: var(--text-primary) !important;
    font-size: 0.9rem !important;
    line-height: 1.7 !important;
    min-height: 300px !important;
}

.output-markdown code {
    font-family: var(--font-mono) !important;
    background: rgba(59, 130, 246, 0.1) !important;
    padding: 2px 6px !important;
    border-radius: 4px !important;
    font-size: 0.82rem !important;
}

.output-markdown details {
    background: rgba(16, 185, 129, 0.05) !important;
    border: 1px solid rgba(16, 185, 129, 0.15) !important;
    border-radius: var(--radius-sm) !important;
    padding: 12px !important;
    margin: 12px 0 !important;
}

.output-markdown details summary {
    cursor: pointer;
    color: var(--accent-green) !important;
    font-weight: 600 !important;
}

/* ===== Image Upload ===== */
.gradio-container .image-container,
.gradio-container .upload-container {
    border: 1px dashed var(--border-subtle) !important;
    border-radius: var(--radius-md) !important;
    background: var(--bg-primary) !important;
    transition: border-color 0.2s ease !important;
}

.gradio-container .image-container:hover,
.gradio-container .upload-container:hover {
    border-color: var(--accent-blue) !important;
}

/* ===== Plot Outputs ===== */
.gradio-container .plot-container,
.gradio-container .image-preview {
    border-radius: var(--radius-md) !important;
    overflow: hidden !important;
}

/* ===== Section Label ===== */
.section-label {
    font-size: 0.75rem;
    font-weight: 600;
    color: var(--text-muted);
    text-transform: uppercase;
    letter-spacing: 0.06em;
    margin-bottom: 10px;
    padding-left: 2px;
}

/* ===== Footer ===== */
.app-footer {
    text-align: center;
    padding: 20px 16px;
    margin-top: 24px;
    border-top: 1px solid var(--border-subtle);
}

.app-footer .footer-text {
    color: var(--text-muted);
    font-size: 0.78rem;
    letter-spacing: 0.01em;
}

.footer-badges {
    display: flex;
    justify-content: center;
    gap: 6px;
    margin-top: 8px;
    flex-wrap: wrap;
}

.footer-badge {
    display: inline-flex;
    align-items: center;
    padding: 3px 8px;
    border-radius: 4px;
    font-size: 0.68rem;
    font-family: var(--font-mono);
    font-weight: 500;
    background: rgba(255, 255, 255, 0.04);
    color: var(--text-muted);
    border: 1px solid var(--border-subtle);
}

/* ===== Dropdown ===== */
.gradio-container .dropdown-container,
.gradio-container select {
    background: var(--bg-primary) !important;
    border: 1px solid var(--border-subtle) !important;
    border-radius: var(--radius-md) !important;
    color: var(--text-primary) !important;
}

/* ===== Label styling ===== */
.gradio-container label {
    color: var(--text-secondary) !important;
    font-size: 0.82rem !important;
    font-weight: 500 !important;
}

/* ===== Responsive ===== */
@media (max-width: 768px) {
    .navbar {
        flex-direction: column;
        gap: 10px;
        text-align: center;
    }
    .navbar-right {
        flex-wrap: wrap;
        justify-content: center;
    }
}
"""

EXAMPLES = [
    "Analyze this tissue patch. What cell types do you observe based on morphology?",
    "Predict the expression levels of CD8A, EPCAM, and COL1A1 in this region.",
    "Is this region likely tumor, stroma, or immune-infiltrated? Explain your reasoning.",
    "What spatial gene expression patterns would you expect in this tissue architecture?",
    "Identify the dominant cell population and predict marker gene expression.",
]

# Determine component status for navbar
_model_status = "Ready" if model is not None else "Loading"
_phoenix_status = "Active" if phoenix_engine is not None else "Inactive"

with gr.Blocks(css=CSS, theme=gr.themes.Base(), title="ChatSpatial Engine") as demo:

    # ===== Top Navbar =====
    gr.HTML(f"""
    <div class="navbar">
        <div class="navbar-brand">
            <span class="dot">&#9673;</span>
            ChatSpatial Engine
        </div>
        <div class="navbar-right">
            <span class="nav-badge blue">Gemma 4 31B</span>
            <span class="nav-badge green">Phoenix</span>
            <span class="nav-badge amber">QLoRA</span>
            <div class="status-indicator">
                <span class="status-dot"></span>
                {_model_status}
            </div>
        </div>
    </div>
    """)

    # ===== Main Layout =====
    with gr.Row(equal_height=False):

        # ----- Left Sidebar (Inputs) -----
        with gr.Column(scale=4):

            # Image Upload Card
            gr.HTML("""
            <div class="card-header">
                <span class="card-icon">&#x1F52C;</span> H&E Tissue Patch
            </div>
            """)
            image_input = gr.Image(
                type="pil",
                label="Drop or click to upload tissue image",
                height=260,
                elem_classes=["ui-card"],
            )

            gr.HTML("<div style='height: 12px;'></div>")

            # Question Card
            gr.HTML("""
            <div class="card-header">
                <span class="card-icon">&#x2753;</span> Analysis Query
            </div>
            """)
            question_input = gr.Textbox(
                label="Question",
                placeholder="Ask about cell types, gene expression, tissue architecture...",
                lines=3,
            )
            example_dropdown = gr.Dropdown(
                choices=EXAMPLES,
                label="Example prompts",
                interactive=True,
            )

            gr.HTML("<div style='height: 12px;'></div>")

            # h5ad optional upload
            gr.HTML("""
            <div class="card-header">
                <span class="card-icon">&#x1F4CA;</span> Spatial Data (optional)
            </div>
            """)
            h5ad_input = gr.File(
                label=".h5ad spatial matrix",
                file_types=[".h5ad"],
            )

            gr.HTML("<div style='height: 16px;'></div>")

            # Run Button
            analyze_btn = gr.Button(
                "▶  Run Analysis",
                variant="primary",
                elem_id="run-btn",
            )

        # ----- Right Panel (Outputs) -----
        with gr.Column(scale=6):
            with gr.Tabs():
                # Tab: Analysis
                with gr.Tab("Analysis"):
                    output_md = gr.Markdown(
                        value="*Upload an H&E tissue image and click* **Run Analysis** *to begin.*",
                        elem_classes=["output-markdown"],
                    )

                # Tab: Expression Profile
                with gr.Tab("Expression Profile"):
                    gr.HTML('<div class="section-label">Gene Expression Visualization</div>')
                    with gr.Row():
                        bar_output = gr.Image(label="Top Expressed Genes", height=320)
                        radar_output = gr.Image(label="Tissue Composition Radar", height=320)
                    heatmap_output = gr.Image(label="Marker Gene Heatmap", height=200)

                # Tab: About
                with gr.Tab("About"):
                    gr.Markdown("""
### Methodology

**ChatSpatial Engine** is a multimodal spatial transcriptomics analysis system built on:

- **Gemma 4 31B Dense** fine-tuned with QLoRA (4-bit quantization) on the STimage-1K4M dataset
- **Phoenix Flow-Matching** for quantitative gene expression prediction from morphology
- **Chain-of-thought reasoning** via `<|think|>` traces for interpretable biological analysis

**Training Data**: 4M+ spatial transcriptomics spots across 1,171 H&E-stained tissue slides
(Visium, ST, VisiumHD technologies) with matched gene expression profiles.

**Marker Panel**: CD8A, CD8B, CD3D, CD4, MS4A1, CD19, CD68, CD163, PTPRC, EPCAM, KRT18,
MKI67, COL1A1, VIM, ACTA2, VEGFA, PDCD1

**Pipeline**: Image patch (224x224) -> Vision encoder (1120 tokens max) -> Gemma 4 reasoning
-> Expression prediction + biological interpretation

---

*Built for the Gemma 4 Good Hackathon (Kaggle)*
                    """)

    # ===== Wire Events =====
    analyze_btn.click(
        fn=analyze,
        inputs=[image_input, question_input, h5ad_input],
        outputs=[output_md, bar_output, radar_output, heatmap_output],
    )

    example_dropdown.change(
        fn=lambda x: x,
        inputs=example_dropdown,
        outputs=question_input,
    )

    # ===== Footer =====
    gr.HTML("""
    <div class="app-footer">
        <div class="footer-text">Powered by</div>
        <div class="footer-badges">
            <span class="footer-badge">Gemma 4 31B</span>
            <span class="footer-badge">STimage-1K4M</span>
            <span class="footer-badge">Unsloth QLoRA</span>
            <span class="footer-badge">Phoenix Flow-Matching</span>
            <span class="footer-badge">HuggingFace Spaces</span>
        </div>
    </div>
    """)

# ---------------------------------------------------------------------------
# Launch
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)