from __future__ import annotations

from dataclasses import dataclass

import torch
from torch import Tensor, nn
import torch.nn.functional as F

from .model import SolPixOutput


@dataclass(frozen=True)
class FlowMatchingConfig:
    """Logit-normal timestep sampling used for the rectified-flow objective."""

    logit_mean: float = 0.0
    logit_std: float = 1.0


class REPAHead(nn.Module):
    """Optional train-only projection from SolPix features to teacher features."""

    def __init__(self, model_width: int = 512, teacher_width: int = 768):
        super().__init__()
        self.projection = nn.Sequential(
            nn.LayerNorm(model_width),
            nn.Linear(model_width, teacher_width),
        )

    def forward(self, features: Tensor) -> Tensor:
        return F.normalize(self.projection(features), dim=-1)


def sample_logit_normal_time(
    batch_size: int,
    device: torch.device,
    config: FlowMatchingConfig | None = None,
) -> Tensor:
    config = config or FlowMatchingConfig()
    logits = torch.randn(batch_size, device=device) * config.logit_std + config.logit_mean
    return logits.sigmoid()


def flow_matching_loss(
    model: nn.Module,
    clean_latents: Tensor,
    text_embeddings: Tensor,
    text_mask: Tensor | None = None,
    *,
    noise: Tensor | None = None,
    time: Tensor | None = None,
    config: FlowMatchingConfig | None = None,
) -> tuple[Tensor, Tensor, Tensor]:
    """Compute rectified-flow MSE for x_t=(1-t)x0+t*noise, v=noise-x0.

    Returns ``(loss, predicted_velocity, time)`` so callers can log timesteps and
    combine this term with auxiliary objectives such as REPA.
    """
    if clean_latents.ndim != 4:
        raise ValueError("clean_latents must have shape [batch, channels, height, width]")
    batch = clean_latents.shape[0]
    noise = torch.randn_like(clean_latents) if noise is None else noise
    if noise.shape != clean_latents.shape:
        raise ValueError("noise must have the same shape as clean_latents")
    time = sample_logit_normal_time(batch, clean_latents.device, config) if time is None else time
    if time.ndim == 0:
        time = time.expand(batch)
    if time.shape != (batch,):
        raise ValueError("time must be a scalar or a [batch] tensor")
    broadcast_time = time.to(clean_latents.dtype).reshape(batch, 1, 1, 1)
    noisy_latents = (1 - broadcast_time) * clean_latents + broadcast_time * noise
    target_velocity = noise - clean_latents
    predicted_velocity = model(noisy_latents, time, text_embeddings, text_mask)
    loss = F.mse_loss(predicted_velocity.float(), target_velocity.float())
    return loss, predicted_velocity, time


def representation_alignment_loss(
    output: SolPixOutput,
    teacher_features: Tensor,
    head: REPAHead,
    token_mask: Tensor | None = None,
) -> Tensor:
    """Cosine alignment loss for clean-image teacher features (REPA-style)."""
    if output.repa_features is None:
        raise ValueError("call SolPix(..., return_features=True) to obtain REPA features")
    prediction = head(output.repa_features)
    target = F.normalize(teacher_features.detach(), dim=-1)
    if prediction.shape != target.shape:
        raise ValueError(
            f"projected model features {tuple(prediction.shape)} do not match teacher features {tuple(target.shape)}"
        )
    per_token = 1 - (prediction * target).sum(dim=-1)
    if token_mask is None:
        return per_token.mean()
    if token_mask.shape != per_token.shape:
        raise ValueError("token_mask must have shape [batch, image_tokens]")
    weights = token_mask.to(per_token.dtype)
    return (per_token * weights).sum() / weights.sum().clamp_min(1)