document.getElementById("analyze-button").addEventListener("click", async () => {
    const text = document.getElementById("input-text").value;
    const model = document.getElementById("model-select").value;
    
    // Show loading state
    const analyzeButton = document.getElementById("analyze-button");
    const buttonSpinner = analyzeButton.querySelector(".button-spinner");
    analyzeButton.classList.add("loading");
    buttonSpinner.classList.add("visible");
    analyzeButton.disabled = true;

    try {
        const response = await fetch("/analyze", {
            method: "POST",
            headers: {
                "Content-Type": "application/json"
            },
            body: JSON.stringify({ text, model })
        });

        const data = await response.json();

        const coloredTextDiv = document.getElementById("colored-text");
        coloredTextDiv.innerHTML = "";

        // Always add the first token
        const firstToken = data.tokens[0];
        const firstTokenSpan = document.createElement("span");
        firstTokenSpan.classList.add("token");
        
        // Handle special tokens and regular tokens differently
        if (firstToken === "<s>" || firstToken === "<|endoftext|>") {
            firstTokenSpan.style.backgroundColor = "#808080"; // Gray for special tokens
            firstTokenSpan.textContent = "■";
            tippy(firstTokenSpan, {
                content: "<div><strong>Beginning of Sequence</strong></div>",
                allowHTML: true,
                theme: 'custom',
                placement: 'top',
                interactive: true
            });
        } else {
            // Handle regular first token
            firstTokenSpan.style.backgroundColor = "#808080"; // or any other color you prefer
            firstTokenSpan.textContent = firstToken;
            tippy(firstTokenSpan, {
                content: `<div><strong>First Token</strong></div>`,
                allowHTML: true,
                theme: 'custom',
                placement: 'top',
                interactive: true
            });
        }
        
        coloredTextDiv.appendChild(firstTokenSpan);

        for (let index = 0; index < data.log_probs.length; index++) {
            const token = data.tokens[index + 1];
            const percentile = data.percentiles[index];
            const logProb = data.log_probs[index];
            const topKPredictions = data.top_k_predictions[index];
            const color = getColor(data.log_probs, logProb);

            const tokenSpan = document.createElement("span");
            tokenSpan.classList.add("token");
            tokenSpan.style.backgroundColor = color;

            let displayToken = token;
            let specialTokenDescription = "";

            // Enhanced special token handling
            if (token === "<s>" || token === "<|endoftext|>") {
                displayToken = "■";
                specialTokenDescription = "Beginning of Sequence";
            } else if (token === "</s>" || token === "<|endoftext|>") {
                displayToken = "■";
                specialTokenDescription = "End of Sequence";
            } else if (token === "<0x0A>") {
                displayToken = "■";
                specialTokenDescription = "Newline";
            } else if (token.startsWith("<") && token.endsWith(">")) {
                displayToken = "■";
                specialTokenDescription = "Special Token: " + token;
            } else {
                // Clean up GPT-2 style tokens (Ġ and Ċ)
                displayToken = displayToken
                    .replace(/\u2581/g, " ")  // Replace underscore token
                    .replace(/Ġ/g, " ")       // Replace GPT-2 space token
                    .replace(/Ċ/g, "\n");     // Replace GPT-2 newline token
            }

            tokenSpan.textContent = displayToken;

            let tooltipContent = "";
            if (specialTokenDescription) {
                tooltipContent += `<div style="font-weight: bold; margin-bottom: 8px;">${specialTokenDescription}</div>`;
            }

            tooltipContent += `<div style="font-weight: bold; margin-bottom: 4px;">Top 5 Predictions:</div>`;
            topKPredictions.forEach(pred => {
                let predToken = pred.token;
                if (predToken === "<0x0A>") {
                    predToken = "\\n";
                } else if (predToken.startsWith("<") && predToken.endsWith(">")) {
                    predToken = "[SPECIAL]";
                } else {
                    predToken = predToken
                        .replace(/\u2581/g, " ")
                        .replace(/Ġ/g, " ")
                        .replace(/Ċ/g, "\n");
                }
                tooltipContent += `<div style="padding-left: 8px;">${predToken}: ${pred.log_prob.toFixed(4)}</div>`;
            });
            
            tooltipContent += `<div style="margin-top: 8px; border-top: 1px solid #555; padding-top: 8px;">
                <div><strong>Stats:</strong></div>
                <div style="padding-left: 8px;">Percentile: ${percentile.toFixed(2)}</div>
                <div style="padding-left: 8px;">Log-Likelihood: ${logProb.toFixed(4)}</div>
            </div>`;

            tippy(tokenSpan, {
                content: tooltipContent,
                allowHTML: true,
                theme: 'custom',
                placement: 'top',
                interactive: true
            });

            coloredTextDiv.appendChild(tokenSpan);
            if (token === "<0x0A>") {
              coloredTextDiv.appendChild(document.createElement("br"));
            }
        }

        document.getElementById("joint-log-likelihood").textContent = data.joint_log_likelihood.toFixed(4);
        document.getElementById("average-log-likelihood").textContent = data.average_log_likelihood.toFixed(4);
    } catch (error) {
        console.error("Error during analysis:", error);
        alert("An error occurred during analysis. Please try again.");
    } finally {
        // Hide loading state
        analyzeButton.classList.remove("loading");
        buttonSpinner.classList.remove("visible");
        analyzeButton.disabled = false;
    }
});

function getColor(allLogProbs, currentLogProb) {
    const minLogProb = Math.min(...allLogProbs);
    const maxLogProb = Math.max(...allLogProbs);

    // Normalize to 0-1 range
    let normalizedLogProb = (currentLogProb - minLogProb) / (maxLogProb - minLogProb);
    normalizedLogProb = Math.max(0, Math.min(1, normalizedLogProb)); // Clamp

    // Optional: Apply a power transformation (adjust the exponent as needed)
    const power = 0.7; // Example: Less than 1 emphasizes differences at lower end
    normalizedLogProb = Math.pow(normalizedLogProb, power);

    const hue = normalizedLogProb * 120; // 0 (red) to 120 (green)
    return `hsl(${hue}, 100%, 50%)`;
}