"""The classifier as a fused Linear with fixed ternary weights and one free bias.

    model = FusedClassifier.from_hub()
    score, present = model(image_tensor)

Constructed without a backbone the module still exposes `head`, the same
decision applied to an already-pooled vector.
"""
import argparse
import json
import sys
from pathlib import Path

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

sys.path.insert(0, str(Path(__file__).resolve().parent))  # repo root, for `common`
from common import BACKBONE, pool  # noqa: E402

HERE = Path(__file__).resolve().parent


class FusedClassifier(nn.Module):
    """Backbone -> 40-dim slice -> ternary linear head -> binary decision.

    `retained_dims` indexes the 768-D pooled vector; `retained_weight` is +1 on
    the person-positive positions and -1 on the person-negative ones. The
    threshold is the only free parameter.
    """

    def __init__(self, backbone, pos_dims, neg_dims, threshold):
        super().__init__()
        self.backbone = backbone
        retained = list(pos_dims) + list(neg_dims)
        self.register_buffer('retained_dims', torch.tensor(retained, dtype=torch.long))
        w = torch.zeros(1, len(retained))
        w[0, :len(pos_dims)] = 1.0
        w[0, len(pos_dims):] = -1.0
        self.register_buffer('retained_weight', w)
        self.threshold = nn.Parameter(torch.tensor(float(threshold)))

    def head(self, pooled):
        """(..., 768) pooled vector -> (score, present), no backbone involved."""
        retained = pooled.index_select(-1, self.retained_dims)
        score = F.linear(retained, self.retained_weight).squeeze(-1)
        return score, score > self.threshold

    @torch.inference_mode()
    def forward(self, x):
        """x: (B, 3, 768, 768) normalized. Returns (score (B,), present (B,))."""
        if self.backbone is None:
            raise RuntimeError('constructed without a backbone; use .head(pooled)')
        dev = 'cuda' if x.is_cuda else 'cpu'
        with torch.autocast(dev, dtype=torch.bfloat16):
            out = self.backbone.forward_features(x)
        return self.head(pool(out['x_norm_patchtokens'].float()))

    @classmethod
    def from_config(cls, backbone=None, classifier_json=None):
        c = json.loads(Path(classifier_json or HERE / 'classifier.json').read_text())
        return cls(backbone, c['pos_dims'], c['neg_dims'], c['threshold'])

    @classmethod
    def from_hub(cls, repo_or_path=None, classifier_json=None):
        from common.models import load_backbone
        return cls.from_config(load_backbone(repo_or_path or BACKBONE), classifier_json)


if __name__ == '__main__':
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument('--classifier', type=Path, default=HERE / 'classifier.json')
    args = ap.parse_args()
    m = FusedClassifier.from_hub(classifier_json=args.classifier).eval()
    n_all = sum(p.numel() for p in m.parameters())
    n_backbone = sum(p.numel() for p in m.backbone.parameters())
    print(f'total params:    {n_all:,}')
    print(f'backbone params: {n_backbone:,}')
    print(f'head params:     {n_all - n_backbone} '
          f'(one learnable threshold; weights are fixed buffers)')
    print(f'retained dims:   {m.retained_dims.numel()}')