contribution_heatmap.py

"""
🔥 GitHub-Style Contribution Heatmap — Custom Gradio 6 Component
Built entirely with gr.HTML: dynamic templates, CSS, JS interactivity, and custom props.
"""

import gradio as gr
import random
import math
from datetime import datetime, timedelta

# ── Color Schemes ────────────────────────────────────────────────────────────
COLOR_SCHEMES = {
    "green":  ["#161b22", "#0e4429", "#006d32", "#26a641", "#39d353"],
    "blue":   ["#161b22", "#0a3069", "#0550ae", "#0969da", "#54aeff"],
    "purple": ["#161b22", "#3b1f72", "#6639a6", "#8957e5", "#bc8cff"],
    "orange": ["#161b22", "#6e3a07", "#9a5b13", "#d4821f", "#f0b040"],
    "pink":   ["#161b22", "#5c1a3a", "#8b2252", "#d63384", "#f472b6"],
    "red":    ["#161b22", "#6e1007", "#9a2013", "#d4401f", "#f06040"],
}

# ── HTML / CSS / JS Templates (shared across all instances) ──────────────────

HEATMAP_HTML = """
<div class="heatmap-container">
    <div class="heatmap-header">
        <h2>${(() => {
            const total = Object.values(value || {}).reduce((a, b) => a + b, 0);
            return '📊 ' + total.toLocaleString() + ' contributions in ' + year;
        })()}</h2>
        <div class="legend">
            <span>Less</span>
            <div class="legend-box" style="background:${c0}"></div>
            <div class="legend-box" style="background:${c1}"></div>
            <div class="legend-box" style="background:${c2}"></div>
            <div class="legend-box" style="background:${c3}"></div>
            <div class="legend-box" style="background:${c4}"></div>
            <span>More</span>
        </div>
    </div>
    <div class="month-labels">
        ${(() => {
            const months = ['Jan','Feb','Mar','Apr','May','Jun','Jul','Aug','Sep','Oct','Nov','Dec'];
            return months.map((m, i) =>
                '<span style="grid-column:' + (Math.floor(i * 4.33) + 2) + '">' + m + '</span>'
            ).join('');
        })()}
    </div>
    <div class="heatmap-grid">
        <div class="day-labels">
            <span></span><span>Mon</span><span></span><span>Wed</span><span></span><span>Fri</span><span></span>
        </div>
        <div class="cells">
            ${(() => {
                const v = value || {};
                const sd = new Date(year, 0, 1);
                const pad = sd.getDay();
                const cells = [];
                for (let i = 0; i < pad; i++) cells.push('<div class="cell empty"></div>');
                const totalDays = Math.floor((new Date(year, 11, 31) - sd) / 86400000) + 1;
                for (let d = 0; d < totalDays; d++) {
                    const dt = new Date(year, 0, 1 + d);
                    const key = dt.getFullYear() + '-' + String(dt.getMonth()+1).padStart(2,'0') + '-' + String(dt.getDate()).padStart(2,'0');
                    const count = v[key] || 0;
                    let lv = 0;
                    if (count > 0) lv = 1;
                    if (count >= 3) lv = 2;
                    if (count >= 6) lv = 3;
                    if (count >= 10) lv = 4;
                    const mn = ['Jan','Feb','Mar','Apr','May','Jun','Jul','Aug','Sep','Oct','Nov','Dec'];
                    const tip = count + ' contributions on ' + mn[dt.getMonth()] + ' ' + dt.getDate() + ', ' + year;
                    cells.push('<div class="cell level-' + lv + '" data-date="' + key + '" data-count="' + count + '" title="' + tip + '"></div>');
                }
                return cells.join('');
            })()}
        </div>
    </div>
    <div class="stats-bar">
        ${(() => {
            const v = value || {};
            const totalDays = Math.floor((new Date(year, 11, 31) - new Date(year, 0, 1)) / 86400000) + 1;
            let streak = 0, maxStreak = 0, total = 0, active = 0, best = 0;
            const vals = [];
            for (let d = 0; d < totalDays; d++) {
                const dt = new Date(year, 0, 1 + d);
                const key = dt.getFullYear() + '-' + String(dt.getMonth()+1).padStart(2,'0') + '-' + String(dt.getDate()).padStart(2,'0');
                const c = v[key] || 0;
                total += c;
                if (c > 0) { streak++; maxStreak = Math.max(maxStreak, streak); active++; best = Math.max(best, c); vals.push(c); }
                else { streak = 0; }
            }
            const avg = vals.length ? (total / vals.length).toFixed(1) : '0';
            const stats = [
                ['🔥', maxStreak, 'Longest Streak'],
                ['📅', active, 'Active Days'],
                ['⚡', best, 'Best Day'],
                ['📈', avg, 'Avg / Active Day'],
                ['🏆', total.toLocaleString(), 'Total'],
            ];
            return stats.map(s =>
                '<div class="stat"><span class="stat-value">' + s[1] + '</span><span class="stat-label">' + s[0] + ' ' + s[2] + '</span></div>'
            ).join('');
        })()}
    </div>
</div>
"""

HEATMAP_CSS = """
    .heatmap-container {
        background: #0d1117;
        border-radius: 12px;
        padding: 24px;
        font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
        color: #c9d1d9;
        overflow-x: auto;
    }
    .heatmap-header {
        display: flex; justify-content: space-between; align-items: center;
        margin-bottom: 12px; flex-wrap: wrap; gap: 10px;
    }
    .heatmap-header h2 { margin: 0; font-size: 16px; font-weight: 500; color: #f0f6fc; }
    .legend { display: flex; align-items: center; gap: 4px; font-size: 11px; color: #8b949e; }
    .legend-box { width: 12px; height: 12px; border-radius: 2px; }
    .month-labels {
        display: grid; grid-template-columns: 30px repeat(53, 1fr);
        font-size: 11px; color: #8b949e; margin-bottom: 4px; padding-left: 2px;
    }
    .heatmap-grid { display: flex; gap: 4px; }
    .day-labels {
        display: grid; grid-template-rows: repeat(7, 1fr);
        font-size: 11px; color: #8b949e; width: 30px; gap: 2px;
    }
    .day-labels span { height: 13px; display: flex; align-items: center; }
    .cells {
        display: grid; grid-template-rows: repeat(7, 1fr);
        grid-auto-flow: column; gap: 2px; flex: 1;
    }
    .cell {
        width: 13px; height: 13px; border-radius: 2px; cursor: pointer;
        transition: all 0.15s ease; outline: 1px solid rgba(27,31,35,0.06);
    }
    .cell:hover {
        outline: 2px solid #58a6ff; outline-offset: -1px;
        transform: scale(1.3); z-index: 1;
    }
    .cell.empty { visibility: hidden; }
    .level-0 { background: ${c0}; }
    .level-1 { background: ${c1}; }
    .level-2 { background: ${c2}; }
    .level-3 { background: ${c3}; }
    .level-4 { background: ${c4}; }
    .stats-bar {
        display: flex; justify-content: space-around; margin-top: 20px;
        padding-top: 16px; border-top: 1px solid #21262d;
        flex-wrap: wrap; gap: 12px;
    }
    .stat { display: flex; flex-direction: column; align-items: center; gap: 4px; }
    .stat-value { font-size: 22px; font-weight: 700; color: ${c4}; }
    .stat-label { font-size: 12px; color: #8b949e; }
"""

HEATMAP_JS = """
    element.addEventListener('click', (e) => {
        if (e.target && e.target.classList.contains('cell') && !e.target.classList.contains('empty')) {
            const date = e.target.dataset.date;
            const cur = parseInt(e.target.dataset.count) || 0;
            const next = cur >= 12 ? 0 : cur + 1;
            const nv = {...(props.value || {})};
            if (next === 0) delete nv[date]; else nv[date] = next;
            props.value = nv;
            trigger('change');
        }
    });
"""


# ── Component Class ──────────────────────────────────────────────────────────

class ContributionHeatmap(gr.HTML):
    """Reusable GitHub-style contribution heatmap built on gr.HTML."""

    def __init__(self, value=None, year=2025, theme="green",
                 c0=None, c1=None, c2=None, c3=None, c4=None, **kwargs):
        if value is None:
            value = {}
        # Use explicit c0-c4 if provided (from gr.HTML updates), else derive from theme
        colors = COLOR_SCHEMES.get(theme, COLOR_SCHEMES["green"])
        super().__init__(
            value=value,
            year=year,
            c0=c0 or colors[0],
            c1=c1 or colors[1],
            c2=c2 or colors[2],
            c3=c3 or colors[3],
            c4=c4 or colors[4],
            html_template=HEATMAP_HTML,
            css_template=HEATMAP_CSS,
            js_on_load=HEATMAP_JS,
            **kwargs,
        )

    def api_info(self):
        return {"type": "object", "description": "Dict mapping YYYY-MM-DD to int counts"}


# ── Helpers ──────────────────────────────────────────────────────────────────

def _theme_update(theme):
    """Return gr.HTML update with only color props — preserves value/year."""
    c = COLOR_SCHEMES.get(theme, COLOR_SCHEMES["green"])
    return gr.HTML(c0=c[0], c1=c[1], c2=c[2], c3=c[3], c4=c[4])


def _full_update(data, year, theme):
    """Return gr.HTML update with value + year + colors all set."""
    c = COLOR_SCHEMES.get(theme, COLOR_SCHEMES["green"])
    return gr.HTML(value=data, year=int(year), c0=c[0], c1=c[1], c2=c[2], c3=c[3], c4=c[4])


# ── Pattern Generators ───────────────────────────────────────────────────────

def _dates(year):
    start = datetime(year, 1, 1)
    end = datetime(year, 12, 31)
    d = start
    while d <= end:
        yield d, d.strftime("%Y-%m-%d")
        d += timedelta(days=1)


def pattern_random(intensity, year):
    data = {}
    for d, key in _dates(year):
        if d > datetime.now(): break
        if random.random() < intensity:
            r = random.random()
            if r < 0.4:   data[key] = random.randint(1, 2)
            elif r < 0.7: data[key] = random.randint(3, 5)
            elif r < 0.9: data[key] = random.randint(6, 9)
            else:          data[key] = random.randint(10, 15)
    return data


def pattern_streak(year):
    data, in_streak = {}, False
    for d, key in _dates(year):
        if d > datetime.now(): break
        if random.random() < 0.08: in_streak = not in_streak
        if in_streak: data[key] = random.randint(2, 14)
    return data


def pattern_weekday(year):
    data = {}
    for d, key in _dates(year):
        if d > datetime.now(): break
        if d.weekday() < 5 and random.random() < 0.75:
            data[key] = random.randint(1, 10)
    return data


def pattern_weekend(year):
    data = {}
    for d, key in _dates(year):
        if d > datetime.now(): break
        if d.weekday() >= 5 and random.random() < 0.85:
            data[key] = random.randint(3, 15)
    return data


def pattern_momentum(year):
    data, total = {}, 365
    for i, (d, key) in enumerate(_dates(year)):
        if d > datetime.now(): break
        prob = 0.1 + 0.8 * (i / total)
        if random.random() < prob:
            data[key] = random.randint(1, max(1, int(1 + 14 * (i / total))))
    return data


def pattern_sine(year):
    data = {}
    for i, (d, key) in enumerate(_dates(year)):
        if d > datetime.now(): break
        w = (math.sin(2 * math.pi * i / 90) + 1) / 2
        if random.random() < 0.3 + 0.6 * w:
            data[key] = max(1, int(w * 14) + random.randint(0, 2))
    return data


def pattern_seasonal(year):
    data = {}
    for d, key in _dates(year):
        if d > datetime.now(): break
        if d.month in (1, 2, 6, 7, 8, 12):
            if random.random() < 0.8: data[key] = random.randint(3, 15)
        elif random.random() < 0.2:
            data[key] = random.randint(1, 4)
    return data


def pattern_burst(year):
    data, days = {}, list(_dates(year))
    for _ in range(random.randint(8, 15)):
        si = random.randint(0, len(days) - 20)
        for j in range(random.randint(3, 14)):
            if si + j < len(days):
                d, key = days[si + j]
                if d > datetime.now(): break
                data[key] = random.randint(5, 15)
    return data


PATTERNS = {
    "🎲 Random":           lambda i, y: pattern_random(i, y),
    "🔥 Streaks":          lambda i, y: pattern_streak(y),
    "💼 Weekday Warrior":  lambda i, y: pattern_weekday(y),
    "🌙 Weekend Hacker":   lambda i, y: pattern_weekend(y),
    "📈 Growing Momentum": lambda i, y: pattern_momentum(y),
    "🌊 Sine Wave":        lambda i, y: pattern_sine(y),
    "☀️ Seasonal":         lambda i, y: pattern_seasonal(y),
    "💥 Burst Mode":       lambda i, y: pattern_burst(y),
}


# ── Gradio App ───────────────────────────────────────────────────────────────

with gr.Blocks(
    title="🔥 Contribution Heatmap",
) as demo:

    gr.Markdown(
        """
    # 🟩 GitHub-Style Contribution Heatmap
    *Built entirely with Gradio 6's `gr.HTML` component — custom templates, CSS, and JS interactivity.*

    **Click any cell** to cycle its intensity (0 → 12 → 0). Use the controls below to generate patterns and switch themes.
    """
    )

    heatmap = ContributionHeatmap(
        value=pattern_random(0.6, 2025), year=2025, theme="green"
    )

    with gr.Row():
        with gr.Column(scale=1):
            theme_dd = gr.Dropdown(
                choices=list(COLOR_SCHEMES.keys()),
                value="green",
                label="🎨 Color Theme",
            )
        with gr.Column(scale=1):
            year_dd = gr.Dropdown(
                choices=[2023, 2024, 2025],
                value=2025,
                label="📅 Year",
            )
        with gr.Column(scale=1):
            pattern_dd = gr.Dropdown(
                choices=list(PATTERNS.keys()),
                value="🎲 Random",
                label="🧬 Pattern",
            )
        with gr.Column(scale=1):
            intensity = gr.Slider(
                0.1, 1.0, value=0.6, step=0.1, label="📊 Intensity"
            )

    with gr.Row():
        generate_btn = gr.Button("✨ Generate", variant="primary", size="lg")
        clear_btn = gr.Button("🗑️ Clear All", variant="stop")

    status = gr.Textbox(label="Status", interactive=False)

    # ── Events ────────────────────────────────────────────────────────────
    # FIX: return gr.HTML(prop=...) to update props on the EXISTING component
    # instead of returning a new ContributionHeatmap(...) instance.

    theme_dd.change(
        fn=lambda theme: _theme_update(theme),
        inputs=[theme_dd],
        outputs=heatmap,
    )

    year_dd.change(
        fn=lambda y, t, i, p: (
            _full_update(PATTERNS.get(p, PATTERNS["🎲 Random"])(i, int(y)), y, t),
            f"📅 Showing {y} — {len((d := PATTERNS.get(p, PATTERNS['🎲 Random'])(i, int(y))))} active days, {sum(d.values()):,} contributions",
        ),
        inputs=[year_dd, theme_dd, intensity, pattern_dd],
        outputs=[heatmap, status],
    )

    def on_generate(int_val, theme, year_val, pat):
        gen = PATTERNS.get(pat, PATTERNS["🎲 Random"])
        data = gen(int_val, int(year_val))
        total = sum(data.values())
        return (
            _full_update(data, year_val, theme),
            f"✨ {pat}  —  {len(data)} active days  ·  {total:,} contributions",
        )

    generate_btn.click(
        fn=on_generate,
        inputs=[intensity, theme_dd, year_dd, pattern_dd],
        outputs=[heatmap, status],
    )

    clear_btn.click(
        fn=lambda t, y: (_full_update({}, y, t), "🗑️ Cleared all contributions"),
        inputs=[theme_dd, year_dd],
        outputs=[heatmap, status],
    )

    heatmap.change(
        fn=lambda d: f"✏️ Edited: {len([v for v in (d or {}).values() if v > 0])} active days · {sum((d or {}).values()):,} contributions" if isinstance(d, dict) else "",
        inputs=heatmap,
        outputs=status,
    )


if __name__ == "__main__":
    demo.launch(
        theme=gr.themes.Soft(primary_hue="green"),
        css="footer { display: none !important; }",
    )