ml_pipeline/content_classifier.py

"""
content_classifier.py
=====================
Heuristic-based ContentClassifier with ML support.
"""

from __future__ import annotations
import re
import os
import joblib
from dataclasses import dataclass, field
from typing import Any, Dict, List

# ── Pattern libraries ──────────────────────────────────────────────────────────
# English play signals
_ACT_RE      = re.compile(r"\bACT\s+([IVX]+|\d+)\b", re.IGNORECASE)
_SCENE_RE    = re.compile(r"\bSCENE\s+([IVX]+|\d+)\b", re.IGNORECASE)
_CUE_RE      = re.compile(r"^[A-Z][A-Z\s'\.]{1,28}[A-Z]\.?\s*$")
_STAGE_RE    = re.compile(r"^[\[(].*[\])]$")
_CAST_RE     = re.compile(r"(DRAMATIS|CAST OF|LIST OF)\s+(PERSONAE|CHARACTERS)", re.IGNORECASE)
_STAGE_VERB  = re.compile(r"\((?:Enter|Exit|Exeunt|Aside|to\s+\w+)[^)]*\)", re.IGNORECASE)

# French play signals
_ACTE_RE        = re.compile(r"\bACTE\s+([IVX]+|\d+)\b", re.IGNORECASE)
_SCENE_FR_RE    = re.compile(r"\bSC[ÈE]NE\s+([IVX]+|\d+)\b", re.IGNORECASE)
_CAST_FR_RE     = re.compile(r"(PERSONNAGES|LISTE DES PERSONNAGES|ACTEURS)", re.IGNORECASE)
_STAGE_FR_VERB  = re.compile(r"\((?:Il entre|Elle entre|Ils entrent|À part|Seul)[^)]*\)", re.IGNORECASE)

# English novel signals
_CHAPTER_RE  = re.compile(r"\b(CHAPTER|PROLOGUE|EPILOGUE|PART)\s+([\dIVX]+|[A-Z][a-z]+)?\b", re.IGNORECASE)
_FIRST_PERS  = re.compile(r"\bI\b(?:'[a-z]+)?\s")
_SAID_RE     = re.compile(r"\b(said|replied|whispered|shouted|murmured|asked)\b")
_PROSE_RE    = re.compile(r"[a-z][a-z\s,;]{60,}[.!?]")

# French novel signals
_CHAPITRE_RE = re.compile(r"\bCHAPITRE\s+([\dIVX]+|[A-Z][a-z]+)?\b", re.IGNORECASE)
_PARTIE_RE   = re.compile(r"\bPARTIE\s+([\dIVX]+|[A-Z][a-z]+)?\b", re.IGNORECASE)
_SAID_FR_RE  = re.compile(r"\b(dit|répondit|murmura|cria|demanda|chuchota)\b")

# Poem signals
_STANZA_GAP_RE  = re.compile(r"\n\s*\n")  # double newlines between stanzas
_RHYME_END_RE   = re.compile(r"[a-z]+[.!?,;]*\s*$", re.IGNORECASE)
_VERSE_LINE_RE  = re.compile(r"^[A-Z][^.!?]{5,60}[,;:]?\s*$")  # Short capitalised lines typical of poetry
_POEM_TITLE_RE  = re.compile(r"\b(POEM|SONNET|ODE|BALLAD|ELEGY|HAIKU|STANZA|VERSE|CANTO)\b", re.IGNORECASE)
_POEM_FR_RE     = re.compile(r"\b(POÈME|SONNET|ODE|BALLADE|ÉLÉGIE|STROPHE|VERS|CHANT)\b", re.IGNORECASE)

_WEIGHTS: Dict[str, Dict[str, float]] = {
    "play": {
        "act_heading":     8.0,
        "scene_heading":   8.0,
        "acte_heading":    8.0,   # French
        "scene_fr_heading":8.0,   # French
        "character_cue":   2.0,
        "stage_direction": 3.0,
        "cast_page":       10.0,
        "cast_fr_page":    10.0,  # French
        "stage_verb":      2.0,
        "stage_fr_verb":   2.0,   # French
    },
    "novel": {
        "chapter_heading": 6.0,
        "chapitre_heading":6.0,   # French
        "partie_heading":  4.0,   # French
        "first_person":    0.3,
        "said_tag":        0.5,
        "said_fr_tag":     0.5,   # French
        "prose_sentence":  0.1,
    },
    "poem": {
        "poem_title":      10.0,
        "poem_fr_title":   10.0,  # French
        "verse_line":      0.3,   # short capitalised lines
        "short_line_ratio": 8.0,  # high ratio of short lines
    },
}

@dataclass
class ClassificationResult:
    type: str              # "play" | "novel" | "poem" | "generic"
    confidence: float
    play_score: float
    novel_score: float
    poem_score: float = 0.0
    ml_probabilities: Dict[str, float] = field(default_factory=dict)
    signals: Dict[str, int] = field(default_factory=dict)

class ContentClassifier:
    MIN_EVIDENCE = 2.0
    MIN_RATIO = 1.2

    def __init__(self):
        self.model = None
        self.vectorizer = None
        self._load_model()

    def _load_model(self):
        try:
            model_dir = os.path.join(os.path.dirname(__file__), "../models")
            model_path = os.path.join(model_dir, "classifier_v1.joblib")
            vec_path = os.path.join(model_dir, "vectorizer_v1.joblib")
            if os.path.exists(model_path):
                self.model = joblib.load(model_path)
                self.vectorizer = joblib.load(vec_path)
        except: pass

    def classify(self, all_blocks: List[Dict[str, Any]]) -> ClassificationResult:
        play_score, novel_score, poem_score = 0.0, 0.0, 0.0
        signals = {k: 0 for sub in _WEIGHTS.values() for k in sub}
        lines = self._extract_lines(all_blocks)
        full_text = " ".join(lines)
        line_count = len(lines)

        for idx, text in enumerate(lines):
            t = text.strip()
            if not t: continue
            is_intro = (idx < line_count * 0.1)

            # ── English play signals ──
            if _ACT_RE.search(t):
                play_score += _WEIGHTS["play"]["act_heading"]
                signals["act_heading"] += 1
            if _SCENE_RE.search(t):
                play_score += _WEIGHTS["play"]["scene_heading"]
                signals["scene_heading"] += 1
            
            # Character cues and stage directions
            if _CUE_RE.match(t) and len(t.split()) <= 5:
                play_score += _WEIGHTS["play"]["character_cue"]
                signals["character_cue"] += 1
            if _STAGE_RE.search(t) or _STAGE_VERB.search(t):
                play_score += _WEIGHTS["play"]["stage_direction"]
                signals["stage_direction"] += 1
            if _CAST_RE.search(t):
                play_score += _WEIGHTS["play"]["cast_page"]
                signals["cast_page"] += 1

            # ── French play signals ──
            if _ACTE_RE.search(t):
                play_score += _WEIGHTS["play"]["acte_heading"]
                signals["acte_heading"] += 1
            if _SCENE_FR_RE.search(t):
                play_score += _WEIGHTS["play"]["scene_fr_heading"]
                signals["scene_fr_heading"] += 1
            if _CAST_FR_RE.search(t):
                play_score += _WEIGHTS["play"]["cast_fr_page"]
                signals["cast_fr_page"] += 1
            if _STAGE_FR_VERB.search(t):
                play_score += _WEIGHTS["play"]["stage_fr_verb"]
                signals["stage_fr_verb"] += 1

            # ── English novel signals ──
            if _CHAPTER_RE.search(t):
                novel_score += _WEIGHTS["novel"]["chapter_heading"]
                signals["chapter_heading"] += 1

            # ── French novel signals ──
            if _CHAPITRE_RE.search(t):
                novel_score += _WEIGHTS["novel"]["chapitre_heading"]
                signals["chapitre_heading"] += 1
            if _PARTIE_RE.search(t):
                novel_score += _WEIGHTS["novel"]["partie_heading"]
                signals["partie_heading"] += 1

            p_mult = 0.2 if is_intro else 1.0
            fp = len(_FIRST_PERS.findall(t))
            if fp:
                novel_score += fp * _WEIGHTS["novel"]["first_person"] * p_mult
                signals["first_person"] += fp
            if _SAID_RE.search(t):
                novel_score += _WEIGHTS["novel"]["said_tag"] * p_mult
                signals["said_tag"] += 1
            if _SAID_FR_RE.search(t):
                novel_score += _WEIGHTS["novel"]["said_fr_tag"] * p_mult
                signals["said_fr_tag"] += 1
            if _PROSE_RE.search(t):
                novel_score += _WEIGHTS["novel"]["prose_sentence"] * p_mult
                signals["prose_sentence"] += 1

            # ── Poem signals ──
            if _POEM_TITLE_RE.search(t):
                poem_score += _WEIGHTS["poem"]["poem_title"]
                signals["poem_title"] += 1
            if _POEM_FR_RE.search(t):
                poem_score += _WEIGHTS["poem"]["poem_fr_title"]
                signals["poem_fr_title"] += 1
            if _VERSE_LINE_RE.match(t) and len(t.split()) <= 12:
                poem_score += _WEIGHTS["poem"]["verse_line"]
                signals["verse_line"] += 1

        # Poem: high ratio of short lines (< 60 chars) is a strong signal
        if line_count > 5:
            short_lines = sum(1 for l in lines if len(l.strip()) < 60 and len(l.strip()) > 3)
            short_ratio = short_lines / line_count
            if short_ratio > 0.7:
                poem_score += _WEIGHTS["poem"]["short_line_ratio"]
                signals["short_line_ratio"] = int(short_ratio * 100)

        # Check heuristic evidence BEFORE ML boost — sparse docs stay generic
        heuristic_total = play_score + novel_score + poem_score

        ml_probs = {"play": 0.0, "novel": 0.0, "poem": 0.0, "generic": 0.0}
        if self.model and self.vectorizer and full_text:
            try:
                vec = self.vectorizer.transform([full_text])
                probs = self.model.predict_proba(vec)[0]
                ml_probs = dict(zip(self.model.classes_, probs))
                play_score  += ml_probs.get("play",  0.0) * 10.0
                novel_score += ml_probs.get("novel", 0.0) * 10.0
            except: pass

        # STRICT OVERRIDE FOR PLAYS (English OR French)
        en_play = (signals["act_heading"] >= 1 and signals["scene_heading"] >= 1) or signals["act_heading"] >= 3
        fr_play = (signals["acte_heading"] >= 1 and signals["scene_fr_heading"] >= 1) or signals["acte_heading"] >= 3
        if en_play or fr_play:
            return ClassificationResult("play", 0.95, play_score, novel_score, poem_score, ml_probs, signals)

        # STRICT OVERRIDE FOR POEMS
        if signals.get("poem_title", 0) >= 1 or (poem_score > play_score and poem_score > novel_score and poem_score >= self.MIN_EVIDENCE):
            total = play_score + novel_score + poem_score
            conf = round(poem_score / max(total, 1), 3) if total > 0 else 0.8
            return ClassificationResult("poem", max(conf, 0.7), play_score, novel_score, poem_score, ml_probs, signals)

        total = play_score + novel_score + poem_score
        if heuristic_total < self.MIN_EVIDENCE:
            return ClassificationResult("generic", 0.5, play_score, novel_score, poem_score, ml_probs, signals)

        scores = {"play": play_score, "novel": novel_score, "poem": poem_score}
        doc_type = max(scores, key=scores.get)
        conf = round(scores[doc_type] / total, 3)
        return ClassificationResult(doc_type, conf, play_score, novel_score, poem_score, ml_probs, signals)

    def _extract_lines(self, blocks: List[Dict[str, Any]]) -> List[str]:
        lines = []
        for b in blocks:
            if b.get("type") != 0: continue
            for l in b.get("lines", []):
                txt = "".join([s.get("text", "") for s in l.get("spans", [])]).strip()
                if txt: lines.append(txt)
        return lines