"""CPU-only alignment and no-clobber artifact contracts."""
from __future__ import annotations

import numpy as np
import pytest

from OID_analysis.v2.alignment import exact_time_pair, valid_delayed_pairs
from OID_analysis.v2.artifacts import (
    ArtifactError, atomic_save_npz, validate_field_artifact, validate_timeline,
)


def test_exact_pair_requires_identical_timeline():
    q = np.zeros((3, 2))
    with pytest.raises(ValueError, match="identical lattice_steps"):
        exact_time_pair(q, [1, 2, 3], q, [1, 3, 4])


def test_delay_does_not_pad_or_cross_segments():
    q = np.arange(6)[:, None]; y = (10 + np.arange(6))[:, None]
    pair = valid_delayed_pairs(q, y, np.arange(0, 60, 10), delay_steps=20, segment_ids=[0,0,0,1,1,1])
    assert pair.left[:, 0].tolist() == [0, 3]
    assert pair.right[:, 0].tolist() == [12, 15]
    assert pair.dropped == 4


def test_atomic_field_artifact_refuses_overwrite(tmp_path):
    steps=np.array([3,6,9],dtype=np.int64)
    payload=dict(ux=np.zeros((3,2,4)),uy=np.zeros((3,2,4)),lattice_steps=steps,
                 control_indices=(steps-1)//8,control_offsets=(steps-1)%8+1,control_interval=np.asarray(8))
    path=atomic_save_npz(tmp_path/'fields.npz',**payload)
    with np.load(path,allow_pickle=False) as data:
        assert validate_field_artifact(data)['timeline']['count']==3
    with pytest.raises(FileExistsError): atomic_save_npz(path,**payload)
    with pytest.raises(ArtifactError): validate_field_artifact({'ux':np.zeros((2,2,2)),'uy':np.zeros((2,2,2))})


@pytest.mark.parametrize("steps", [
    np.array([1, 2, 2], dtype=np.int64),
    np.array([3, 2, 1], dtype=np.int64),
    np.array([1.0, 2.0, 3.0]),
    np.array([[1, 2, 3]], dtype=np.int64),
])
def test_delayed_pair_rejects_malformed_lattice_steps(steps):
    q = np.zeros((3, 1))
    with pytest.raises(ValueError, match="lattice_steps"):
        valid_delayed_pairs(q, q, steps, delay_steps=1)


def _timeline(**updates):
    steps = np.array([3, 6, 9], dtype=np.int64)
    payload = {
        "lattice_steps": steps,
        "sample_ids": np.array([10, 11, 12], dtype=np.int64),
        "control_indices": (steps - 1) // 8,
        "control_offsets": (steps - 1) % 8 + 1,
        "intra_control_offsets": (steps - 1) % 8 + 1,
        "control_interval": np.asarray(8, dtype=np.int64),
    }
    payload.update(updates)
    return payload


def test_timeline_accepts_current_integer_schema_and_optional_alias():
    result = validate_timeline(_timeline(), expected_control_interval=8)
    assert result == {"count": 3, "first": 3, "last": 9, "control_interval": 8}


@pytest.mark.parametrize(("key", "value"), [
    ("control_interval", np.array([8], dtype=np.int64)),
    ("control_interval", np.asarray(8.0)),
    ("control_interval", np.asarray(0, dtype=np.int64)),
    ("lattice_steps", np.array([3.0, 6.0, 9.0])),
    ("lattice_steps", np.array([3, 9, 6], dtype=np.int64)),
    ("lattice_steps", np.array([3, 6, 6], dtype=np.int64)),
    ("lattice_steps", np.array([0, 3, 6], dtype=np.int64)),
    ("sample_ids", np.array([10.0, 11.0, 12.0])),
    ("sample_ids", np.array([10, 10, 12], dtype=np.int64)),
    ("control_indices", np.array([0.0, 0.0, 1.0])),
    ("control_offsets", np.array([[3, 6, 1]], dtype=np.int64)),
    ("intra_control_offsets", np.array([3, 6, 2], dtype=np.int64)),
])
def test_timeline_rejects_coercion_and_malformed_arrays(key, value):
    with pytest.raises(ArtifactError):
        validate_timeline(_timeline(**{key: value}))
