semantic.py

import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import json
import numpy as np
from enum import Enum
from dataclasses import dataclass
from typing import List, Dict, Tuple, Any, Optional

from sentence_transformers import SentenceTransformer

EMBEDDING_MODEL_NAME = "all-mpnet-base-v2"


class SimilarityDimension(Enum):
    CHARACTER_SEMANTIC = "character_semantic"
    THEMATIC = "thematic"
    EMOTIONAL_ARC = "emotional_arc"
    CAUSAL_STRUCTURE = "causal_structure"
    NARRATIVE_FUNCTION = "narrative_function"
    RELATIONSHIP_DYNAMICS = "relationship_dynamics"
    EVENT_SEMANTIC = "event_semantic"


@dataclass
class SimilarityScore:
    dimension: SimilarityDimension
    score: float
    details: Dict[str, Any]  # should include per-component raw scores for training


# Fixed public order of top-level dimensions (7)
DIM_ORDER = [
    "thematic",
    "character_semantic",
    "emotional_arc",
    "causal_structure",
    "narrative_function",
    "relationship_dynamics",
    "event_semantic",
]

# Map string name -> Enum
_DIM_TO_ENUM = {
    "thematic": SimilarityDimension.THEMATIC,
    "character_semantic": SimilarityDimension.CHARACTER_SEMANTIC,
    "emotional_arc": SimilarityDimension.EMOTIONAL_ARC,
    "causal_structure": SimilarityDimension.CAUSAL_STRUCTURE,
    "narrative_function": SimilarityDimension.NARRATIVE_FUNCTION,
    "relationship_dynamics": SimilarityDimension.RELATIONSHIP_DYNAMICS,
    "event_semantic": SimilarityDimension.EVENT_SEMANTIC,
}

# Component order (18) for training intra-dimension weights
# (dim, component_key) pairs — keys must match those put in SimilarityScore.details
COMPONENT_ORDER: List[Tuple[str, str]] = [
    ("character_semantic", "role"),
    ("character_semantic", "attribute"),
    ("character_semantic", "agency"),

    ("thematic", "theme"),
    ("thematic", "abstraction"),

    ("emotional_arc", "arc"),
    ("emotional_arc", "distribution"),
    ("emotional_arc", "change"),

    ("causal_structure", "relation"),
    ("causal_structure", "strength"),
    ("causal_structure", "gap"),

    ("narrative_function", "structure"),
    ("narrative_function", "beats"),

    ("relationship_dynamics", "type"),
    ("relationship_dynamics", "evolution"),
    ("relationship_dynamics", "power"),

    ("event_semantic", "type"),
    ("event_semantic", "tone"),
]


def _normalize_weighted_sum(components: Dict[str, float], weights: Dict[str, float], default_weights: Dict[str, float]) -> float:
    """
    Combine component scores with provided weights (or defaults), normalized by sum of weights.
    """
    if weights is None:
        weights = default_weights
    # ensure all component keys exist in weights (missing -> 0)
    keys = list(components.keys())
    w = np.array([max(0.0, float(weights.get(k, 0.0))) for k in keys], dtype=float)
    x = np.array([float(components[k]) for k in keys], dtype=float)
    sw = float(w.sum())
    if sw <= 0:
        # if everything is zero, fall back to equal weights
        w = np.ones_like(w, dtype=float)
        sw = float(w.sum())
    return float((w * x).sum() / sw)


class StorySemanticAnalyzer:
    """
    Analyzes semantic similarity between stories across multiple narrative dimensions.
    Supports optional per-dimension 'intra' weights to combine sub-components.
    """

    def __init__(self, embedding_model=None, model_name=EMBEDDING_MODEL_NAME):
        if embedding_model is not None:
            self.embedding_model = embedding_model
        else:
            self.embedding_model = SentenceTransformer(model_name)
            print(f"Loaded embedding model: {model_name}")

    # ---------- Public API ----------

    def compute_story_similarity(
        self,
        story1,
        story2,
        intra: Optional[Dict[str, Dict[str, float]]] = None,
    ) -> Dict[SimilarityDimension, SimilarityScore]:
        """
        Compute semantic similarity; if 'intra' is provided, it can override
        the per-dimension component mixture (e.g., structure vs beats).
        """
        sims: Dict[SimilarityDimension, SimilarityScore] = {}

        sims[SimilarityDimension.CHARACTER_SEMANTIC] = self._character_semantic_similarity(
            story1, story2, intra.get("character_semantic") if intra else None
        )
        sims[SimilarityDimension.THEMATIC] = self._thematic_similarity(
            story1, story2, intra.get("thematic") if intra else None
        )
        sims[SimilarityDimension.EMOTIONAL_ARC] = self._emotional_arc_similarity(
            story1, story2, intra.get("emotional_arc") if intra else None
        )
        sims[SimilarityDimension.CAUSAL_STRUCTURE] = self._causal_structure_similarity(
            story1, story2, intra.get("causal_structure") if intra else None
        )
        sims[SimilarityDimension.NARRATIVE_FUNCTION] = self._narrative_function_similarity(
            story1, story2, intra.get("narrative_function") if intra else None
        )
        sims[SimilarityDimension.RELATIONSHIP_DYNAMICS] = self._relationship_dynamics_similarity(
            story1, story2, intra.get("relationship_dynamics") if intra else None
        )
        sims[SimilarityDimension.EVENT_SEMANTIC] = self._event_semantic_similarity(
            story1, story2, intra.get("event_semantic") if intra else None
        )

        return sims

    # ---------- Dimension scorers (each returns component details + combined score) ----------

    def _character_semantic_similarity(self, story1, story2, intra_dim: Optional[Dict[str, float]]) -> SimilarityScore:
        if not story1.characters or not story2.characters:
            return SimilarityScore(SimilarityDimension.CHARACTER_SEMANTIC, 0.0, {})

        roles1 = self._extract_character_roles(story1.characters)
        roles2 = self._extract_character_roles(story2.characters)
        role_similarity = self._compute_role_overlap(roles1, roles2)

        attrs1 = self._extract_character_attributes(story1.characters)
        attrs2 = self._extract_character_attributes(story2.characters)
        attr_similarity = self._compute_attribute_similarity(attrs1, attrs2)

        agency1 = self._extract_agency_distribution(story1.characters)
        agency2 = self._extract_agency_distribution(story2.characters)
        agency_similarity = self._compute_agency_similarity(agency1, agency2)

        components = {
            "role": role_similarity,
            "attribute": attr_similarity,
            "agency": agency_similarity,
        }
        default_w = {"role": 0.4, "attribute": 0.4, "agency": 0.2}
        combined = _normalize_weighted_sum(components, intra_dim, default_w)

        details = {
            **components,
            "roles_1": roles1,
            "roles_2": roles2,
            "common_attributes": self._find_common_attributes(attrs1, attrs2),
        }
        return SimilarityScore(SimilarityDimension.CHARACTER_SEMANTIC, combined, details)

    def _thematic_similarity(self, story1, story2, intra_dim: Optional[Dict[str, float]]) -> SimilarityScore:
        if not story1.thematic_elements or not story2.thematic_elements:
            return SimilarityScore(SimilarityDimension.THEMATIC, 0.0, {})

        themes1 = [elem.get('theme', '') for elem in story1.thematic_elements]
        themes2 = [elem.get('theme', '') for elem in story2.thematic_elements]
        theme_similarity = self._compute_text_similarity(themes1, themes2)

        abstractions1 = [elem.get('abstraction_level', '') for elem in story1.thematic_elements]
        abstractions2 = [elem.get('abstraction_level', '') for elem in story2.thematic_elements]
        abstraction_similarity = self._compute_categorical_similarity(abstractions1, abstractions2)

        components = {
            "theme": theme_similarity,
            "abstraction": abstraction_similarity,
        }
        default_w = {"theme": 0.8, "abstraction": 0.2}
        combined = _normalize_weighted_sum(components, intra_dim, default_w)

        details = {
            **components,
            "themes_1": themes1,
            "themes_2": themes2,
            "theme_overlap": self._find_theme_overlap(themes1, themes2),
            "abstraction_similarity": abstraction_similarity,
        }
        return SimilarityScore(SimilarityDimension.THEMATIC, combined, details)

    def _emotional_arc_similarity(self, story1, story2, intra_dim: Optional[Dict[str, float]]) -> SimilarityScore:
        if not story1.emotional_trajectories or not story2.emotional_trajectories:
            return SimilarityScore(SimilarityDimension.EMOTIONAL_ARC, 0.0, {})

        arc_similarity = self._compare_emotional_progressions(story1.emotional_trajectories, story2.emotional_trajectories)
        emotions1 = self._extract_emotion_distribution(story1.emotional_trajectories)
        emotions2 = self._extract_emotion_distribution(story2.emotional_trajectories)
        emotion_similarity = self._compute_distribution_similarity(emotions1, emotions2)
        changes1 = self._extract_change_patterns(story1.emotional_trajectories)
        changes2 = self._extract_change_patterns(story2.emotional_trajectories)
        change_similarity = self._compute_categorical_similarity(changes1, changes2)

        components = {
            "arc": arc_similarity,
            "distribution": emotion_similarity,
            "change": change_similarity,
        }
        default_w = {"arc": 0.5, "distribution": 0.3, "change": 0.2}
        combined = _normalize_weighted_sum(components, intra_dim, default_w)

        details = {
            **components,
            "emotion_distribution_1": emotions1,
            "emotion_distribution_2": emotions2,
            "common_emotional_changes": self._find_common_changes(changes1, changes2),
        }
        return SimilarityScore(SimilarityDimension.EMOTIONAL_ARC, combined, details)

    def _causal_structure_similarity(self, story1, story2, intra_dim: Optional[Dict[str, float]]) -> SimilarityScore:
        if not story1.causal_chains or not story2.causal_chains:
            return SimilarityScore(SimilarityDimension.CAUSAL_STRUCTURE, 0.0, {})

        relations1 = [chain.get('relationship_type', '') for chain in story1.causal_chains]
        relations2 = [chain.get('relationship_type', '') for chain in story2.causal_chains]
        relation_similarity = self._compute_categorical_similarity(relations1, relations2)

        strengths1 = [chain.get('strength', '') for chain in story1.causal_chains]
        strengths2 = [chain.get('strength', '') for chain in story2.causal_chains]
        strength_similarity = self._compute_categorical_similarity(strengths1, strengths2)

        gaps1 = [chain.get('temporal_gap', '') for chain in story1.causal_chains]
        gaps2 = [chain.get('temporal_gap', '') for chain in story2.causal_chains]
        gap_similarity = self._compute_categorical_similarity(gaps1, gaps2)

        components = {
            "relation": relation_similarity,
            "strength": strength_similarity,
            "gap": gap_similarity,
        }
        default_w = {"relation": 0.5, "strength": 0.3, "gap": 0.2}
        combined = _normalize_weighted_sum(components, intra_dim, default_w)

        details = {
            **components,
            "causal_types_1": relations1,
            "causal_types_2": relations2,
            "common_causal_patterns": self._find_common_elements(relations1, relations2),
            "strength_similarity": strength_similarity,
        }
        return SimilarityScore(SimilarityDimension.CAUSAL_STRUCTURE, combined, details)

    def _narrative_function_similarity(self, story1, story2, intra_dim: Optional[Dict[str, float]]) -> SimilarityScore:
        if not story1.events or not story2.events:
            return SimilarityScore(SimilarityDimension.NARRATIVE_FUNCTION, 0.0, {})

        functions1 = [event.get('narrative_function', '') for event in story1.events]
        functions2 = [event.get('narrative_function', '') for event in story2.events]

        structure_similarity = self._compute_narrative_structure_similarity(functions1, functions2)
        beats1 = set(functions1)
        beats2 = set(functions2)
        beat_overlap = len(beats1.intersection(beats2)) / len(beats1.union(beats2)) if beats1.union(beats2) else 0.0

        components = {
            "structure": structure_similarity,
            "beats": beat_overlap,
        }
        default_w = {"structure": 0.6, "beats": 0.4}
        combined = _normalize_weighted_sum(components, intra_dim, default_w)

        details = {
            **components,
            "narrative_functions_1": functions1,
            "narrative_functions_2": functions2,
            "common_beats": list(beats1.intersection(beats2)),
            "unique_to_story1": list(beats1 - beats2),
            "unique_to_story2": list(beats2 - beats1),
        }
        return SimilarityScore(SimilarityDimension.NARRATIVE_FUNCTION, combined, details)

    def _relationship_dynamics_similarity(self, story1, story2, intra_dim: Optional[Dict[str, float]]) -> SimilarityScore:
        if not story1.relationships or not story2.relationships:
            return SimilarityScore(SimilarityDimension.RELATIONSHIP_DYNAMICS, 0.0, {})

        types1 = [rel.get('relationship_type', '') for rel in story1.relationships]
        types2 = [rel.get('relationship_type', '') for rel in story2.relationships]
        type_similarity = self._compute_categorical_similarity(types1, types2)

        evolutions1 = [rel.get('relationship_evolution', '') for rel in story1.relationships]
        evolutions2 = [rel.get('relationship_evolution', '') for rel in story2.relationships]
        evolution_similarity = self._compute_categorical_similarity(evolutions1, evolutions2)

        powers1 = [rel.get('power_dynamic', '') for rel in story1.relationships]
        powers2 = [rel.get('power_dynamic', '') for rel in story2.relationships]
        power_similarity = self._compute_categorical_similarity(powers1, powers2)

        components = {
            "type": type_similarity,
            "evolution": evolution_similarity,
            "power": power_similarity,
        }
        default_w = {"type": 0.4, "evolution": 0.3, "power": 0.3}
        combined = _normalize_weighted_sum(components, intra_dim, default_w)

        details = {
            **components,
            "relationship_types_overlap": self._find_common_elements(types1, types2),
            "evolution_patterns_overlap": self._find_common_elements(evolutions1, evolutions2),
            "power_dynamics_overlap": self._find_common_elements(powers1, powers2),
        }
        return SimilarityScore(SimilarityDimension.RELATIONSHIP_DYNAMICS, combined, details)

    def _event_semantic_similarity(self, story1, story2, intra_dim: Optional[Dict[str, float]]) -> SimilarityScore:
        if not story1.events or not story2.events:
            return SimilarityScore(SimilarityDimension.EVENT_SEMANTIC, 0.0, {})

        types1 = [event.get('event_type', '') for event in story1.events]
        types2 = [event.get('event_type', '') for event in story2.events]
        type_similarity = self._compute_categorical_similarity(types1, types2)

        tones1 = [event.get('emotional_tone', '') for event in story1.events]
        tones2 = [event.get('emotional_tone', '') for event in story2.events]
        tone_similarity = self._compute_categorical_similarity(tones1, tones2)

        components = {
            "type": type_similarity,
            "tone": tone_similarity,
        }
        # Old code used 0.3 + 0.2 (unnormalized 0.5 total). We normalize to be consistent.
        default_w = {"type": 0.3, "tone": 0.2}
        combined = _normalize_weighted_sum(components, intra_dim, default_w)

        details = {
            **components,
            "event_types_overlap": self._find_common_elements(types1, types2),
            "emotional_tones_overlap": self._find_common_elements(tones1, tones2),
        }
        return SimilarityScore(SimilarityDimension.EVENT_SEMANTIC, combined, details)

    # -------------------------------
    # COMPUTATION HELPERS
    # -------------------------------

    def _extract_character_roles(self, characters):
        roles = []
        for char in characters:
            roles.extend(char.get('functional_roles', []))
        return roles

    def _extract_character_attributes(self, characters):
        attributes = []
        for char in characters:
            attributes.extend(char.get('key_attributes', []))
        return attributes

    def _extract_agency_distribution(self, characters):
        agencies = [char.get('agency_level', '') for char in characters]
        return {level: agencies.count(level) for level in set(agencies)}

    def _compute_role_overlap(self, roles1, roles2):
        set1, set2 = set(roles1), set(roles2)
        return len(set1.intersection(set2)) / len(set1.union(set2)) if set1.union(set2) else 0.0

    def _compute_attribute_similarity(self, attrs1, attrs2):
        set1, set2 = set(attrs1), set(attrs2)
        return len(set1.intersection(set2)) / len(set1.union(set2)) if set1.union(set2) else 0.0

    def _compute_agency_similarity(self, agency1, agency2):
        all_levels = set(agency1.keys()).union(set(agency2.keys()))
        if not all_levels:
            return 0.0
        total1 = sum(agency1.values())
        total2 = sum(agency2.values())
        sim = 0.0
        for level in all_levels:
            p1 = agency1.get(level, 0) / total1 if total1 > 0 else 0.0
            p2 = agency2.get(level, 0) / total2 if total2 > 0 else 0.0
            sim += 1.0 - abs(p1 - p2)
        return sim / len(all_levels)

    def _compute_text_similarity(self, texts1, texts2):
        if not texts1 or not texts2:
            return 0.0
        e1 = self.embedding_model.encode(texts1)
        e2 = self.embedding_model.encode(texts2)
        max_sims = []
        for v1 in e1:
            sims = [self._cosine_similarity(v1, v2) for v2 in e2]
            max_sims.append(max(sims) if sims else 0.0)
        return float(np.mean(max_sims)) if max_sims else 0.0

    def _compute_categorical_similarity(self, cat1, cat2):
        set1, set2 = set(cat1), set(cat2)
        return len(set1.intersection(set2)) / len(set1.union(set2)) if set1.union(set2) else 0.0

    def _compute_distribution_similarity(self, dist1, dist2):
        all_keys = set(dist1.keys()).union(set(dist2.keys()))
        if not all_keys:
            return 0.0
        total1 = sum(dist1.values())
        total2 = sum(dist2.values())
        sim = 0.0
        for k in all_keys:
            p1 = dist1.get(k, 0) / total1 if total1 > 0 else 0.0
            p2 = dist2.get(k, 0) / total2 if total2 > 0 else 0.0
            sim += 1.0 - abs(p1 - p2)
        return sim / len(all_keys)

    def _find_common_attributes(self, attrs1, attrs2):
        return list(set(attrs1).intersection(set(attrs2)))

    def _find_theme_overlap(self, themes1, themes2):
        if not themes1 or not themes2:
            return []
        overlaps = []
        e1 = self.embedding_model.encode(themes1)
        e2 = self.embedding_model.encode(themes2)
        for theme1, v1 in zip(themes1, e1):
            for theme2, v2 in zip(themes2, e2):
                sim = self._cosine_similarity(v1, v2)
                if sim > 0.7:
                    overlaps.append((theme1, theme2, float(sim)))
        return overlaps

    def _find_common_elements(self, list1, list2):
        return list(set(list1).intersection(set(list2)))

    def _extract_emotion_distribution(self, trajectories):
        emotions = []
        for traj in trajectories:
            for point in traj.get('trajectory', []):
                emotions.append(point.get('emotional_state', ''))
        return {e: emotions.count(e) for e in set(emotions)}

    def _extract_change_patterns(self, trajectories):
        changes = []
        for traj in trajectories:
            for point in traj.get('trajectory', []):
                changes.append(point.get('emotional_change', ''))
        return changes

    def _find_common_changes(self, changes1, changes2):
        return list(set(changes1).intersection(set(changes2)))

    def _compare_emotional_progressions(self, traj1, traj2):
        """Sequence similarity (normalized LCS) between emotional progressions."""
        if not traj1 or not traj2:
            return 0.0
        seqs1 = self._extract_emotional_sequences(traj1)
        seqs2 = self._extract_emotional_sequences(traj2)
        if not seqs1 or not seqs2:
            return 0.0
        sims = []
        for s1 in seqs1:
            best = 0.0
            for s2 in seqs2:
                lcs = self._longest_common_subsequence(s1, s2)
                denom = max(len(s1), len(s2))
                best = max(best, (lcs / denom) if denom > 0 else 0.0)
            sims.append(best)
        return float(np.mean(sims)) if sims else 0.0

    def _extract_emotional_sequences(self, trajectories):
        seqs: List[List[str]] = []
        for traj in trajectories:
            pts = traj.get('trajectory', [])
            if pts:
                seqs.append([p.get('emotional_state', '') for p in pts])
        return seqs

    def _longest_common_subsequence(self, seq1, seq2):
        m, n = len(seq1), len(seq2)
        dp = [[0] * (n + 1) for _ in range(m + 1)]
        for i in range(1, m + 1):
            for j in range(1, n + 1):
                if seq1[i - 1] == seq2[j - 1]:
                    dp[i][j] = dp[i - 1][j - 1] + 1
                else:
                    dp[i][j] = max(dp[i - 1][j], dp[i][j - 1])
        return dp[m][n]

    def _compute_narrative_structure_similarity(self, functions1, functions2):
        if not functions1 or not functions2:
            return 0.0
        narrative_order = {
            'inciting_incident': 1,
            'rising_action': 2,
            'climax': 3,
            'falling_action': 4,
            'resolution': 5,
            'exposition': 0,
        }
        seq1 = self._create_narrative_sequence(functions1, narrative_order)
        seq2 = self._create_narrative_sequence(functions2, narrative_order)
        structure_sim = self._sequence_alignment_similarity(seq1, seq2)

        set1, set2 = set(functions1), set(functions2)
        presence_sim = len(set1.intersection(set2)) / len(set1.union(set2)) if set1.union(set2) else 0.0

        # Keep the same semantics as before (0.6/0.4 normalized)
        return 0.6 * structure_sim + 0.4 * presence_sim

    def _create_narrative_sequence(self, functions, order_mapping):
        ordered = []
        for f in functions:
            if f in order_mapping:
                ordered.append((order_mapping[f], f))
            else:
                ordered.append((2.5, f))
        ordered.sort(key=lambda x: x[0])
        return [f for _, f in ordered]

    def _sequence_alignment_similarity(self, seq1, seq2):
        if not seq1 or not seq2:
            return 0.0
        edit = self._edit_distance(seq1, seq2)
        denom = max(len(seq1), len(seq2))
        return 1.0 - (edit / denom) if denom > 0 else 0.0

    def _edit_distance(self, seq1, seq2):
        m, n = len(seq1), len(seq2)
        dp = [[0] * (n + 1) for _ in range(m + 1)]
        for i in range(m + 1):
            dp[i][0] = i
        for j in range(n + 1):
            dp[0][j] = j
        for i in range(1, m + 1):
            for j in range(1, n + 1):
                if seq1[i - 1] == seq2[j - 1]:
                    dp[i][j] = dp[i - 1][j - 1]
                else:
                    dp[i][j] = 1 + min(dp[i - 1][j], dp[i][j - 1], dp[i - 1][j - 1])
        return dp[m][n]

    def _cosine_similarity(self, v1, v2):
        if len(v1) == 0 or len(v2) == 0:
            return 0.0
        d = float(np.dot(v1, v2))
        n1 = float(np.linalg.norm(v1))
        n2 = float(np.linalg.norm(v2))
        if n1 == 0.0 or n2 == 0.0:
            return 0.0
        return d / (n1 * n2)


# ---------- Public helpers used by training ----------

def similarities_to_vector(similarities) -> np.ndarray:
    """
    Convert dict[SimilarityDimension -> SimilarityScore] to 7-dim vector
    (top-level dim scores) in DIM_ORDER.
    """
    return np.array([
        similarities[_DIM_TO_ENUM[name]].score for name in DIM_ORDER
    ], dtype=float)


def similarities_to_components(similarities) -> np.ndarray:
    """
    Convert dict[SimilarityDimension -> SimilarityScore] to 18-dim component vector
    in COMPONENT_ORDER. Reads component keys from SimilarityScore.details.
    Missing components default to 0.0.
    """
    out = []
    for dim_name, comp in COMPONENT_ORDER:
        dim_enum = _DIM_TO_ENUM[dim_name]
        det = similarities[dim_enum].details if similarities.get(dim_enum) else {}
        out.append(float(det.get(comp, 0.0)))
    return np.array(out, dtype=float)


def analyze_story_pair(story1, story2, embedding_model=None, weights_dict: Optional[Dict[str, Any]] = None):
    """
    Calculate per-dimension SimilarityScore and an overall score.
    If weights_dict has `intra`, use it to mix subcomponents inside each dimension.
    If weights_dict has top-level keys for DIM_ORDER, use them to mix across dimensions.
    """
    analyzer = StorySemanticAnalyzer(embedding_model=embedding_model, model_name=EMBEDDING_MODEL_NAME)

    intra = None
    if isinstance(weights_dict, dict) and "intra" in weights_dict:
        # expect: { dim_name: {component_key: weight} }
        intra = weights_dict.get("intra")

    similarities = analyzer.compute_story_similarity(story1, story2, intra=intra)

    # Default top-level 7-dim weights
    default_weights = {
        SimilarityDimension.THEMATIC: 0.25,
        SimilarityDimension.CHARACTER_SEMANTIC: 0.20,
        SimilarityDimension.EMOTIONAL_ARC: 0.20,
        SimilarityDimension.CAUSAL_STRUCTURE: 0.15,
        SimilarityDimension.NARRATIVE_FUNCTION: 0.10,
        SimilarityDimension.RELATIONSHIP_DYNAMICS: 0.05,
        SimilarityDimension.EVENT_SEMANTIC: 0.05,
    }

    # Learned 7-dim override if provided
    weights = default_weights
    if isinstance(weights_dict, dict):
        # accept either top-level floats or missing -> default
        wsum = 0.0
        cand = {}
        for name in DIM_ORDER:
            val = float(weights_dict.get(name, 0.0))
            if val < 0:
                val = 0.0
            cand[_DIM_TO_ENUM[name]] = val
            wsum += val
        if wsum > 0.0:
            # normalize to sum 1
            weights = {k: cand[k] / wsum for k in cand}

    overall_score = 0.0
    for dim, w in weights.items():
        overall_score += similarities[dim].score * w

    return similarities, float(overall_score)