"""Standalone VocBulwark Cage watermark detector.

Loadable with `AutoModel.from_pretrained(path, trust_remote_code=True)` without
the training repository. Wraps the Cage extractor and compares the extracted
bits to a fixed provenance watermark baked into the config.
"""
from dataclasses import dataclass
from math import comb

import torch
from transformers import PreTrainedModel
from transformers.modeling_outputs import ModelOutput

from .configuration_detector import CageDetectorConfig
from .cage_config import CageExtractorConfig
from .cage_extractor import CageExtractor


@dataclass
class DetectionOutput(ModelOutput):
    bits: torch.Tensor = None
    logits: torch.Tensor = None


class CageDetector(PreTrainedModel):
    config_class = CageDetectorConfig

    def __init__(self, config):
        super().__init__(config)
        ext_cfg = CageExtractorConfig(
            watermark_bits=config.watermark_bits,
            in_channels=config.in_channels,
            fine_kernel=config.fine_kernel,
            mid_kernel=config.mid_kernel,
            coarse_kernel=config.coarse_kernel,
            sample_rate=config.sample_rate,
        )
        self.cage_extractor = CageExtractor(ext_cfg)

    def _prep(self, audio, input_sample_rate):
        """Coerce input to [B, T] mono at the detector sample rate."""
        if not torch.is_tensor(audio):
            audio = torch.as_tensor(audio)
        p = next(self.parameters())
        audio = audio.to(device=p.device, dtype=p.dtype)
        if audio.dim() == 1:                       # [T] -> [1, T]
            audio = audio.unsqueeze(0)
        elif audio.dim() == 3:                     # [B, C, T] -> [B, T]
            audio = audio.mean(dim=1) if audio.shape[1] > 1 else audio.squeeze(1)
        if input_sample_rate is not None and input_sample_rate != self.config.sample_rate:
            import torchaudio
            audio = torchaudio.functional.resample(
                audio, input_sample_rate, self.config.sample_rate)
        return audio

    @staticmethod
    def _p_value(matches, n):
        """P(Binomial(n, 0.5) >= matches): chance an unrelated clip matches this well."""
        return sum(comb(n, i) for i in range(int(matches), n + 1)) / (2 ** n)

    @torch.no_grad()
    def extract(self, audio, input_sample_rate=None):
        """Return (bits, logits). bits: [B, n] in {0,1}; logits: [B, n]."""
        audio = self._prep(audio, input_sample_rate)
        logits = self.cage_extractor(audio)
        bits = (logits > 0).to(torch.int64)
        return bits, logits

    @torch.no_grad()
    def detect(self, audio, input_sample_rate=None):
        """Detect the fixed watermark. Returns a dict (or list of dicts for B>1):
        {matches, n_bits, p_value}. `matches` is how many of the model's fixed
        bits the clip carries; `p_value` is the chance an unrelated clip matches
        at least this well under Binomial(n_bits, 0.5). A clip from our model
        matches all n_bits (p_value ~ 0); an unrelated clip matches about half."""
        bits, _ = self.extract(audio, input_sample_rate)
        wstar = torch.tensor(self.config.fixed_watermark,
                             device=bits.device, dtype=bits.dtype)
        matches = (bits == wstar).sum(dim=-1).tolist()
        n = int(self.config.watermark_bits)
        results = [{
            "matches": int(m),
            "n_bits": n,
            "p_value": self._p_value(m, n),
        } for m in matches]
        return results[0] if len(results) == 1 else results

    @torch.no_grad()
    def forward(self, audio=None, input_sample_rate=None, **kwargs):
        bits, logits = self.extract(audio, input_sample_rate)
        return DetectionOutput(bits=bits, logits=logits)