"""Hugging Face Space for dejanai/sentiment-classifier: paste text, see sentence-by-sentence sentiment
as a heatmap, or as a sortable table with CSV export."""
import base64
import csv
import html
import re
import tempfile
from pathlib import Path

import gradio as gr
from transformers import pipeline

MODEL = "dejanai/sentiment-classifier"
LABELS = ["very positive", "positive", "somewhat positive", "neutral", "somewhat negative", "negative", "very negative"]
# Diverging scale: blue = positive, grey = neutral, red = negative. Text colour keeps contrast on each.
COLORS = ["#2a78d6", "#6da7ec", "#b7d3f6", "#ecebe7", "#f6c4c2", "#ec8b89", "#d63c3b"]
TEXT = ["#ffffff", "#0b0b0b", "#0b0b0b", "#0b0b0b", "#0b0b0b", "#0b0b0b", "#ffffff"]
MAX_SEGMENTS = 500

LOGO = base64.b64encode((Path(__file__).parent / "dejan-logo.png").read_bytes()).decode()
clf = pipeline("text-classification", model=MODEL, device=-1)


def segments(text):
    """Split text into headings, list items, and sentences: the units the model was trained on."""
    for line in text.splitlines():
        line = re.sub(r"^\s*(#+|[-*+•]|\d+[.)])\s*", "", line).replace("**", "").strip()
        if line:
            yield from (s for s in re.split(r"(?<=[.!?])\s+(?=[A-Z0-9\"'])", line) if s)


def heatmap(parts, ids):
    spans = "".join(
        f'<span class="seg" title="{LABELS[i]}" style="background:{COLORS[i]};color:{TEXT[i]}">{html.escape(p)}</span> '
        for p, i in zip(parts, ids))
    legend = "".join(f'<span class="key"><i style="background:{c}"></i>{lab}</span>' for lab, c in zip(LABELS, COLORS))
    return f'<div class="legend">{legend}</div><div class="heat">{spans}</div>'


def analyse(text):
    parts = list(segments(text or ""))[:MAX_SEGMENTS]
    if not parts:
        return "", [], None, gr.update(visible=False)
    results = clf(parts, batch_size=32)
    ids = [LABELS.index(r["label"]) for r in results]
    rows = [[n + 1, p, LABELS[i], 3 - i] for n, (p, i) in enumerate(zip(parts, ids))]
    out = Path(tempfile.mkdtemp()) / "sentiment.csv"
    with open(out, "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["#", "sentence", "sentiment", "score"])
        w.writerows(rows)
    return heatmap(parts, ids), rows, str(out), gr.update(visible=True)


CSS = """
:root, .gradio-container {
  --color-accent: #3c7f12; --color-accent-soft: #eef5e8;
  --primary-50: #eef5e8; --primary-100: #dcebd0; --primary-200: #b9d7a1; --primary-300: #96c372;
  --primary-400: #6aa63d; --primary-500: #3c7f12; --primary-600: #336406; --primary-700: #2a5205;
  --input-border-color-focus: #3c7f12; --border-color-accent: #3c7f12; --border-color-accent-subdued: #b9d7a1;
  --loader-color: #3c7f12; --link-text-color: #3c7f12;
}
.tab-wrapper button.selected, button[role="tab"][aria-selected="true"] { color: #3c7f12 !important; }
.tab-wrapper button.selected::after, button[role="tab"][aria-selected="true"]::after { background: #3c7f12 !important; }
.gradio-container { background: #ffffff !important; max-width: 980px !important; margin: 0 auto !important; }
footer { display: none !important; }
#header { display: flex; align-items: center; gap: 16px; padding: 8px 0 0; }
#header img { height: 40px; }
#explainer { color: #52514e; font-size: 16px; margin: 4px 0 8px; }
#go { background: #3c7f12 !important; border-color: #3c7f12 !important; color: #ffffff !important; }
#go:hover { background: #336406 !important; }
.legend { display: flex; flex-wrap: wrap; gap: 12px; font-size: 12px; color: #52514e; margin-bottom: 12px; }
.legend .key { display: inline-flex; align-items: center; gap: 6px; }
.legend i { width: 14px; height: 14px; border-radius: 3px; display: inline-block; border: 1px solid #e3e2de; }
.heat { font-size: 16px; line-height: 2.1; color: #0b0b0b; }
.heat .seg { padding: 3px 4px; border-radius: 4px; box-decoration-break: clone; -webkit-box-decoration-break: clone; }
"""
FORCE_LIGHT = "() => { document.body.classList.remove('dark'); }"

with gr.Blocks(title="DEJAN AI Sentiment Classifier", css=CSS, theme=gr.themes.Default(), js=FORCE_LIGHT) as demo:
    gr.HTML(f'<div id="header"><a href="https://dejan.ai" target="_blank">'
            f'<img src="data:image/png;base64,{LOGO}" alt="DEJAN"></a></div>'
            '<div id="explainer">Paste any text to see the sentiment of every sentence, from very positive to very negative.</div>')
    text = gr.Textbox(show_label=False, lines=10, placeholder="Paste a review, a page, or any text...")
    go = gr.Button("Analyse", elem_id="go")
    with gr.Column(visible=False) as results:             # hidden until there are results
        with gr.Tabs():
            with gr.Tab("Heatmap"):
                heat = gr.HTML()
            with gr.Tab("Table"):
                table = gr.Dataframe(headers=["#", "Sentence", "Sentiment", "Score"],
                                     datatype=["number", "str", "str", "number"], wrap=True, interactive=False,
                                     column_widths=["6%", "64%", "18%", "12%"])
                export = gr.DownloadButton("Export CSV", value=None)
    go.click(analyse, inputs=text, outputs=[heat, table, export, results])

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