"""CPU contract for an oracle-free learned flow-latent encoder.

The test decoder is intentionally tiny. It maps required continuous Music3
flow shape [B,128,86] to [B,2,44032], never RVQ/native tokens.
"""
from __future__ import annotations

from dataclasses import asdict
from pathlib import Path

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

from music3lab.codec import learned_flow_encoder as flow

SAMPLES = 44_032
SEEDS = (101, 103, 107, 109)


class TinyFrozenFlowDecoder(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.gain = nn.Parameter(torch.tensor([1.25, 0.75]), requires_grad=False)

    def forward(self, z: torch.Tensor) -> torch.Tensor:
        assert z.shape[1:] == (128, 86)
        audio = F.interpolate(z[:, :2], size=SAMPLES, mode="linear", align_corners=False)
        return audio * self.gain.view(1, 2, 1)


def _latents(count: int = 12) -> torch.Tensor:
    time = torch.linspace(-1.0, 1.0, 86)
    values = []
    for i in range(count):
        z = torch.zeros(128, 86)
        z[0] = (i + 1) / count * time
        z[1] = .25 * torch.cos((i + 1) * torch.pi * time)
        values.append(z)
    return torch.stack(values)


def _examples(count: int = 12) -> tuple[flow.AudioExample, ...]:
    with torch.no_grad():
        audio = TinyFrozenFlowDecoder()(_latents(count))
    return tuple(flow.AudioExample(flow.deterministic_audio_id(a), a.clone()) for a in audio)


def _config() -> flow.LearnedFlowEncoderConfig:
    return flow.LearnedFlowEncoderConfig(
        latent_channels=128, latent_frames=86, audio_channels=2, audio_samples=SAMPLES,
        hidden_channels=8, learning_rate=.03, batch_size=3, seeds=SEEDS,
    )


def test_exact_geometry_frozen_decoder_and_encoder_gradients() -> None:
    encoder, decoder = flow.LearnedFlowEncoder(_config()), TinyFrozenFlowDecoder()
    audio = torch.randn(2, 2, SAMPLES)
    z = encoder(audio)
    assert z.shape == (2, 128, 86)
    flow.audio_reconstruction_loss(decoder(z), audio).backward()
    assert all(p.grad is None for p in decoder.parameters())
    assert any(p.grad is not None and p.grad.abs().sum() > 0 for p in encoder.parameters())


def test_deterministic_ids_and_four_seed_splits_are_complete_disjoint_and_leak_free() -> None:
    examples = _examples()
    changed = examples[0].audio.clone()
    changed[0, 0] += 1e-4
    assert flow.deterministic_audio_id(changed) != examples[0].sample_id
    first = flow.build_four_seed_splits(examples, seeds=SEEDS, validation_fraction=.25)
    assert first == flow.build_four_seed_splits(examples, seeds=SEEDS, validation_fraction=.25)
    assert tuple(first) == SEEDS
    all_ids = {x.sample_id for x in examples}
    for seed, split in first.items():
        assert seed in SEEDS and split.train_ids and split.validation_ids
        assert set(split.train_ids).isdisjoint(split.validation_ids)
        assert set(split.train_ids) | set(split.validation_ids) == all_ids
        assert split.train_ids == tuple(sorted(split.train_ids))
        assert split.validation_ids == tuple(sorted(split.validation_ids))
    with pytest.raises((TypeError, ValueError), match="latent|unknown|field"):
        flow.AudioExample(examples[0].sample_id, examples[0].audio, latent_target=torch.zeros(128, 86))


def test_loss_identity_order_and_heldout_metric_are_audio_only() -> None:
    decoder, target = TinyFrozenFlowDecoder(), _examples(2)[0].audio.unsqueeze(0)
    exact = flow.audio_reconstruction_loss(target, target)
    distorted = flow.audio_reconstruction_loss(target + .1, target)
    inverted = flow.audio_reconstruction_loss(-target, target)
    assert exact.item() == pytest.approx(0.0, abs=1e-12)
    assert distorted > exact and inverted > distorted
    metric = flow.evaluate_heldout(flow.LearnedFlowEncoder(_config()), decoder, _examples(4), torch.zeros(1, 128, 86))
    assert set(metric) == {"encoder_loss", "prior_mean_loss", "improvement_fraction", "sample_count"}
    assert metric["sample_count"] == 4 and metric["encoder_loss"] >= 0 and metric["prior_mean_loss"] > 0
    assert metric["improvement_fraction"] == pytest.approx(1 - metric["encoder_loss"] / metric["prior_mean_loss"])
    assert "token" not in " ".join(metric).lower()


def test_short_cpu_training_reduces_heldout_loss_below_prior_mean() -> None:
    torch.manual_seed(7)
    examples, config, decoder = _examples(), _config(), TinyFrozenFlowDecoder()
    split = flow.build_four_seed_splits(examples, seeds=SEEDS, validation_fraction=.25)[101]
    encoder = flow.LearnedFlowEncoder(config)
    before = flow.evaluate_heldout(encoder, decoder, split.validation_examples(examples), torch.zeros(1, 128, 86))
    run = flow.train_encoder(encoder, decoder, split.training_examples(examples), split.validation_examples(examples), config, steps=36)
    after = flow.evaluate_heldout(encoder, decoder, split.validation_examples(examples), torch.zeros(1, 128, 86))
    assert run.training_losses[0] > run.training_losses[-1]
    assert after["encoder_loss"] < before["encoder_loss"]
    assert after["encoder_loss"] < after["prior_mean_loss"] and after["improvement_fraction"] > 0


def test_optional_refinement_is_bounded_and_never_claims_native_tokens() -> None:
    decoder, target = TinyFrozenFlowDecoder(), _examples(1)[0].audio.unsqueeze(0)
    initial = torch.zeros(1, 128, 86)
    result = flow.optional_refine_latents(decoder, initial, target, steps=4, learning_rate=.1, max_steps=4)
    assert result.latents.shape == (1, 128, 86)
    assert 0 <= result.steps_used <= 4 and result.final_loss <= result.initial_loss
    assert result.kind == "continuous_flow_latent" and "token" not in result.kind
    with pytest.raises((TypeError, ValueError), match="max_steps|bound"):
        flow.optional_refine_latents(decoder, initial, target, steps=5, learning_rate=.1, max_steps=4)


def test_checkpoint_contains_only_encoder_config_and_provenance(tmp_path: Path) -> None:
    config, encoder = _config(), flow.LearnedFlowEncoder(_config())
    path = tmp_path / "learned-flow-encoder.pt"
    provenance = {"base_commit": "9606448", "dataset_ids": tuple(x.sample_id for x in _examples()), "split_seeds": SEEDS, "decoder_identity": "frozen-test-decoder"}
    flow.save_encoder_checkpoint(path, encoder, config, provenance)
    payload = torch.load(path, map_location="cpu", weights_only=False)
    assert set(payload) == {"schema_version", "encoder_state_dict", "config", "provenance"}
    assert payload["config"] == asdict(config) and payload["provenance"] == provenance
    serialized = repr(payload).lower()
    for word in ("decoder_state", "latent_target", "native", "rvq"):
        assert word not in serialized