rhythma_analysis.py

import logging
import os
from dataclasses import asdict, dataclass

import numpy as np

from rhythma_engine import RhythmaModulationEngine

try:
    from groq import Groq

    GROQ_AVAILABLE = True
except ImportError:
    Groq = None
    GROQ_AVAILABLE = False


LOGGER = logging.getLogger(__name__)


@dataclass
class AnalysisResult:
    emotional_state: str = "neutral"
    rhythm_pattern: str = "calm"
    transcription: str = ""
    session_profile: dict | None = None
    error: str | None = None

    def to_dict(self):
        return asdict(self)


@dataclass(frozen=True)
class SessionProfile:
    key: str
    title: str
    emotional_tone: str
    tone_center: float
    pattern: str
    modulation_type: str
    guidance: str
    reflection: str
    duration_hint: int
    brightness: float
    density: float
    shimmer: float
    breath_rate: float

    def to_dict(self):
        return asdict(self)


def _cosine_similarity(left, right):
    denominator = np.linalg.norm(left) * np.linalg.norm(right)
    if denominator == 0:
        return -1.0
    return float(np.dot(left, right) / denominator)


SESSION_PRESETS = {
    "anxious": SessionProfile(
        key="anxious",
        title="Grounding Tide",
        emotional_tone="Settling and steady",
        tone_center=396.0,
        pattern="calm",
        modulation_type="sine",
        guidance="Let your breath fall behind the pulse until the session feels steady.",
        reflection="This session favors stability over intensity.",
        duration_hint=15,
        brightness=0.25,
        density=0.45,
        shimmer=0.12,
        breath_rate=0.08,
    ),
    "stressed": SessionProfile(
        key="stressed",
        title="Soft Landing",
        emotional_tone="Unwinding and spacious",
        tone_center=417.0,
        pattern="relaxed",
        modulation_type="sine",
        guidance="Let the longer exhale soften the edges of the session.",
        reflection="This session eases pressure by widening the pulse.",
        duration_hint=18,
        brightness=0.22,
        density=0.38,
        shimmer=0.1,
        breath_rate=0.07,
    ),
    "calm": SessionProfile(
        key="calm",
        title="Quiet Harbor",
        emotional_tone="Easeful and settled",
        tone_center=432.0,
        pattern="calm",
        modulation_type="sine",
        guidance="Rest inside the repeating tone until it feels effortless.",
        reflection="This session keeps motion light to support an even mood.",
        duration_hint=15,
        brightness=0.32,
        density=0.28,
        shimmer=0.11,
        breath_rate=0.09,
    ),
    "sad": SessionProfile(
        key="sad",
        title="Low Ember",
        emotional_tone="Tender and reflective",
        tone_center=341.3,
        pattern="relaxed",
        modulation_type="sine",
        guidance="Allow the lower tone to hold the feeling without forcing it to lift.",
        reflection="This session gives weight and warmth to slower emotion.",
        duration_hint=16,
        brightness=0.18,
        density=0.33,
        shimmer=0.08,
        breath_rate=0.07,
    ),
    "angry": SessionProfile(
        key="angry",
        title="Ember Release",
        emotional_tone="Directed and discharging",
        tone_center=528.0,
        pattern="active",
        modulation_type="pulse",
        guidance="Track the sharper pulse until it turns from force into direction.",
        reflection="This session channels intensity into movement rather than compression.",
        duration_hint=12,
        brightness=0.5,
        density=0.62,
        shimmer=0.16,
        breath_rate=0.14,
    ),
    "fearful": SessionProfile(
        key="fearful",
        title="Shelter Light",
        emotional_tone="Protected and steadying",
        tone_center=384.0,
        pattern="calm",
        modulation_type="sine",
        guidance="Stay with the nearest tone and let it make the room feel smaller and safer.",
        reflection="This session reduces motion so attention can settle close to the body.",
        duration_hint=14,
        brightness=0.24,
        density=0.31,
        shimmer=0.09,
        breath_rate=0.08,
    ),
    "confused": SessionProfile(
        key="confused",
        title="North Star",
        emotional_tone="Clarifying and composed",
        tone_center=480.0,
        pattern="focused",
        modulation_type="sine",
        guidance="Follow one repeating detail until the rest of the field begins to organize.",
        reflection="This session simplifies the soundstage to support orientation.",
        duration_hint=14,
        brightness=0.34,
        density=0.3,
        shimmer=0.13,
        breath_rate=0.1,
    ),
    "happy": SessionProfile(
        key="happy",
        title="Bright Current",
        emotional_tone="Open and buoyant",
        tone_center=576.0,
        pattern="active",
        modulation_type="pulse",
        guidance="Enjoy the lift in the rhythm without pushing it faster.",
        reflection="This session keeps energy lively while protecting headroom.",
        duration_hint=12,
        brightness=0.56,
        density=0.4,
        shimmer=0.24,
        breath_rate=0.15,
    ),
    "focused": SessionProfile(
        key="focused",
        title="Clear Horizon",
        emotional_tone="Attentive and composed",
        tone_center=512.0,
        pattern="focused",
        modulation_type="sine",
        guidance="Stay with one thought and let the pulse keep the edges quiet.",
        reflection="This session narrows motion to support sustained attention.",
        duration_hint=20,
        brightness=0.4,
        density=0.35,
        shimmer=0.18,
        breath_rate=0.12,
    ),
    "relaxed": SessionProfile(
        key="relaxed",
        title="Open Meadow",
        emotional_tone="Loose and restorative",
        tone_center=444.0,
        pattern="relaxed",
        modulation_type="sine",
        guidance="Let the slow sway in the session keep your attention unforced.",
        reflection="This session favors softness and lingering resonance.",
        duration_hint=18,
        brightness=0.28,
        density=0.26,
        shimmer=0.12,
        breath_rate=0.08,
    ),
    "active": SessionProfile(
        key="active",
        title="Kinetic Bloom",
        emotional_tone="Motivated and rhythmic",
        tone_center=648.0,
        pattern="active",
        modulation_type="pulse",
        guidance="Let the pulse carry forward motion without turning rushed.",
        reflection="This session keeps energy articulated and bright.",
        duration_hint=10,
        brightness=0.6,
        density=0.48,
        shimmer=0.2,
        breath_rate=0.16,
    ),
    "neutral": SessionProfile(
        key="neutral",
        title="Still Current",
        emotional_tone="Balanced and open",
        tone_center=432.0,
        pattern="calm",
        modulation_type="sine",
        guidance="Listen for the simplest pulse and let it set the pace.",
        reflection="This session leaves space for your attention to settle naturally.",
        duration_hint=12,
        brightness=0.3,
        density=0.3,
        shimmer=0.1,
        breath_rate=0.1,
    ),
}


class RhythmaSymphAICore:
    """
    Interprets text and audio input to determine emotional state and rhythm pattern.
    """

    def __init__(self, use_groq=True, use_embeddings=True):
        self.emotional_states = [
            "anxious",
            "stressed",
            "calm",
            "sad",
            "angry",
            "fearful",
            "confused",
            "happy",
            "neutral",
            "focused",
            "relaxed",
            "active",
        ]
        self.rhythm_patterns = list(RhythmaModulationEngine.RHYTHM_CONFIGS.keys())

        self.groq_client = None
        self.use_groq = use_groq and GROQ_AVAILABLE
        self.use_embeddings = use_embeddings
        self.embedding_model = None
        self.emotional_embeddings = {}
        self.rhythm_embeddings = {}
        self._embedding_init_attempted = False

        if self.use_groq:
            self._initialize_groq_client()

    def _initialize_groq_client(self):
        api_key = os.environ.get("GROQ_API_KEY")
        if not api_key:
            LOGGER.warning("GROQ_API_KEY not found. Groq features disabled.")
            self.use_groq = False
            return

        try:
            self.groq_client = Groq(api_key=api_key)
        except Exception:
            LOGGER.exception("Failed to initialize Groq client.")
            self.use_groq = False

    def _ensure_embeddings_loaded(self):
        if not self.use_embeddings or self._embedding_init_attempted:
            return

        self._embedding_init_attempted = True
        try:
            from sentence_transformers import SentenceTransformer

            self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
            self.emotional_embeddings = {
                state: self.embedding_model.encode([state])[0]
                for state in self.emotional_states
            }
            self.rhythm_embeddings = {
                pattern: self.embedding_model.encode([pattern])[0]
                for pattern in self.rhythm_patterns
            }
        except ImportError:
            LOGGER.info(
                "SentenceTransformer not installed. Falling back to keyword matching."
            )
            self.use_embeddings = False
        except Exception:
            LOGGER.exception("Failed to initialize SentenceTransformer embeddings.")
            self.use_embeddings = False
            self.embedding_model = None
            self.emotional_embeddings = {}
            self.rhythm_embeddings = {}

    def detect_emotion_with_groq(self, input_text):
        if not self.use_groq or not self.groq_client:
            return None

        prompt = (
            "Analyze the user's feeling described below.\n"
            "Identify the single MOST prominent emotional state or intention from the following list:\n"
            f"{', '.join(self.emotional_states)}\n"
            "Focus on the core feeling expressed. Respond with ONLY the chosen state/intention from the list.\n"
            f"User's feeling: \"{input_text}\"\n"
            "State/Intention:"
        )

        try:
            chat_completion = self.groq_client.chat.completions.create(
                messages=[{"role": "user", "content": prompt}],
                model="llama-3.3-70b-versatile",
                max_tokens=15,
                temperature=0.2,
                stop=["\n"],
            )
            detected_emotion = chat_completion.choices[0].message.content.strip().lower()
            if detected_emotion in self.emotional_states:
                return detected_emotion
            return self.get_closest_emotional_state(detected_emotion)
        except Exception:
            LOGGER.exception("Groq emotion detection failed.")
            return None

    def get_closest_emotional_state(self, input_text):
        if not input_text:
            return "neutral"

        input_text_lower = input_text.lower()
        words = set(input_text_lower.split())
        for state in self.emotional_states:
            if state in words or state in input_text_lower:
                return state

        if "focus" in input_text_lower or "deep work" in input_text_lower:
            return "focused"

        self._ensure_embeddings_loaded()
        if self.embedding_model and self.emotional_embeddings:
            try:
                input_embedding = self.embedding_model.encode([input_text])[0]
                return max(
                    self.emotional_embeddings,
                    key=lambda state: _cosine_similarity(
                        input_embedding, self.emotional_embeddings[state]
                    ),
                )
            except Exception:
                LOGGER.exception("Semantic emotion matching failed.")

        return "neutral"

    def get_closest_rhythm_pattern(self, input_text=None, emotional_state=None):
        if emotional_state:
            mapping = {
                "anxious": "calm",
                "stressed": "relaxed",
                "calm": "calm",
                "sad": "relaxed",
                "angry": "active",
                "fearful": "calm",
                "confused": "focused",
                "happy": "active",
                "neutral": "calm",
                "focused": "focused",
                "relaxed": "relaxed",
                "active": "active",
            }
            return mapping.get(emotional_state, "calm")

        self._ensure_embeddings_loaded()
        if input_text and self.embedding_model and self.rhythm_embeddings:
            try:
                input_embedding = self.embedding_model.encode([input_text])[0]
                return max(
                    self.rhythm_embeddings,
                    key=lambda pattern: _cosine_similarity(
                        input_embedding, self.rhythm_embeddings[pattern]
                    ),
                )
            except Exception:
                LOGGER.exception("Semantic rhythm matching failed.")

        return "calm"

    def build_session_profile(self, emotional_state, rhythm_pattern):
        if emotional_state in SESSION_PRESETS:
            preset = SESSION_PRESETS[emotional_state]
        else:
            preset = SESSION_PRESETS["neutral"]
        profile = preset.to_dict()
        profile["pattern"] = rhythm_pattern or preset.pattern
        return profile

    def apply_profile_overrides(
        self,
        profile,
        tone_center=None,
        modulation_type=None,
        session_pattern=None,
    ):
        shaped_profile = dict(profile)
        if tone_center is not None and tone_center > 0:
            shaped_profile["tone_center"] = tone_center
        if modulation_type:
            shaped_profile["modulation_type"] = modulation_type
        if session_pattern:
            shaped_profile["pattern"] = session_pattern
        return shaped_profile

    def transcribe_audio(self, audio_path):
        if not self.use_groq or not self.groq_client:
            return None, "Transcription disabled: Groq client not available or API key missing."

        if not audio_path or not os.path.exists(audio_path):
            return None, "Transcription failed: Audio file path is invalid or missing."

        try:
            with open(audio_path, "rb") as audio_file:
                response = self.groq_client.audio.transcriptions.create(
                    file=(os.path.basename(audio_path), audio_file.read()),
                    model="whisper-large-v3",
                    response_format="json",
                )
            return response.text, None
        except Exception as exc:
            LOGGER.exception("Groq transcription failed.")
            return None, f"Error during Groq transcription: {exc}"

    def analyze_input(self, input_text=None, audio_path=None):
        result = AnalysisResult()
        text_to_analyze = None

        try:
            if audio_path and self.use_groq:
                transcribed_text, transcription_error = self.transcribe_audio(audio_path)
                if transcription_error:
                    result.error = transcription_error
                    result.transcription = f"[Transcription Error: {transcription_error}]"
                elif transcribed_text:
                    result.transcription = transcribed_text
                    text_to_analyze = transcribed_text

            if not text_to_analyze and input_text:
                text_to_analyze = input_text

            if text_to_analyze:
                detected_emotion = None
                if self.use_groq:
                    detected_emotion = self.detect_emotion_with_groq(text_to_analyze)

                result.emotional_state = detected_emotion or self.get_closest_emotional_state(
                    text_to_analyze
                )
            else:
                result.emotional_state = "neutral"

            result.rhythm_pattern = self.get_closest_rhythm_pattern(
                input_text=text_to_analyze,
                emotional_state=result.emotional_state,
            )
            result.session_profile = self.build_session_profile(
                result.emotional_state,
                result.rhythm_pattern,
            )
        except Exception as exc:
            LOGGER.exception("Unexpected error during input analysis.")
            result = AnalysisResult(
                session_profile=self.build_session_profile("neutral", "calm"),
                transcription=result.transcription,
                error=f"Unexpected error during input analysis: {exc}",
            )

        return result.to_dict()