#!/usr/bin/env python3
"""
AETHER-Micro Quality Head

4-dimensional quality assessment:
- Coherence (일관성)
- Relevance (관련성)
- Specificity (구체성)
- Helpfulness (유용성)
"""

import torch
import torch.nn as nn

from .configuration_aether_micro import AETHERMicroConfig


class AETHERMicroQualityHead(nn.Module):
    """
    Quality Head: 4-dimensional quality assessment

    Outputs:
        quality_scores: (batch_size, seq_length, 4)
            - Coherence (dim 0)
            - Relevance (dim 1)
            - Specificity (dim 2)
            - Helpfulness (dim 3)
    """

    def __init__(self, config: AETHERMicroConfig):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.quality_dim = config.quality_head_dim if hasattr(config, 'quality_head_dim') else 4

        # Quality projection
        self.quality_proj = nn.Linear(self.hidden_size, self.quality_dim, bias=True)

        # Initialize bias to encourage positive quality scores
        nn.init.constant_(self.quality_proj.bias, 0.5)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """
        Args:
            hidden_states: (batch_size, seq_length, hidden_size)

        Returns:
            quality_scores: (batch_size, seq_length, quality_dim)
                Values in [0, 1] range (sigmoid activation)
        """
        # Project to quality scores
        quality_scores = self.quality_proj(hidden_states)  # (batch, seq, 4)

        # Sigmoid to [0, 1] range
        quality_scores = torch.sigmoid(quality_scores)

        return quality_scores