app.py

import gradio as gr
import pandas as pd
import numpy as np
import re
from scipy.spatial.distance import pdist, squareform
from scipy.cluster.hierarchy import linkage
import plotly.figure_factory as ff
import plotly.graph_objects as go
import tempfile
import os

def extract_features(text):
    # Basic tokens
    tokens = re.findall(r'\b[a-zA-Z]+\b', text.lower())
    total_words = len(tokens)
    
    if total_words < 10:
        return None
        
    unique_words = len(set(tokens))
    
    # 1. Type-Token Ratio (Vocabulary Richness)
    ttr = unique_words / total_words
    
    # 2. Average Word Length
    avg_word_len = np.mean([len(w) for w in tokens])
    
    # 3. Sentence Length Metrics
    sentences = re.split(r'[.!?]+', text)
    sentences = [s.strip() for s in sentences if s.strip()]
    sent_lengths = [len(re.findall(r'\b[a-zA-Z]+\b', s)) for s in sentences]
    avg_sent_len = np.mean(sent_lengths) if sent_lengths else 0
    std_sent_len = np.std(sent_lengths) if sent_lengths else 0
    
    # 4. Punctuation frequencies
    commas = len(re.findall(r',', text)) / total_words
    semicolons = len(re.findall(r';', text)) / total_words
    exclamations = len(re.findall(r'!', text)) / total_words
    questions = len(re.findall(r'\?', text)) / total_words
    
    return {
        "Vocabulary Richness (TTR)": ttr,
        "Average Word Length": avg_word_len,
        "Average Sentence Length": avg_sent_len,
        "Sentence Length Variation (STD)": std_sent_len,
        "Commas Frequency": commas,
        "Semicolons Frequency": semicolons,
        "Exclamation Frequency": exclamations,
        "Question Frequency": questions
    }

def run_stylometry(files, chosen_features):
    if not files or len(files) < 3:
        return "Please upload at least 3 distinct text files (TXT format) to perform comparative stylometry.", None, None, None
        
    documents = {}
    doc_features = {}
    
    for file_obj in files:
        # Extract filename as label
        label = os.path.splitext(os.path.basename(file_obj.name))[0]
        try:
            with open(file_obj.name, "r", encoding="utf-8", errors="ignore") as f:
                text = f.read()
            feats = extract_features(text)
            if feats:
                documents[label] = text
                doc_features[label] = feats
        except Exception as e:
            return f"Error reading file '{label}': {str(e)}", None, None, None
            
    if len(doc_features) < 3:
        return "At least 3 files must have valid text content of at least 10 words.", None, None, None
        
    df_features = pd.DataFrame(doc_features).transpose()
    
    # Subset by chosen features
    if not chosen_features:
        chosen_features = list(df_features.columns)
        
    df_sub = df_features[chosen_features]
    
    # Standardize features (Z-Score)
    df_norm = (df_sub - df_sub.mean()) / df_sub.std()
    # Replace possible NaNs from zero standard deviation
    df_norm = df_norm.fillna(0)
    
    labels = list(df_norm.index)
    
    # Calculate Distance Matrix (Euclidean)
    distances = pdist(df_norm.values, metric='euclidean')
    dist_matrix = squareform(distances)
    
    # Max similarity corresponds to 0 distance
    # Convert distance to a similarity score between 0 and 100%
    max_d = np.max(dist_matrix) if np.max(dist_matrix) > 0 else 1.0
    sim_matrix = (1.0 - (dist_matrix / max_d)) * 100
    
    # 1. Similarity Heatmap Plotly
    fig_heatmap = go.Figure(data=go.Heatmap(
        z=sim_matrix,
        x=labels,
        y=labels,
        colorscale='Hot',
        text=[[f"{val:.1f}%" for val in row] for row in sim_matrix],
        texttemplate="%{text}",
        hoverinfo='z'
    ))
    
    fig_heatmap.update_layout(
        title="Stylometric Style Similarity Heatmap (%)",
        paper_bgcolor='#16100c',
        plot_bgcolor='#16100c',
        font_color='#f4eee6',
        xaxis=dict(gridcolor='rgba(255,255,255,0.05)'),
        yaxis=dict(gridcolor='rgba(255,255,255,0.05)'),
        margin=dict(l=40, r=40, t=50, b=40)
    )
    
    # 2. Hierarchical Cluster Dendrogram
    try:
        Z = linkage(distances, 'ward')
        fig_dendro = ff.create_dendrogram(df_norm.values, orientation='left', labels=labels, linkagefun=lambda x: Z)
        fig_dendro.update_layout(
            title="Hierarchical Stylistic Cluster Dendrogram",
            paper_bgcolor='#16100c',
            plot_bgcolor='#16100c',
            font_color='#f4eee6',
            xaxis=dict(showgrid=True, gridcolor='rgba(255,255,255,0.05)', title="Distance (Ward threshold)"),
            yaxis=dict(gridcolor='rgba(255,255,255,0.05)'),
            margin=dict(l=80, r=40, t=50, b=40)
        )
        # Match colors to style system
        for trace in fig_dendro.data:
            if 'color' in trace:
                trace.line.color = '#ff7043'
    except Exception as e:
        # Fallback to simple placeholder scatter if dendrogram linkage fails (e.g. from numerical identical elements)
        fig_dendro = go.Figure()
        fig_dendro.add_annotation(text=f"Hierarchical linkage omitted: {str(e)}", showarrow=False, font=dict(size=14))
        fig_dendro.update_layout(paper_bgcolor='#16100c', font_color='#f4eee6')

    # Prep output tables
    df_features_out = df_features.round(4).reset_index().rename(columns={"index": "Document Label"})
    
    # Save CSV
    out_csv = tempfile.mktemp(suffix=".csv")
    df_features_out.to_csv(out_csv, index=False)
    
    return "", fig_heatmap, fig_dendro, df_features_out, gr.update(value=out_csv, visible=True)

theme = gr.themes.Default(
    primary_hue="orange",
    neutral_hue="stone"
).set(
    body_background_fill="#0d0907",
    body_text_color="#c4bbae",
    block_background_fill="#16100c",
    block_border_width="1px",
    block_label_text_color="#f4eee6"
)

all_features = [
    "Vocabulary Richness (TTR)",
    "Average Word Length",
    "Average Sentence Length",
    "Sentence Length Variation (STD)",
    "Commas Frequency",
    "Semicolons Frequency",
    "Exclamation Frequency",
    "Question Frequency"
]

with gr.Blocks(theme=theme, title="Comparative Stylometry Analyzer") as demo:
    gr.Markdown(
        """
        # 🖋️ Comparative Stylometry & Authorship Analyzer
        ### Extract, analyze, and map the quantitative writing styles of multiple documents. Perfect for authorship debates, forensic linguistics, and distant reading comparisons.
        """
    )
    
    error_msg = gr.Markdown("", visible=False)
    
    with gr.Row():
        with gr.Column(scale=1):
            file_objs = gr.File(
                label="Upload Text Files (Select at least 3 TXT files)", 
                file_types=[".txt"], 
                file_count="multiple"
            )
            
            chosen_features = gr.CheckboxGroup(
                choices=all_features,
                value=all_features[:5],
                label="Stylometric Features to Compare",
                info="Toggle features to customize the mathematical fingerprint vector."
            )
            
            btn = gr.Button("Calculate Writing Footprints", variant="primary")
            
        with gr.Column(scale=2):
            with gr.Tabs():
                with gr.TabItem("Style Similarity Matrix"):
                    plot_heatmap = gr.Plot()
                with gr.TabItem("Clustering Dendrogram"):
                    plot_dendro = gr.Plot()
                with gr.TabItem("Extracted Stylometrics Table"):
                    table_features = gr.Dataframe()
                    download_btn = gr.File(label="Download Full Stylometrics CSV", visible=False)

    def process(files, features):
        err, heatmap, dendro, table, csv_path = run_stylometry(files, features)
        if err:
            return gr.update(value=err, visible=True), None, None, None, gr.update(visible=False)
        return gr.update(visible=False), heatmap, dendro, table, csv_path

    btn.click(
        process,
        inputs=[file_objs, chosen_features],
        outputs=[error_msg, plot_heatmap, plot_dendro, table_features, download_btn]
    )

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