#!/usr/bin/env python3
"""Compare channel-major vs time-major layout for torch.compile scalar x slice."""
import os

os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
import torch  # noqa: E402

torch.set_num_threads(1)
import time  # noqa: E402

B, T, klen, half_length, factor = 1000, 1040, 17, 8, 2
n_win = T - 2 * half_length


# ── scalar x slice inner (channel-major) ──
def _inner_cm(data, kv, half_length, n_win):
    klen = 2 * half_length + 1
    acc = kv[0] * data[:, 0:n_win]
    for k in range(1, klen):
        acc = acc + kv[k] * data[:, k : k + n_win]
    return acc


_compiled_cm = torch.compile(_inner_cm)


# ── scalar x slice inner (time-major) ──
def _inner_tm(data, kv, half_length, n_win):
    klen = 2 * half_length + 1
    acc = kv[0] * data[0:n_win, :]
    for k in range(1, klen):
        acc = acc + kv[k] * data[k : k + n_win, :]
    return acc


_compiled_tm = torch.compile(_inner_tm)


# ── full functions ──
def resample_cm(data_cm, kernel_t):
    output = torch.empty(B, n_win * factor, dtype=data_cm.dtype)
    output[:, 0::factor] = data_cm[:, half_length : half_length + n_win]
    for f in range(1, factor):
        output[:, f::factor] = _compiled_cm(data_cm, kernel_t[f], half_length, n_win)
    return output  # (B, T_out)


def resample_tm(data_tm, kernel_t):
    output = torch.empty(n_win * factor, B, dtype=data_tm.dtype)
    output[0::factor, :] = data_tm[half_length : half_length + n_win, :]
    for f in range(1, factor):
        output[f::factor, :] = _compiled_tm(data_tm, kernel_t[f], half_length, n_win)
    return output  # (T_out, B) — stays time-major


def resample_tm_to_cm(data_tm, kernel_t):
    output = resample_tm(data_tm, kernel_t)
    return output.T.contiguous()  # convert back to (B, T_out)


# ── data ──
kernel_t = torch.randn(factor, klen)
data_cm = torch.randn(B, T)
data_tm = data_cm.T.contiguous()  # (T, B)

N = 200
W = 20

# warmup both compiled functions
for _ in range(W):
    resample_cm(data_cm, kernel_t)
    resample_tm(data_tm, kernel_t)
    resample_tm_to_cm(data_tm, kernel_t)

# benchmark
t0 = time.perf_counter()
for _ in range(N):
    resample_cm(data_cm, kernel_t)
t1 = time.perf_counter()
print(f"channel-major (B,T) in, (B,T_out) out:         {(t1-t0)/N*1e6:.1f} us")

t0 = time.perf_counter()
for _ in range(N):
    resample_tm(data_tm, kernel_t)
t1 = time.perf_counter()
print(f"time-major (T,B) in, (T_out,B) out:            {(t1-t0)/N*1e6:.1f} us")

t0 = time.perf_counter()
for _ in range(N):
    resample_tm_to_cm(data_tm, kernel_t)
t1 = time.perf_counter()
print(f"time-major (T,B) in, transpose to (B,T_out):   {(t1-t0)/N*1e6:.1f} us")

# verify correctness
ref = resample_cm(data_cm, kernel_t)
out_tm = resample_tm(data_tm, kernel_t)
print(f"\nmatch: {torch.allclose(ref, out_tm.T, atol=1e-5)}")

# time just transposes
t0 = time.perf_counter()
for _i in range(N):
    data_cm.T.contiguous()
t1 = time.perf_counter()
print(f"\ntranspose (B,T) to (T,B) alone:               {(t1-t0)/N*1e6:.1f} us")

t0 = time.perf_counter()
for _i in range(N):
    out = torch.empty(n_win * factor, B)
    out.T.contiguous()
t1 = time.perf_counter()
print(f"transpose (T_out,B) to (B,T_out) alone:        {(t1-t0)/N*1e6:.1f} us")
