Proxy-based contrastive learning for retinal ganglion cell embeddings using response repeats

import functools
import json
import logging
import math
from pathlib import Path
from typing import List, Literal, Sequence
import einops
import kdai._logging
import kdai.train
import numpy as np
from numpy.typing import ArrayLike
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data
import torchinfo


_logger = logging.getLogger(__name__)

OUTPUT_ROOT_DIR = "./out/exp"
LOG_FILENAME = "train.log"


class StimSegment:
    def __init__(
        self, stim_type_id, cell_ids, repeat_starts, spikes, sample_freq, label
    ):
        if len(cell_ids) != len(spikes):
            raise ValueError(f"{len(cell_ids)=} != {len(spikes)=}")
        self.n_cells = len(cell_ids)
        # Each stimulus can have a different number of repeats.
        self.n_repeats = [len(spikes[i]) for i in range(len(spikes))]
        self._len = sum(self.n_repeats)
        self._cumulative_starts = np.cumsum(
            np.concatenate([[0], self.n_repeats])
        )
        self.stim_type_id = stim_type_id
        self.cell_ids = cell_ids
        self.repeat_starts = repeat_starts
        self.spikes = spikes
        self.sample_freq = sample_freq
        self.label = label

    def __str__(self):
        return (
            f"StimSegment: {self.label}, "
            f"n_cells: {self.n_cells}, "
            f"len: {self.__len__()}, "
            f"first repeat t (s): {self.repeat_starts[0]/self.sample_freq}, "
        )

    def __len__(self):
        return self._len

    def max_rep(self):
        return max(self.n_repeats)

    def max_spikes_in_rep(self):
        res = max(len(rep) for c in self.spikes for rep in c)
        return res

    def __getitem__(self, idx):
        if idx < 0:
            raise ValueError("Negative indexes not yet supported")
        cell_idx = (
            np.searchsorted(self._cumulative_starts, idx, side="right") - 1
        )
        repeat = idx - self._cumulative_starts[cell_idx]
        try:
            spikes = self.spikes[cell_idx][repeat]
        except IndexError:
            print(
                f"{len(self)=}, {idx=}, {cell_idx=}, {repeat=}, {len(self.spikes)=}"
            )
            raise IndexError
        return cell_idx, repeat, spikes

    def cell_idx(self, cell_id):
        return self.cell_ids.index(cell_id)

    def spikes_by_cell(self, cell_idx):
        res = self.spikes[cell_idx]
        return res


def load_stim_segments(json_file_path):
    segs = []
    with open(json_file_path) as f:
        js = json.load(f)
        sample_freq = js["sample_freq"]
        n_cells = js["meta"]["n_cells"]
        n_stims = len(js["meta"]["stim_types"])
        for s in js["segments"]:
            segment = StimSegment(
                s["stim_type_id"],
                s["cell_ids"],
                s["repeat_starts"],
                s["spikes"],
                sample_freq,
                s["label"],
            )
            segs.append(segment)
    return segs, n_cells, n_stims, sample_freq


class Dataset:
    def __init__(
        self, segments: Sequence[StimSegment], sample_freq, augment=False
    ):
        # Currently, we impose that all cells must be in all segments
        cells = set(segments[0].cell_ids)
        for s in segments:
            if set(s.cell_ids) != cells:
                raise ValueError(f"{s.cell_ids} != {cells}")
        self.segments = segments
        self.sample_freq = sample_freq
        self.augment = augment
        self.spike_jitter_ms = 1.0  # SD of per-spike jitter (ms)
        self.start_t_jitter_s = 20 * 60  # SD of start_t jitter (s = 20 min)
        self._len = sum(len(s) for s in self.segments)
        self._cumulative_starts = np.cumsum(
            np.concatenate([[0], [len(c) for c in self.segments]])
        )

    def frame_to_ms(self, f):
        return (f / self.sample_freq) * 1000

    def frame_to_sec(self, f):
        return f / self.sample_freq

    def __len__(self):
        return self._len

    def n_cells(self):
        return self.segments[0].n_cells

    def __getitem__(self, idx):
        # Don't try the following line:
        ## cell_idx = np.searchsorted(self.cumulative_starts, idx, side="left")
        # Why? Example case: [0, 5, 12], idx=4 => 1 (correct is 0)
        seg_idx = (
            np.searchsorted(self._cumulative_starts, idx, side="right") - 1
        )
        inner_idx = idx - self._cumulative_starts[seg_idx]
        assert seg_idx >= 0, seg_idx
        assert inner_idx >= 0, inner_idx
        seg = self.segments[seg_idx]
        cell_idx, rep, spikes = seg[inner_idx]
        spikes = torch.tensor(spikes, dtype=torch.float32)
        spikes = self.frame_to_ms(spikes)  # ms
        start_t = self.frame_to_sec(
            self.segments[seg_idx].repeat_starts[0]
        )  # seconds
        if self.augment:
            # Jitter. Use torch RNG (not np.random) so noise is seeded 
            # per-worker, not identical across dataloader workers. 
            spikes = spikes + torch.randn_like(spikes) * self.spike_jitter_ms
            spikes = spikes.clamp_min(0.0)  # times stay >= 0
            # No harm in also jittering the start time used as context.
            start_t = start_t + float(torch.randn(())) * self.start_t_jitter_s
        res = {
            "cell_idx": cell_idx,
            "spikes": spikes,
            "stim_idx": seg.stim_type_id,
            "repeat": rep,
            "start_t": start_t,
            "label": seg.label,
        }
        return res


def collate_fn(batch):
    B = len(batch)
    max_n_spikes = max(len(b["spikes"]) for b in batch)

    def spikes_and_mask(b_idx):
        spikes = batch[b_idx]["spikes"]
        out_spikes = torch.zeros((max_n_spikes,), dtype=torch.float32)
        mask = torch.zeros((max_n_spikes,), dtype=torch.bool)
        out_spikes[0 : len(spikes)] = spikes
        mask[: len(spikes)] = True
        return out_spikes, mask

    spikes, mask = zip(*(spikes_and_mask(b) for b in range(B)))
    spikes = torch.stack(spikes)
    mask = torch.stack(mask)

    to_tensor = lambda key, dtype: torch.tensor(
        [b[key] for b in batch], dtype=dtype
    )
    res = (
        to_tensor("cell_idx", torch.long),
        spikes,
        mask,
        to_tensor("stim_idx", torch.long),
        to_tensor("repeat", torch.long),
        to_tensor("start_t", torch.float32),
    )
    return res


class DsMgr(kdai.train.DatasetManager):
    def __init__(self, segs, val_seg_idx, sample_freq, augment):
        self.segs = segs
        self.val_seg_idx = val_seg_idx
        self.sample_freq = sample_freq
        self.augment = augment
        self._train_ds = Dataset(
            [self.segs[i] for i in range(len(segs)) if i != val_seg_idx],
            sample_freq,
            augment=self.augment,
        )
        self._val_ds = Dataset(
            [segs[self.val_seg_idx]], sample_freq, augment=self.augment
        )
        self._test_ds = []

    def train_ds(self) -> torch.utils.data.Dataset:
        return self._train_ds

    def val_ds(self) -> torch.utils.data.Dataset:
        return self._val_ds

    def test_ds(self) -> torch.utils.data.Dataset:
        return self._test_ds

    @functools.cache
    def block_ds(self, block_id, augment=None):
        if augment is None:
            augment = self.augment
        return Dataset([self.segs[block_id]], self.sample_freq, augment=augment)


def get_optim_param_groups(module):
    """Return parameters with and without decay, in separate groups.

    This function signature is designated by train.py.
    """
    with_grad = {n: p for n, p in module.named_parameters() if p.requires_grad}
    # Collect set of all embedding parameters.
    embedding_params = set()
    for module in module.modules():
        if isinstance(module, nn.Embedding):
            for param in module.parameters(recurse=False):
                embedding_params.add(id(param))

    def with_decay_fn(param):
        single_dim = param.dim() == 1
        is_embedding = id(param) in embedding_params
        with_decay = not (single_dim or is_embedding)
        return with_decay

    with_decay_params = []
    without_decay_params = []
    with_decay_names = []
    without_decay_names = []
    for n, p in with_grad.items():
        if with_decay_fn(p):
            with_decay_params.append(p)
            with_decay_names.append(n)
        else:
            without_decay_params.append(p)
            without_decay_names.append(n)
    _logger.info(f"Params with decay: {', '.join(with_decay_names)}")
    _logger.info(f"Params without decay: {', '.join(without_decay_names)}")
    return [with_decay_params, without_decay_params]


class Normalize(nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer("mean", torch.zeros(size=[]))
        self.register_buffer("sd", torch.ones(size=[]))

    def set_mean_sd(self, m: ArrayLike, sd: ArrayLike):
        self.mean = torch.tensor(m)
        self.sd = torch.tensor(sd)
        if self.mean.ndim != 0:
            raise ValueError(
                f"Expected mean to be a scalar; got ({self.mean.ndim}) dims."
            )
        if self.sd.ndim != 0:
            raise ValueError(
                f"Expected sd to be a scalar; got ({self.sd.ndim}) dims."
            )

    def forward(self, x, shift=True, scale=True):
        shift = self.mean if shift else 0
        scale = self.sd if scale else 1
        return (x - shift) / scale

    def denorm(self, x):
        return x * self.sd + self.mean


class IdxEmbed(nn.Module):
    def __init__(self, length, n_embd):
        super().__init__()
        self.n_embd = n_embd
        self.length = length
        self.pos_embd = nn.Embedding(self.length, self.n_embd)
        self.register_buffer(
            "range_cache",
            einops.rearrange(torch.arange(self.length), "L -> 1 L"),
            persistent=False,
        )

    def forward(self):
        return self.pos_embd(self.range_cache)


class ValueEmbed(nn.Module):
    """Convert scalars to a fixed sized vector.

    Can be used for timestamps or single dimension data.
    """

    def __init__(self, n_embd):
        super().__init__()
        self.n_embd = n_embd
        self.register_buffer("slope", torch.zeros(self.n_embd // 2))

    @torch.no_grad()
    def init_weights(self, max_range, epsilon):
        # Use up most of the range [0, 2π] by landing on 3/4 * 2π.
        min_scale = (3 / 2 * math.pi) / max_range
        max_scale = (3 / 2 * math.pi) / epsilon
        self.slope.copy_(
            torch.exp(
                torch.linspace(
                    math.log(min_scale),
                    math.log(max_scale),
                    self.n_embd // 2,
                    dtype=torch.float32,
                    requires_grad=False,
                )
            )
        )

    def forward(self, t):
        """
        If using a batch of values, you will want to send in (B, 1)

        Often you see:
            t = self.slope * t.unsqueeze(-1)

        but to support various uses, not just for timesteps, we don't do this
        internally here.
        """
        t = self.slope * t
        t = torch.cat([t.sin(), t.cos()], dim=-1).to(dtype=torch.float32)
        return t


class Attention(nn.Module):
    def __init__(self, n_embd, heads=4, dropout=0.0, project_out=True):
        super().__init__()
        self.project_out = project_out
        self.dim_head = int(n_embd / heads)
        if self.dim_head != n_embd / heads:
            raise ValueError(
                f"n_embd ({n_embd}) must be divisible by heads ({heads})"
            )
        self.heads = heads
        self.to_qkv = nn.Linear(n_embd, n_embd * 3, bias=False)
        self.to_out = (
            nn.Sequential(
                nn.Linear(n_embd, n_embd),
                nn.Dropout(dropout),
            )
            if self.project_out
            else nn.Identity()
        )
        self.norm = nn.RMSNorm(n_embd)
        self.dropout_rate = dropout

    def forward(self, x, mask=None):
        x = self.norm(x)
        qkv = self.to_qkv(x).chunk(3, dim=-1)
        q, k, v = map(
            lambda t: einops.rearrange(t, "b t (h c) -> b h t c", h=self.heads),
            qkv,
        )

        # q = q * self.scale
        # raw_attn = torch.einsum("b h i d, b h j d -> b h i j", q, k)
        # attn = raw_attn.softmax(dim=-1)
        # attn = self.dropout(attn)
        # out = torch.einsum("b h i j, b h j d -> b h i d", attn, v)
        dropout_rate = self.dropout_rate if self.training else 0.0
        out = F.scaled_dot_product_attention(
            q, k, v, dropout_p=dropout_rate, attn_mask=mask
        )
        out = einops.rearrange(out, "b h t d -> b t (h d)")
        out = self.to_out(out)
        return out


class FeedForward(nn.Module):
    def __init__(self, n_embd, hidden_dim, dropout=0.0):
        super().__init__()
        self.input_norm = nn.LayerNorm(n_embd)
        self.linear1 = nn.Linear(n_embd, hidden_dim)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(hidden_dim, n_embd)

    def forward(self, x):
        x = self.input_norm(x)
        x = self.linear1(x)
        x = self.dropout(F.gelu(x))
        x = self.linear2(x)
        x = self.dropout(x)
        return x


class Transformer(nn.Module):
    def __init__(self, dim, depth, heads, mlp_dim, dropout=0.0):
        super().__init__()
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(
                nn.ModuleList(
                    [
                        Attention(
                            dim,
                            heads,
                            dropout,
                            project_out=True,
                        ),
                        FeedForward(dim, mlp_dim, dropout=dropout),
                    ]
                )
            )

    def forward(self, x, attn_mask=None):
        if attn_mask is not None:
            attn_mask = einops.rearrange(attn_mask, "b s -> b 1 1 s")
        for attn, ff in self.layers:
            x = attn(x, attn_mask) + x
            x = ff(x) + x
        return x


class ProxySLL(nn.Module):

    def __init__(
        self,
        n_cells,
        n_stim,
        n_rep,
        n_layer,
        n_head,
        head_dim,
        out_dim,
        dropout=0.0,
    ):
        super().__init__()
        self.n_cells = n_cells
        self.out_dim = out_dim
        self.emb_d = n_head * head_dim
        self.n_special_tokens = 4  # 3 context and 1 cls token
        if self.emb_d % 2 != 0:
            raise ValueError("Embed length must be even.")
        mlp_expansion = 4
        self.mlp_d = self.emb_d * mlp_expansion
        # Learnable
        self.tf = Transformer(
            self.emb_d, n_layer, n_head, self.mlp_d, dropout=dropout
        )
        self.cls_token = nn.Parameter(torch.zeros(self.emb_d))
        self.repeat_embed = nn.Embedding(n_rep, self.emb_d)
        self.stim_embed = nn.Embedding(n_stim, self.emb_d)
        self.start_t_embed = ValueEmbed(self.emb_d)
        self.spike_t_embed = ValueEmbed(self.emb_d)
        self.out_proj = nn.Linear(self.emb_d, out_dim)
        self.out_norm = nn.RMSNorm(self.out_dim)
        ## The output embedding
        self.cell_embd = nn.Embedding(n_cells, self.out_dim)
        self._init_weights()

    @torch.no_grad()
    def _init_weights(self):
        self.cls_token.normal_(0.0, 1.0)
        # Defaults:
        # repeat_embed, stim_embd, rec_t_embed all get N(0, 1) inits.

    def optim_param_groups(self):
        """Return parameters with and without decay, in separate groups.

        This function signature is designated by train.py.
        """
        res = get_optim_param_groups(self)
        return res

    def cell_embeddings(self):
        return self.cell_embd.weight

    @torch.no_grad()
    def init_input_scale(
        self,
        spike_t_max_range,
        spike_t_epsilon,
        start_t_max_range,
        start_t_epsilon,
    ):
        self.spike_t_embed.init_weights(spike_t_max_range, spike_t_epsilon)
        self.start_t_embed.init_weights(start_t_max_range, start_t_epsilon)

    def embed_ctx(self, stim_id, repeat, start_t):
        e_s = einops.rearrange(self.stim_embed(stim_id), "b e -> b 1 e")
        e_r = einops.rearrange(self.repeat_embed(repeat), "b e-> b 1 e")
        e_rt = self.start_t_embed(einops.rearrange(start_t, "b -> b 1 1"))
        res = torch.cat([e_s, e_r, e_rt], dim=1)
        return res

    def forward(self, spike_ts, valid_mask, stim_id, repeat, start_t):
        """Encode a list of spike times given as a padded batch.

        Args:
            spike_ts: spike times, in ms.
        """
        B, S = spike_ts.shape
        dv = spike_ts.device
        assert stim_id.shape == repeat.shape == start_t.shape == (B,)
        assert valid_mask.dtype == torch.bool
        assert valid_mask.shape == (B, S), valid_mask.shape
        n_used = S + self.n_special_tokens
        ctx = self.embed_ctx(stim_id, repeat, start_t)
        assert ctx.shape == (B, 3, self.emb_d)
        x = torch.cat(
            [
                self.spike_t_embed(einops.rearrange(spike_ts, "b s -> b s 1")),
                ctx,
                einops.repeat(self.cls_token, "d -> b 1 d", b=B),
            ],
            dim=1,
        )
        valid_mask = torch.cat(
            [
                valid_mask,
                torch.ones(
                    B, self.n_special_tokens, dtype=valid_mask.dtype, device=dv
                ),
            ],
            dim=1,
        )
        assert x.shape == (B, n_used, self.emb_d)
        x = self.tf(x, valid_mask)
        # Take the value from the token line.
        x = self.out_proj(x[:, -1])
        x = self.out_norm(x)
        return x


EvalMode = Literal["loss", "info"]


class JoinEmbedTrainable(kdai.train.BaseTrainable):
    ds_mgr: DsMgr

    def __init__(self, ds_mgr, model, label, eval_mode: EvalMode = "loss"):
        super().__init__(ds_mgr, model, label)
        # collate_fn is used by BaseTrainable
        self.collate_fn = collate_fn

        self.n_cells = model.n_cells
        self.loss_temperature = 0.5
        self.eval_mode = eval_mode
        self._init_weights()

    @torch.no_grad()
    def _init_weights(self):
        # Hard-coded ranges.
        # Per-repeat spike times: max=10 min, epsilon=0.1 ms
        self.spike_t_max_range = 10 * 60 * 1000  # milliseconds
        self.spike_t_epsilon = 0.1  # milliseconds
        # Repeat start time: max=24 hours, epsilon=5min
        self.t_max_range = 24 * 60 * 60  # seconds
        self.t_epsilon = 5 * 60  # seconds
        self.model.init_input_scale(
            self.spike_t_max_range,
            self.spike_t_epsilon,
            self.t_max_range,
            self.t_epsilon,
        )

    def similarity(self, resp_embd, cell_embd, cell_idx):
        B, E = resp_embd.shape
        assert cell_embd.shape == (self.n_cells, E), cell_embd.shape
        assert cell_idx.shape == (B,)
        assert cell_idx.dtype == torch.long

        resp_embd = F.normalize(resp_embd, dim=-1)
        cell_embd = F.normalize(cell_embd, dim=-1)
        rc_sim = torch.einsum("be,ce -> bc", resp_embd, cell_embd)
        return rc_sim

    def _loss(self, resp_embd, cell_embd, cell_idx, temperature):
        assert temperature > 0
        sim = self.similarity(resp_embd, cell_embd, cell_idx)
        logits = sim / temperature
        l = F.cross_entropy(logits, cell_idx)
        return l, {}

    def _forward_no_loss(self, sample):
        cell_idx, spike_ts, valid_mask, stim_idx, repeat, start_t = (
            s.cuda() for s in sample
        )
        resp_embd = self.model(spike_ts, valid_mask, stim_idx, repeat, start_t)
        return resp_embd, cell_idx

    def _forward(self, sample):
        resp_embd, cell_idx = self._forward_no_loss(sample)
        loss, named_losses = self._loss(
            resp_embd,
            self.model.cell_embeddings(),
            cell_idx,
            self.loss_temperature,
        )
        return resp_embd, loss, named_losses

    def forward(self, sample):
        x, loss, _ = self._forward(sample)
        return x, loss

    def eval_loss(self, dl):
        """Evaluate loss only."""
        loss_meter = kdai._logging.Meter("loss")
        for sample in dl:
            N = len(sample[0])
            _, loss = self.forward(sample)
            loss_meter.update(loss.item(), N)
        results = {
            "metrics": [
                kdai._logging.loss_metric(loss_meter.avg),
            ],
        }
        return results

    def empirical_centroids(self, dl):
        sums = torch.zeros(
            self.n_cells,
            self.model.out_dim,
            dtype=torch.float32,
            device=self.in_device(),
        )
        counts = torch.zeros(
            self.n_cells,
            dtype=torch.long,
            device=self.in_device(),
        )

        for sample in dl:
            embd, cell_idx = self._forward_no_loss(sample)
            sums.index_add_(0, cell_idx, embd)
            counts.index_add_(0, cell_idx, torch.ones_like(cell_idx))

        valid = counts > 0
        centroids = torch.zeros_like(sums)
        centroids[valid] = sums[valid] / counts[valid, None]
        return centroids, valid

    def neighborhood_agreement(self, emb_a, emb_b, valid_a, valid_b, k):
        valid = valid_a & valid_b
        a = F.normalize(emb_a[valid], dim=-1)
        b = F.normalize(emb_b[valid], dim=-1)

        # similarity between embeddings (within a block), and ignore self.
        a_sim = a @ a.T
        b_sim = b @ b.T
        a_sim.fill_diagonal_(-torch.inf)
        b_sim.fill_diagonal_(-torch.inf)

        # Each embedding's nearest k neighboring embeddings
        nn_a = a_sim.topk(k, dim=-1).indices
        nn_b = b_sim.topk(k, dim=-1).indices

        nn_a = einops.rearrange(nn_a, "c x -> c x 1")
        nn_b = einops.rearrange(nn_b, "c x -> c 1 x")
        eqs = (nn_a == nn_b).any(dim=-1)
        frac_eq = eqs.float().mean()
        return frac_eq

    def cross_block_agreement(
        self, dl_fn, b_a_idx, b_b_idx, ks=None
    ) -> List[float]:
        if ks is None:
            ks = [5, 10, 20]
        dl_a = dl_fn(self.ds_mgr.block_ds(b_a_idx), collate_fn=self.collate_fn)
        dl_b = dl_fn(self.ds_mgr.block_ds(b_b_idx), collate_fn=self.collate_fn)
        emb_a, valid_a = self.empirical_centroids(dl_a)
        emb_b, valid_b = self.empirical_centroids(dl_b)
        res = []
        for k in ks:
            res_k = self.neighborhood_agreement(
                emb_a, emb_b, valid_a, valid_b, k
            ).item()
            res.append(res_k)
        return res

    def eval_metrics(self, dl):
        meters = {
            name: kdai._logging.Meter(name)
            for name in [
                "loss",
                "acc",
                "similarity",
                "hard-neg-sim",
                "target-hard-margin",
                "top-5",
                "top-10",
                "mean-rank",
                "median-rank",
            ]
        }
        for sample in dl:
            N = len(sample[0])
            # Recreate forward and loss.
            # Note, this is a source of errors if it doesn't keep
            # in sync with forward.
            resp_embd, cell_idx = self._forward_no_loss(sample)
            sim = self.similarity(
                resp_embd,
                self.model.cell_embeddings(),
                cell_idx,
            )
            logits = sim / self.loss_temperature
            loss = F.cross_entropy(logits, cell_idx)
            meters["loss"].update(loss.item(), N)
            pred = logits.argmax(dim=-1)
            acc = (pred == cell_idx).float().mean().item()
            meters["acc"].update(acc, N)
            # Similarity to true
            b_idx = torch.arange(N, device=sim.device)
            target_sim = sim[b_idx, cell_idx]
            meters["similarity"].update(target_sim.mean().item(), N)
            # Ranks
            rank = (sim > target_sim[:, None]).sum(dim=1) + 1
            meters["top-5"].update((rank <= 5).float().mean().item(), N)
            meters["top-10"].update((rank <= 10).float().mean().item(), N)
            meters["mean-rank"].update(rank.float().mean().item(), N)
            meters["median-rank"].update(rank.float().median().item(), N)

            neg_sim = sim.clone()
            neg_sim[b_idx, cell_idx] = -torch.inf
            hardest_neg_sim = neg_sim.max(dim=-1).values
            margin = target_sim - hardest_neg_sim
            meters["hard-neg-sim"].update(hardest_neg_sim.mean().item(), N)
            meters["target-hard-margin"].update(margin.mean().item(), N)

        results = {name: meter.avg for name, meter in meters.items()}
        return results

    @torch.no_grad()
    def evaluate_train(self, dl_fn):
        dl = dl_fn(self.ds_mgr.train_ds(), collate_fn=self.collate_fn)
        # Evaluate on the train ds currently only supports basic eval.
        loss_metrics = self.eval_loss(dl)
        results = {"metrics": loss_metrics}
        return results

    @torch.no_grad()
    def eval_detail(self, dl_fn):
        dl = dl_fn(self.ds_mgr.val_ds(), collate_fn=self.collate_fn)
        metrics = self.eval_metrics(dl)

        def to_log_metric(m, v):
            if m == "loss":
                return kdai._logging.loss_metric(v)
            else:
                return kdai._logging.Metric(m, v, increasing=False)

        metrics = [to_log_metric(m, v) for m, v in metrics.items()]

        # Metrics that use their own dl.
        b_a_idx = 3
        b_b_idx = 10
        ks = [5, 10, 20]
        cx_name = lambda k: f"b{b_a_idx}xb{b_b_idx}_agreement_k{ks[k]}"
        cx_ks = self.cross_block_agreement(dl_fn, b_a_idx, b_b_idx, ks)
        metrics.extend(
            [
                kdai._logging.Metric(cx_name(i), cx_ks[i], increasing=True)
                for i in range(len(ks))
            ]
        )

        results = {
            "metrics": metrics,
            "embedding": kdai._logging.Embeddings(
                self.model.cell_embeddings().cpu().numpy(),
                [str(i) for i in range(self.model.n_cells)],
            ),
        }
        return results

    @torch.no_grad()
    def evaluate_val(self, dl_fn):
        if self.eval_mode == "loss":
            dl = dl_fn(self.ds_mgr.val_ds(), collate_fn=self.collate_fn)
            return self.eval_loss(dl)
        elif self.eval_mode == "info":
            res = self.eval_detail(dl_fn)
        else:
            raise ValueError(f"Unknown eval mode: {self.eval_mode}")
        return res

    @torch.no_grad()
    def model_summary(self, batch_size: int) -> str:
        dl = torch.utils.data.DataLoader(
            self.train_ds(), batch_size=batch_size, collate_fn=self.collate_fn
        )
        sample = next(iter(dl))
        cell_idx, spike_ts, valid_mask, stim_idx, repeat, start_t = (
            s.cuda() for s in sample
        )
        res = torchinfo.summary(
            self.model,
            input_data=(spike_ts, valid_mask, stim_idx, repeat, start_t),
            col_names=["input_size", "output_size", "mult_adds", "num_params"],
            device=self.in_device(),
            depth=4,
        )
        return res


def create_tr():
    path = Path("data/dataset-4-0-0-0.json")
    segs, n_cells, n_stim, sample_freq = load_stim_segments(path)
    max_n_spikes = max(s.max_spikes_in_rep() for s in segs)
    _logger.info(f"Maximum spikes in a repeat: {max_n_spikes}")
    n_rep = max(s.max_rep() for s in segs)
    b05_idx = 5
    augment = True
    ds_mgr = DsMgr(segs, b05_idx, sample_freq, augment)
    # Model
    n_layer = 8
    n_head = 8
    head_dim = 8
    out_dim = 64
    dropout = 0.15
    model = ProxySLL(
        n_cells,
        n_stim,
        n_rep,
        n_layer,
        n_head,
        head_dim,
        out_dim,
        dropout=dropout,
    )
    model.compile()
    # Trainable
    tr = JoinEmbedTrainable(
        ds_mgr, model, "basic-contrastive", eval_mode="info"
    )
    tr.forward = torch.compile(tr.forward)
    return tr


def start_logging():
    ver_parts = kdai._logging.version_labels_from_script_dir()
    root_logger = kdai._logging.setup_logging(logging.INFO)
    out_dir = kdai._logging.get_outdir(OUTPUT_ROOT_DIR, ver_parts)
    ver_parts.append(out_dir.name)
    kdai._logging.enable_file_logging(out_dir / LOG_FILENAME)
    kdai._logging.snapshot_importing_script(out_dir)
    kdai._logging.snapshot_modules(out_dir, ["kdai"])
    return out_dir, ver_parts


def main():
    out_dir, ver_parts = start_logging()
    ver_str = "-".join(str(v) for v in ver_parts)

    tr = create_tr()
    train_dir = out_dir / tr.label

    n_epochs = 256
    batch_size = 128
    steps_per_epoch = len(tr.train_ds()) / batch_size
    total_n_steps = n_epochs * steps_per_epoch
    assert total_n_steps == int(total_n_steps)
    lr = 3e-4

    def _scheduler_fn(_optimizer):
        return torch.optim.lr_scheduler.OneCycleLR(
            _optimizer,
            max_lr=lr,
            total_steps=int(total_n_steps),
            three_phase=False,
        )

    kdai.train.train(
        tr,
        n_epochs=n_epochs,
        batch_size=batch_size,
        lr=lr,
        weight_decay=0.01,
        out_dir=train_dir,
        n_workers=10,
        eval_batch_size=2048,
        scheduler_fn=_scheduler_fn,
    )


if __name__ == "__main__":
    main()