use hurray_core::layout::LayoutDescriptor;

let rm = LayoutDescriptor::RowMajor;
let cm = LayoutDescriptor::ColMajor;
assert_eq!(rm.tag(), 0x01);
assert_eq!(cm.tag(), 0x02);

import hurray

rm = hurray.RowMajorLayout()
cm = hurray.ColMajorLayout()
assert rm.tag == 0x01
assert cm.tag == 0x02

use hurray_core::layout::{LayoutDescriptor, StridedLayout};

// Row-major strides for a 3×4 tensor: last dim varies fastest.
let rm_strides = LayoutDescriptor::Strided(StridedLayout::new(vec![4, 1]));

// Same tensor with first dimension reversed.
let reversed = LayoutDescriptor::Strided(StridedLayout::new(vec![-4, 1]));

// Broadcast along dimension 0: all rows map to row 0.
let broadcast = LayoutDescriptor::Strided(StridedLayout::new(vec![0, 1]));

import hurray

# Row-major strides for a 3x4 tensor: last dim varies fastest.
rm_strides = hurray.StridedLayout([4, 1])

# Same tensor with first dimension reversed.
reversed_ = hurray.StridedLayout([-4, 1])

# Broadcast along dimension 0: all rows map to row 0.
broadcast = hurray.StridedLayout([0, 1])

assert rm_strides.strides == (4, 1)

use hurray_core::layout::{LayoutDescriptor, TiledLayout};

let tiled = LayoutDescriptor::Tiled(Box::new(
    TiledLayout::new(
        vec![2, 4], // tile_shape
        0x01,       // outer_layout: row-major
        0x02,       // inner_layout: column-major
        None,       // outer_strides: None (implicit for row-major outer)
        None,       // inner_strides: None
        None,       // inner_tiled: None (not recursive)
    ).unwrap(),
));

import hurray

tiled = hurray.TiledLayout(
    [2, 4],                    # tile_shape
    outer_layout="row_major",
    inner_layout="col_major",
)

use hurray_core::layout::{LayoutDescriptor, OuterStrides, TiledLayout};

let tiled_strided = LayoutDescriptor::Tiled(Box::new(
    TiledLayout::new(
        vec![2, 2],
        0x03, // strided outer
        0x01, // row-major inner
        Some(OuterStrides::new(vec![2, 1])), // tile-grid strides in units of tiles
        None,
        None,
    ).unwrap(),
));

import hurray

tiled_strided = hurray.TiledLayout(
    [2, 2],
    outer_layout="strided",            # tile-grid strides required
    inner_layout="row_major",
    outer_strides=[2, 1],              # in units of tiles
)

use hurray_core::layout::TiledLayout;

let inner = TiledLayout::new(vec![4, 4], 0x01, 0x01, None, None, None).unwrap();
let outer = TiledLayout::new(
    vec![32, 32],
    0x01,
    0x04, // inner_layout is itself tiled
    None,
    None,
    Some(Box::new(inner)),
).unwrap();

import hurray

inner = hurray.TiledLayout([4, 4], "row_major", "row_major")
outer = hurray.TiledLayout(
    [32, 32],
    outer_layout="row_major",
    inner_layout="tiled",              # the inner layout is itself tiled
    inner_tiled=inner,
)

use hurray_core::layout::{CooLayout, LayoutDescriptor};

let coo = LayoutDescriptor::Coo(CooLayout::new(
    42,   // nnz
    true, // is_sorted: non-zeros in lexicographic order
));
assert_eq!(coo.buffer_count().map(|n| n.get()), Some(2));

import hurray

coo = hurray.CooLayout(
    nnz=42,
    is_sorted=True,     # non-zeros in lexicographic order
)
assert coo.buffer_count == 2

use hurray_core::layout::{CsrLayout, LayoutDescriptor};

let csr = LayoutDescriptor::Csr(CsrLayout::new(100)); // nnz = 100
assert_eq!(csr.buffer_count().map(|n| n.get()), Some(3));

import hurray

csr = hurray.CsrLayout(nnz=100)
assert csr.buffer_count == 3

use hurray_core::layout::{CscLayout, LayoutDescriptor};

let csc = LayoutDescriptor::Csc(CscLayout::new(100));
assert_eq!(csc.buffer_count().map(|n| n.get()), Some(3));

import hurray

csc = hurray.CscLayout(nnz=100)
assert csc.buffer_count == 3

use hurray_core::layout::{CsfLayout, LayoutDescriptor};
use hurray_core::Shape;

// Rank-3 sparse tensor, identity mode order, 4 non-zeros.
let csf = LayoutDescriptor::Csf(CsfLayout::new(4, vec![0, 1, 2]));
assert_eq!(csf.tag(), 0x09);
assert_eq!(csf.buffer_count().map(|n| n.get()), Some(7)); // 2*3 + 1

// rank ≥ 3 only; CSR/CSC own rank-2.
let shape = Shape::new(vec![2, 3, 4]).unwrap();
assert!(csf.validate_against_shape(&shape).is_ok());
assert!(csf
    .validate_against_shape(&Shape::new(vec![3, 4]).unwrap())
    .is_err());

import hurray

# Rank-3 sparse tensor, identity mode order, 4 non-zeros.
csf = hurray.CsfLayout(nnz=4, mode_order=[0, 1, 2])
assert csf.tag == 0x09
assert csf.buffer_count == 7        # 2*3 + 1

# rank >= 3 only; CSR/CSC own rank-2.
csf.validate_against_shape([2, 3, 4])
try:
    csf.validate_against_shape([3, 4])
    raise AssertionError("rank-2 should be refused")
except hurray.InvalidDescriptorError:
    pass

use hurray_core::layout::{LayoutDescriptor, MortonLayout};
use hurray_core::Shape;

// 4×4 tensor: each dim needs 2 bits (4 <= 2^2).
let morton = LayoutDescriptor::Morton(MortonLayout::new(vec![2, 2]).unwrap());
let shape = Shape::new(vec![4, 4]).unwrap();
morton.validate_against_shape(&shape).unwrap();

import hurray

# 4x4 tensor: each dim needs 2 bits (4 <= 2^2).
morton = hurray.MortonLayout([2, 2])
morton.validate_against_shape([4, 4])

use hurray_core::layout::{HilbertLayout, LayoutDescriptor};
use hurray_core::Shape;

// 8×8×8 tensor: order=3 (8 = 2^3), rank=3.
let hilbert = LayoutDescriptor::Hilbert(HilbertLayout::new(3, 3).unwrap());
let shape = Shape::new(vec![8, 8, 8]).unwrap();
hilbert.validate_against_shape(&shape).unwrap();

import hurray

# 8x8x8 tensor: order=3 (8 = 2^3), rank=3.
hilbert = hurray.HilbertLayout(hilbert_order=3, hilbert_rank=3)
hilbert.validate_against_shape([8, 8, 8])

use hurray_core::layout::{
    validate_layout_tag_strict, is_invalid_tag, is_named_tag, is_reserved_tag,
    is_private_tag, LayoutDescriptor, UnknownLayout,
};
use hurray_core::Error;

// Check individual tag categories without constructing a descriptor.
// 0x10 is a genuinely unassigned tag in the Tier-1 reserved range.
assert!(is_invalid_tag(0x00));
assert!(is_named_tag(0x07));      // CSR — this crate knows how to check it
assert!(is_reserved_tag(0x10));
assert!(is_private_tag(0xF3));

// Strict-mode validation: rejects invalid, reserved, and private tags.
assert!(validate_layout_tag_strict(0x01).is_ok());
assert!(matches!(validate_layout_tag_strict(0x00), Err(Error::InvalidLayoutTag(0x00))));
assert!(matches!(validate_layout_tag_strict(0x10), Err(Error::ReservedLayoutTag(0x10))));
assert!(matches!(validate_layout_tag_strict(0xF0), Err(Error::PrivateLayoutTag(0xF0))));

// Permissive mode: wrap unrecognised tags in Unknown for passthrough.
// The reader must NOT dereference the tensor data buffer for Unknown layouts.
let unknown = LayoutDescriptor::Unknown(UnknownLayout::new(0x10, vec![]).unwrap());
assert_eq!(unknown.tag(), 0x10);
assert!(unknown.buffer_count().is_none());

// Only genuinely unrecognised tags: "unknown" is a claim, and it has to be true.
// A named tag wrapped this way would skip every check its own variant applies
// while still encoding to that tag on the wire.
assert!(matches!(UnknownLayout::new(0x07, vec![]), Err(Error::NamedLayoutTag(0x07))));
assert!(matches!(UnknownLayout::new(0xF0, vec![]), Err(Error::PrivateLayoutTag(0xF0))));

import hurray

# Classify a tag without constructing a descriptor. The four categories partition
# the byte space, so one call answers the question four predicates would.
assert hurray.layout_tag_kind(0x07) == "named"      # CSR
assert hurray.layout_tag_kind(0x10) == "reserved"
assert hurray.layout_tag_kind(0xF3) == "private"
assert hurray.layout_tag_kind(0x00) == "invalid"

# Permissive mode: wrap an unrecognised tag for passthrough.
# The reader must NOT dereference the tensor data buffer for an Unknown layout.
unknown = hurray.UnknownLayout(0x10, b"")
assert unknown.tag == 0x10
assert unknown.buffer_count is None

# Only genuinely unrecognised tags: "unknown" is a claim, and it has to be true.
# A named tag wrapped this way would skip every check its own class applies while
# still encoding to that tag on the wire.
for taken in (0x07, 0xF0):
    try:
        hurray.UnknownLayout(taken, b"")
        raise AssertionError(f"0x{taken:02X} is not unknown")
    except ValueError as exc:
        print(exc)

use hurray_core::layout::{CsrLayout, LayoutDescriptor};
use hurray_core::Shape;

let csr = LayoutDescriptor::Csr(CsrLayout::new(5));

// Rank-2: valid.
assert!(csr.validate_against_shape(&Shape::new(vec![4, 5]).unwrap()).is_ok());

// Rank-3: rejected — CSR is only defined for rank-2 tensors.
assert!(csr.validate_against_shape(&Shape::new(vec![2, 3, 4]).unwrap()).is_err());

import hurray

csr = hurray.CsrLayout(nnz=5)

# Rank-2: valid.
csr.validate_against_shape([4, 5])

# Rank-3: rejected - CSR is only defined for rank-2 tensors.
try:
    csr.validate_against_shape([2, 3, 4])
    raise AssertionError("rank-3 should be refused")
except hurray.InvalidDescriptorError:
    pass

use hurray_core::layout::{LayoutDescriptor, PrivateExtensionLayout};

let private = LayoutDescriptor::PrivateExtension(
    PrivateExtensionLayout::new(
        0xF0,                    // tag: must be 0xF0–0xFE
        0xDEAD_BEEF_0000_0001,   // implementation-defined layout ID
        vec![0x01, 0x00, 0x04],  // opaque metadata
    ).unwrap(),
);
// buffer_count is None: the format doesn't know how many buffers this needs.
assert!(private.buffer_count().is_none());

import hurray

private = hurray.PrivateExtensionLayout(
    0xF0,                       # tag: must be 0xF0-0xFE
    0xDEAD_BEEF_0000_0001,      # implementation-defined layout ID
    b"\x01\x00\x04",            # opaque metadata
)
# buffer_count is None: the format doesn't know how many buffers this needs.
assert private.buffer_count is None