model.py

"""
🎭 Story DNA Analyzer — model.py
=================================
Uses HuggingFace zero-shot NLI to dissect any story text into:
  • Genre         (8 categories)
  • Narrative Tropes (12 tropes)
  • Mood           (6 tones)
  • Target Audience (4 groups)
  • Villain Archetype (6 types)
  • Story Outcome Prediction (3 endings)

Model: cross-encoder/nli-MiniLM2-L6-H768  (~90 MB, fast CPU inference)
"""

import sys
from transformers import pipeline
from dataclasses import dataclass, field
from typing import Dict, List
import json

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")

# ──────────────────────────────────────────
# LABEL DEFINITIONS
# ──────────────────────────────────────────

GENRES = [
    "epic fantasy with magic and mythical creatures",
    "science fiction with technology and space",
    "psychological horror and dread",
    "slow-burn romance and emotional longing",
    "crime thriller and detective mystery",
    "historical drama set in the past",
    "dark comedy and satirical humor",
    "coming-of-age and personal growth",
]

TROPES = [
    "the chosen one who must save the world",
    "enemies who fall in love",
    "a betrayal by a trusted ally",
    "a hero's journey of self-discovery",
    "found family replacing blood family",
    "a dark secret hidden for years",
    "the mentor who dies to motivate the hero",
    "redemption arc for a former villain",
    "love triangle causing conflict",
    "an unreliable narrator hiding the truth",
    "a dystopian society being overthrown",
    "time travel or parallel universe twist",
]

MOODS = [
    "dark, gritty and hopeless",
    "warm, cozy and optimistic",
    "tense, suspenseful and unpredictable",
    "whimsical, playful and magical",
    "melancholic, bittersweet and nostalgic",
    "intense, passionate and dramatic",
]

AUDIENCES = [
    "young children under age 10",
    "teenagers and young adults",
    "mature adults dealing with complex themes",
    "general audiences of all ages",
]

VILLAIN_ARCHETYPES = [
    "a power-hungry tyrant who craves control",
    "a misunderstood anti-hero with tragic backstory",
    "a charismatic manipulator hiding in plain sight",
    "an ancient evil force or monster",
    "a corrupt institution or faceless system",
    "no clear villain — internal struggle is the enemy",
]

ENDINGS = [
    "a triumphant happy ending where good wins",
    "a tragic ending with sacrifice and loss",
    "an ambiguous open ending left to interpretation",
]

# ──────────────────────────────────────────
# PRETTY LABEL MAPPING
# ──────────────────────────────────────────

GENRE_DISPLAY = {
    "epic fantasy with magic and mythical creatures":        "🧙 Epic Fantasy",
    "science fiction with technology and space":             "🚀 Science Fiction",
    "psychological horror and dread":                        "👁️ Psychological Horror",
    "slow-burn romance and emotional longing":               "💔 Romance",
    "crime thriller and detective mystery":                  "🔍 Crime Thriller",
    "historical drama set in the past":                      "🏛️ Historical Drama",
    "dark comedy and satirical humor":                       "🎭 Dark Comedy",
    "coming-of-age and personal growth":                     "🌱 Coming-of-Age",
}

TROPE_DISPLAY = {
    "the chosen one who must save the world":               "⚡ The Chosen One",
    "enemies who fall in love":                             "💘 Enemies to Lovers",
    "a betrayal by a trusted ally":                        "🗡️ The Big Betrayal",
    "a hero's journey of self-discovery":                   "🗺️ Hero's Journey",
    "found family replacing blood family":                  "🤝 Found Family",
    "a dark secret hidden for years":                       "🔒 Hidden Dark Secret",
    "the mentor who dies to motivate the hero":             "💀 Dead Mentor",
    "redemption arc for a former villain":                  "🌅 Redemption Arc",
    "love triangle causing conflict":                       "💞 Love Triangle",
    "an unreliable narrator hiding the truth":              "🎭 Unreliable Narrator",
    "a dystopian society being overthrown":                 "✊ Revolution Arc",
    "time travel or parallel universe twist":               "🌀 Time/Dimension Twist",
}

MOOD_DISPLAY = {
    "dark, gritty and hopeless":                            "🖤 Dark & Gritty",
    "warm, cozy and optimistic":                            "☀️ Warm & Hopeful",
    "tense, suspenseful and unpredictable":                 "⚡ Tense & Suspenseful",
    "whimsical, playful and magical":                       "✨ Whimsical & Magical",
    "melancholic, bittersweet and nostalgic":               "🌧️ Melancholic",
    "intense, passionate and dramatic":                     "🔥 Passionate & Dramatic",
}

AUDIENCE_DISPLAY = {
    "young children under age 10":                         "🧒 Children (Under 10)",
    "teenagers and young adults":                          "🧑 Young Adult (13–25)",
    "mature adults dealing with complex themes":           "🧔 Adult (18+)",
    "general audiences of all ages":                       "👨‍👩‍👧‍👦 All Ages",
}

VILLAIN_DISPLAY = {
    "a power-hungry tyrant who craves control":            "👑 The Tyrant",
    "a misunderstood anti-hero with tragic backstory":     "🥀 The Tragic Anti-Hero",
    "a charismatic manipulator hiding in plain sight":     "🎭 The Hidden Manipulator",
    "an ancient evil force or monster":                    "🐉 Ancient Evil",
    "a corrupt institution or faceless system":            "🏛️ The System",
    "no clear villain — internal struggle is the enemy":   "🪞 Internal Conflict",
}

ENDING_DISPLAY = {
    "a triumphant happy ending where good wins":           "🏆 Happy Ending",
    "a tragic ending with sacrifice and loss":             "💧 Tragic Ending",
    "an ambiguous open ending left to interpretation":     "❓ Ambiguous Ending",
}


# ──────────────────────────────────────────
# RESULT DATACLASS
# ──────────────────────────────────────────

@dataclass
class StoryDNA:
    genre:           Dict[str, float] = field(default_factory=dict)
    top_tropes:      List[Dict]       = field(default_factory=list)
    mood:            Dict[str, float] = field(default_factory=dict)
    audience:        Dict[str, float] = field(default_factory=dict)
    villain:         Dict[str, float] = field(default_factory=dict)
    predicted_ending:Dict[str, float] = field(default_factory=dict)

    def top(self, d: dict, n=1):
        return sorted(d.items(), key=lambda x: -x[1])[:n]

    def summary(self) -> str:
        lines = ["=" * 52, "🎭  STORY DNA ANALYSIS", "=" * 52]

        lines.append(f"\n📚 GENRE")
        for label, score in self.top(self.genre, 3):
            bar = "█" * int(score * 20)
            lines.append(f"  {label:<30} {bar} {score:.0%}")

        lines.append(f"\n🧬 TOP NARRATIVE TROPES")
        for t in self.top_tropes[:4]:
            lines.append(f"  {t['label']:<30} {t['score']:.0%}")

        lines.append(f"\n🎨 MOOD")
        for label, score in self.top(self.mood, 2):
            lines.append(f"  {label:<30} {score:.0%}")

        lines.append(f"\n👥 TARGET AUDIENCE")
        for label, score in self.top(self.audience, 2):
            lines.append(f"  {label:<30} {score:.0%}")

        lines.append(f"\n😈 VILLAIN ARCHETYPE")
        for label, score in self.top(self.villain, 2):
            lines.append(f"  {label:<30} {score:.0%}")

        lines.append(f"\n🔮 PREDICTED ENDING")
        for label, score in self.top(self.predicted_ending, 2):
            lines.append(f"  {label:<30} {score:.0%}")

        lines.append("=" * 52)
        return "\n".join(lines)

    def to_dict(self) -> dict:
        return {
            "genre":            self.genre,
            "top_tropes":       self.top_tropes,
            "mood":             self.mood,
            "audience":         self.audience,
            "villain":          self.villain,
            "predicted_ending": self.predicted_ending,
        }


# ──────────────────────────────────────────
# ANALYZER CLASS
# ──────────────────────────────────────────

class StoryDNAAnalyzer:
    """
    Zero-shot story analyzer using HuggingFace NLI.
    No training needed — works out of the box.
    """

    MODEL_ID = "cross-encoder/nli-MiniLM2-L6-H768"

    def __init__(self, device: int = -1):
        print(f"⏳ Loading model: {self.MODEL_ID} ...")
        self.clf = pipeline(
            "zero-shot-classification",
            model=self.MODEL_ID,
            device=device,          # -1 = CPU, 0 = first GPU
        )
        print("✅ Model ready!\n")

    def _classify(self, text: str, labels: list, display_map: dict) -> dict:
        out = self.clf(text, candidate_labels=labels, multi_label=False)
        return {
            display_map[label]: round(score, 4)
            for label, score in zip(out["labels"], out["scores"])
        }

    def _classify_multi(self, text: str, labels: list, display_map: dict) -> list:
        out = self.clf(text, candidate_labels=labels, multi_label=True)
        results = [
            {"label": display_map[label], "score": round(score, 4)}
            for label, score in zip(out["labels"], out["scores"])
        ]
        return sorted(results, key=lambda x: -x["score"])

    def analyze(self, text: str) -> StoryDNA:
        if len(text.strip()) < 20:
            raise ValueError("Please provide at least 20 characters of story text.")

        print("🔬 Analyzing Story DNA...")

        dna = StoryDNA(
            genre            = self._classify(text, GENRES,            GENRE_DISPLAY),
            top_tropes       = self._classify_multi(text, TROPES,      TROPE_DISPLAY),
            mood             = self._classify(text, MOODS,             MOOD_DISPLAY),
            audience         = self._classify(text, AUDIENCES,         AUDIENCE_DISPLAY),
            villain          = self._classify(text, VILLAIN_ARCHETYPES,VILLAIN_DISPLAY),
            predicted_ending = self._classify(text, ENDINGS,           ENDING_DISPLAY),
        )
        return dna


# ──────────────────────────────────────────
# QUICK TEST
# ──────────────────────────────────────────

if __name__ == "__main__":
    analyzer = StoryDNAAnalyzer()

    sample = """
    In a crumbling empire ruled by an immortal king who feeds on fear,
    a young blind girl discovers she can see the true faces of the dead.
    Hunted by the king's shadow-wolves, she escapes into the forbidden forest
    where a disgraced general with blood on his hands offers her an uneasy alliance.
    Together they must unlock the secret buried beneath the throne — 
    a truth that will either free the kingdom or doom it forever.
    """

    result = analyzer.analyze(sample)
    print(result.summary())
    print("\n📦 JSON output:")
    print(json.dumps(result.to_dict(), indent=2))