# SPDX-License-Identifier: Apache-2.0
# Copyright 2026 Lucasgitguy-co
import math

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

from .data import grid_centers


class Student(nn.Module):
    """Small image-conditioned coordinate MLP; fixed cells avoid unordered point regression.

    15 Gaussian outputs + 2 auxiliary geometry logits. This is a coarse baseline,
    not an architecture expected to reproduce million-splat teacher detail.
    """
    def __init__(self, grid=16, width=128):
        super().__init__()
        if grid not in (8, 16, 32) or width < 16:
            raise ValueError("Unsupported grid or width")
        self.config = {"grid": grid, "width": width}
        self.grid = grid
        blocks, channels = [], 3
        for out in (24, 48, 96, 128):
            blocks.extend([nn.Conv2d(channels, out, 3, stride=2, padding=1), nn.GroupNorm(8, out), nn.SiLU()])
            channels = out
        self.encoder = nn.Sequential(*blocks, nn.AdaptiveAvgPool2d(4), nn.Flatten(),
                                     nn.Linear(128*4*4, width), nn.SiLU())
        centers = torch.from_numpy(grid_centers(grid))
        freq = 2**torch.arange(4, dtype=torch.float32) * math.pi * 2
        angles = centers[..., None] * freq
        self.register_buffer("centers", centers)
        self.register_buffer("coords", torch.cat([centers, angles.sin().flatten(1), angles.cos().flatten(1)], dim=1))
        self.query = nn.Sequential(nn.Linear(27, width), nn.SiLU())
        self.decoder = nn.Sequential(nn.Linear(width, width), nn.SiLU(), nn.Linear(width, width), nn.SiLU(), nn.Linear(width, 17))
        with torch.no_grad():
            self.decoder[-1].bias.zero_()
            self.decoder[-1].bias[0] = -2

    def forward(self, images):
        features = self.encoder(images)
        raw = self.decoder(features[:, None, :] + self.query(self.coords)[None])
        quat = raw[..., 10:14] + raw.new_tensor([1, 0, 0, 0])
        return {"occ": raw[..., :1], "xyz": self.centers[None] + raw[..., 1:4].tanh() / (2*self.grid),
                "rgb": raw[..., 4:7].sigmoid(),
                "scale": (raw[..., 7:10].clamp(-8, 5) - math.log(self.grid*3)).exp(),
                "quat": F.normalize(quat, dim=-1, eps=1e-8), "opacity": raw[..., 14:15].sigmoid(),
                "geometry": raw[..., 15:17]}


def balanced_bce(logits, target):
    error = F.binary_cross_entropy_with_logits(logits, target, reduction="none")
    positive = (error * target).sum() / target.sum().clamp_min(1)
    negative = (error * (1-target)).sum() / (1-target).sum().clamp_min(1)
    return (positive + negative) / 2


def distillation_loss(pred, target, grid):
    mask = target["occ"]
    def masked(value):
        return (value * mask).sum() / (mask.sum().clamp_min(1) * value.shape[-1])
    terms = {"occupancy": balanced_bce(pred["occ"], mask),
             "position": masked(((pred["xyz"] - target["xyz"]) * grid).square()),
             "colour": masked((pred["rgb"] - target["rgb"].clamp(0, 1)).square()),
             "scale": masked((pred["scale"].log() - target["scale"].clamp_min(1e-8).log()).square()),
             "rotation": masked(1 - (pred["quat"] * target["quat"]).sum(-1, keepdim=True).square().clamp(0, 1)),
             "opacity": masked((pred["opacity"] - target["opacity"]).square())}
    geometry = pred["geometry"].sum() * 0
    for i in range(2):
        available = target["geometry_valid"][:, i] > 0
        if available.any():
            geometry = geometry + balanced_bce(pred["geometry"][available, :, i:i+1], target["geometry"][available, :, i:i+1])
    terms["geometry"] = geometry
    total = terms["occupancy"] + terms["position"] + 2*terms["colour"] + .1*terms["scale"] + .1*terms["rotation"] + terms["opacity"] + .25*geometry
    return total, terms