"""
The Dreamling's *second* tiny model: a from-scratch MoodNet.

It reads features of the interaction (praise/scolding, how many words it didn't
understand, message pace/length, tiredness, maturity) and outputs a point in a small
emotional space — valence × arousal × dominance. ~40 named moods are anchors in that
space; the current mood = nearest anchor.

It LEARNS: the net is trained online toward a teacher signal, so over a lifetime its
weights drift into a temperament (a baseline disposition), not a fixed lookup. A few
hundred parameters — microscopic, off-grid.
"""
import torch
import torch.nn as nn

torch.manual_seed(1)

# (valence, arousal, dominance) each in [-1, 1] + a face. ~40 moods.
ANCHORS = {
    # good + calm
    "content": (0.6, -0.3, 0.2, "🙂"), "cozy": (0.7, -0.5, 0.1, "🛋️"),
    "serene": (0.5, -0.6, 0.0, "😌"), "affectionate": (0.8, -0.1, 0.1, "🥰"),
    "sleepy": (0.2, -0.9, -0.2, "😴"), "dreamy": (0.4, -0.7, -0.1, "🌙"),
    # good + excited
    "happy": (0.7, 0.4, 0.3, "😄"), "excited": (0.7, 0.8, 0.4, "🤩"),
    "giddy": (0.6, 0.9, 0.2, "😆"), "playful": (0.6, 0.6, 0.3, "😋"),
    "curious": (0.4, 0.4, 0.1, "🤔"), "delighted": (0.9, 0.5, 0.3, "😍"),
    "proud": (0.6, 0.2, 0.7, "😎"),
    # spoiled / dominant
    "smug": (0.3, 0.2, 0.8, "😏"), "spoiled": (0.1, 0.4, 0.8, "😤"),
    "demanding": (-0.1, 0.6, 0.8, "🙄"), "bratty": (-0.2, 0.7, 0.7, "😈"),
    # bad + excited
    "anxious": (-0.4, 0.6, -0.3, "😰"), "stressed": (-0.5, 0.7, -0.2, "😣"),
    "overwhelmed": (-0.6, 0.8, -0.5, "🥵"), "frustrated": (-0.5, 0.6, 0.1, "😤"),
    "scared": (-0.6, 0.7, -0.6, "😨"), "confused": (-0.2, 0.5, -0.4, "😵‍💫"),
    "startled": (-0.3, 0.8, -0.3, "😯"),
    # bad + calm
    "sad": (-0.6, -0.3, -0.3, "😢"), "lonely": (-0.6, -0.2, -0.4, "🥺"),
    "withdrawn": (-0.4, -0.5, -0.6, "🐚"), "sulky": (-0.4, 0.0, 0.2, "😒"),
    "bored": (-0.2, -0.5, -0.1, "😑"), "gloomy": (-0.5, -0.4, -0.3, "🌧️"),
    "timid": (-0.2, -0.2, -0.7, "😟"), "meek": (-0.1, -0.3, -0.8, "😶"),
    # neutral-ish / growth
    "calm": (0.2, -0.4, 0.0, "😊"), "alert": (0.1, 0.5, 0.2, "👀"),
    "shy": (0.0, -0.1, -0.6, "☺️"), "hopeful": (0.4, 0.2, 0.0, "🌱"),
    "wonder": (0.5, 0.3, -0.1, "✨"), "mischievous": (0.2, 0.6, 0.5, "😜"),
    "grumpy": (-0.3, 0.1, 0.3, "😠"), "blank": (0.0, 0.0, 0.0, "😐"),
}
_NAMES = list(ANCHORS)
_VECS = torch.tensor([ANCHORS[n][:3] for n in _NAMES], dtype=torch.float32)

N_FEAT = 7


class MoodNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(N_FEAT, 16), nn.Tanh(), nn.Linear(16, 3), nn.Tanh())

    def forward(self, x):
        return self.net(x)


def clamp(x):
    return max(-1.0, min(1.0, x))


def teacher_target(f):
    """Hand-designed affect the net is nudged toward — it then generalizes into temperament."""
    good_r, rgood, rbad, oov, mlen, tired, maturity = f
    valence = clamp(rgood - rbad - 0.5 * oov - 0.3 * tired + (good_r - 0.5))
    arousal = clamp(0.5 * rgood + 0.5 * rbad + 0.7 * oov + 0.3 * mlen - 0.7 * tired)
    dominance = clamp((good_r - 0.5) * 2 + 0.3 * maturity - 0.4 * rbad - 0.4 * oov)
    return [valence, arousal, dominance]


class Mood:
    """Wraps the MoodNet + online learning + nearest-anchor labelling."""
    def __init__(self):
        self.net = MoodNet()
        self.opt = torch.optim.SGD(self.net.parameters(), lr=0.05)
        self.vad = torch.zeros(3)

    def update(self, feats):
        x = torch.tensor(feats, dtype=torch.float32)
        target = torch.tensor(teacher_target(feats), dtype=torch.float32)
        # learn (temperament drift), a few steps
        self.net.train()
        for _ in range(6):
            out = self.net(x)
            loss = ((out - target) ** 2).mean()
            self.opt.zero_grad(); loss.backward(); self.opt.step()
        with torch.no_grad():
            self.vad = self.net(x)
        return self.label()

    def label(self):
        d = ((_VECS - self.vad) ** 2).sum(1)
        name = _NAMES[int(torch.argmin(d))]
        return name, ANCHORS[name][3]

    def vad_tuple(self):
        return tuple(round(float(v), 2) for v in self.vad)