"""One focused check of real V2 crops, persistent workers and checkpoint resume."""

import contextlib
import copy
import io
import json
import signal
import sys
from pathlib import Path
from unittest.mock import patch

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "src"))

import numpy as np
import torch

from vimeml.training import train
from vimeml.training.data import SentenceWindowDataset, make_loader, prepare_indexes, write_json
from vimeml.training.data_v2 import PrefixCropWindowDataset


def snapshots(loader):
    return [{key: value.clone() for key, value in batch.items()} for batch in loader]


def identical(left, right):
    assert len(left) == len(right)
    for a, b in zip(left, right):
        assert a.keys() == b.keys()
        assert all(torch.equal(a[key], b[key]) for key in a)


def fixture_resume(output):
    token_dir, index_dir = output / "fixture-tokens", output / "fixture-indexes"
    token_dir.mkdir()
    sequences = [
        [2, *[4 + i % 60 for i in range(length)], 3] for length in (6, 8, 14, 15, 16, 17, 30, 41)
    ]
    offsets = np.cumsum([0, *map(len, sequences)], dtype=np.uint64)
    splits = {}
    for split in ("train", "validation", "test"):
        np.asarray([token for sequence in sequences for token in sequence], dtype="<u2").tofile(
            token_dir / f"{split}.tokens.bin"
        )
        offsets.astype("<u8").tofile(token_dir / f"{split}.offsets.bin")
        (token_dir / f"{split}.sources.bin").write_bytes(bytes(len(sequences)))
        splits[split] = {
            "sentences": len(sequences),
            "stored_tokens": int(offsets[-1]),
            "prediction_pairs": int(offsets[-1]) - len(sequences),
            "max_sequence_tokens": max(map(len, sequences)),
        }
    write_json(
        token_dir / "manifest.json",
        {
            "status": "complete",
            "format": "vimeml_sentence_tokens_v1",
            "token_dtype": "uint16_le",
            "offset_dtype": "uint64_le",
            "offset_unit": "tokens",
            "vocab_size": 64,
            "special_ids": dict(pad=0, unk=1, bos=2, eos=3),
            "source_ids": {"fineweb": 0},
            "splits": splits,
            "output_sha256": {},
        },
    )
    index = prepare_indexes(token_dir, index_dir, 16, verification="metadata")
    base = {
        "architecture": "tiny_gpt_v2",
        "token_dir": str(token_dir),
        "index_dir": str(index_dir),
        "model": dict(
            vocab_size=64,
            context_length=16,
            d_model=16,
            n_heads=2,
            n_layers=1,
            d_ff=32,
            dropout=0.1,
        ),
        "verification": {"mode": "metadata"},
        "training": dict(
            device="cpu",
            precision="fp32",
            seed=42,
            cpu_threads=2,
            batch_size=4,
            num_workers=0,
            bucket_multiplier=4,
            gradient_accumulation=1,
            epochs=2,
            max_steps=0,
            learning_rate=0.001,
            min_learning_rate=0.0001,
            warmup_steps=1,
            weight_decay=0.1,
            beta1=0.9,
            beta2=0.95,
            grad_clip=1.0,
            log_every=999,
            eval_every=999,
            eval_batches=1,
            checkpoint_every=999,
            benchmark_warmup_steps=0,
            validate_full_at_end=True,
            prefix_crop_probability=0.30,
            prefix_crop_min_remaining_tokens=8,
        ),
    }
    configs = []
    for name in ("reference", "resumed"):
        config = copy.deepcopy(base)
        config.update(output_dir=str(output / name), log_dir=str(output / f"{name}-events"))
        configs.append(config)
    stop_step = (index["splits"]["train"]["windows"] + 3) // 4 + 1
    original_rate = train.learning_rate

    def stop_in_second_epoch(step, settings):
        if step == stop_step:
            signal.getsignal(signal.SIGINT)(signal.SIGINT, None)
        return original_rate(step, settings)

    stream = io.StringIO()
    with contextlib.redirect_stdout(stream):
        reference = train.run(configs[0])
        with patch.object(train, "learning_rate", side_effect=stop_in_second_epoch):
            interrupted = train.run(configs[1])
        assert interrupted["status"] == "interrupted"
        resumed = train.run(configs[1], resume=True)
    (output / "fixture-training.log").write_text(stream.getvalue(), encoding="utf-8")
    left = torch.load(output / "reference/last.pt", weights_only=True)
    right = torch.load(output / "resumed/last.pt", weights_only=True)
    for key in (
        "step",
        "epoch",
        "batch_cursor",
        "total_tokens",
        "total_windows",
        "total_dropped_tokens",
        "total_cropped_windows",
    ):
        assert left[key] == right[key]
    for name in left["model"]:
        torch.testing.assert_close(left["model"][name], right["model"][name], rtol=0, atol=0)
    assert reference["completed_data_passes"] == resumed["completed_data_passes"] == 2.0
    assert resumed["prefix_crop_dropped_tokens"] > 0
    assert resumed["uncropped_prediction_pairs_seen"] == 2 * splits["train"]["prediction_pairs"]
    assert (
        resumed["full_validation"]["last"]["prediction_pairs"]
        == splits["validation"]["prediction_pairs"]
    )
    return {
        "exact_checkpoint_resume": True,
        "dropout": 0.1,
        "epochs": 2,
        "updates": resumed["step"],
        "interrupted_step": interrupted["step"],
        "trained_tokens": resumed["total_trained_tokens"],
        "dropped_tokens": resumed["prefix_crop_dropped_tokens"],
        "full_epoch_accounting": "passed",
        "validation_unaugmented": True,
    }


def main():
    torch.set_num_threads(2)
    output = ROOT / "outputs/model-checks/v2-prefix-crop"
    if output.exists() and any(output.iterdir()):
        raise FileExistsError("Use a new output for a new acceptance run.")
    output.mkdir(parents=True, exist_ok=True)
    token_dir, index_dir = (
        ROOT / "artifacts/token-data/corpus-v2-16k",
        ROOT / "artifacts/training-data/corpus-v2-c128",
    )
    original = SentenceWindowDataset(token_dir, index_dir)
    cropped = PrefixCropWindowDataset(token_dir, index_dir)
    forced = PrefixCropWindowDataset(token_dir, index_dir, probability=1.0)
    worker_loader = None
    try:
        indices = list(range(4096))
        indices.extend(sorted({int(cropped.first[0]), int(cropped.first[0]) + 1, len(cropped) - 1}))
        counts = {}
        for epoch in (0, 1):
            cropped.set_epoch(epoch)
            lengths = cropped.window_lengths(indices)
            eligible = applied = dropped = 0
            offsets = []
            for index, length in zip(indices, lengths):
                before, after = original[index], cropped[(epoch, index)]
                offset = after["crop_offset"]
                assert len(after["input_ids"]) == int(length)
                assert after["input_ids"] == before["input_ids"][offset:]
                assert after["labels"] == before["labels"][offset:]
                assert after == cropped[(epoch, index)]
                if offset:
                    assert before["window_start"] == 0 and len(after["labels"]) >= 8
                    assert after["input_ids"][0] != 2
                if before["window_start"] > 0 or len(before["labels"]) <= 8:
                    assert offset == 0
                if before["window_start"] == 0 and len(before["labels"]) > 8:
                    eligible += 1
                    force = forced[(epoch, index)]
                    assert force["crop_offset"] > 0 and len(force["labels"]) >= 8
                offsets.append(offset)
                applied += offset > 0
                dropped += offset
            counts[str(epoch)] = {
                "samples": len(indices),
                "eligible": eligible,
                "cropped": applied,
                "crop_rate_among_eligible": applied / eligible,
                "dropped_tokens": dropped,
            }
            if epoch == 0:
                first_offsets = offsets
            else:
                assert first_offsets != offsets
        selected = list(range(256))
        worker_loader = make_loader(
            cropped,
            batch_size=32,
            num_workers=2,
            indices=selected,
            shuffle=False,
            bucket_multiplier=4,
        )
        for epoch in (0, 1):
            worker_loader.batch_sampler.set_epoch(epoch)
            actual = snapshots(worker_loader)
            reference_loader = make_loader(
                cropped,
                batch_size=32,
                num_workers=0,
                indices=selected,
                shuffle=False,
                bucket_multiplier=4,
            )
            reference_loader.batch_sampler.set_epoch(epoch)
            reference = snapshots(reference_loader)
            identical(actual, reference)
            resume_loader = make_loader(
                cropped,
                batch_size=32,
                num_workers=0,
                indices=selected,
                shuffle=False,
                bucket_multiplier=4,
                start_batch=3,
            )
            resume_loader.batch_sampler.set_epoch(epoch, start_batch=3)
            identical(actual[3:], snapshots(resume_loader))
            del reference_loader, resume_loader
        for split in ("validation", "test"):
            unchanged = SentenceWindowDataset(token_dir, index_dir, split)
            try:
                sample = unchanged[0]
                assert "crop_offset" not in sample
            finally:
                unchanged.close()
        report = {
            "status": "passed",
            "probability": 0.30,
            "min_remaining_tokens": 8,
            "real_corpus_samples": counts,
            "scalar_vector_lengths_match": True,
            "same_epoch_reproducible": True,
            "different_epochs_change_crops": True,
            "persistent_worker_epochs": [0, 1],
            "worker_counts_compared": [0, 2],
            "prefetched_resume_batches_exact": True,
            "short_and_continuation_windows_unchanged": True,
            "validation_test_unchanged": True,
            "fixture_training": fixture_resume(output),
            "formal_training_started": False,
        }
        write_json(output / "report.json", report)
        print(json.dumps(report, ensure_ascii=False, indent=2))
    finally:
        del worker_loader
        original.close()
        cropped.close()
        forced.close()


if __name__ == "__main__":
    main()