"""Static-only frozen encoder extraction and exact request-schedule generation."""
import ast
import hashlib
from pathlib import Path
from typing import Any
import numpy as np


def frozen_pn_drive(path, expected_sha):
    """Execute the exact one pure function, not unrelated historical experiments."""
    import pandas as pd
    source = Path(path).read_bytes()
    assert hashlib.sha256(source).hexdigest() == expected_sha
    tree = ast.parse(source, filename=str(path))
    nodes = [node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == 'pn_drive']
    assert len(nodes) == 1 and not nodes[0].decorator_list
    module = ast.Module(body=nodes, type_ignores=[])
    namespace = {'np': np, 'pd': pd, 'Any': Any}
    exec(compile(module, str(path), 'exec'), namespace)
    return namespace['pn_drive'], hashlib.sha256(ast.dump(nodes[0], include_attributes=False).encode()).hexdigest()


def requests_from_forced(forced, targets):
    targets = np.asarray(targets)
    assert targets.dtype == np.int32 and targets.ndim == 1 and len(np.unique(targets)) == len(targets)
    assert forced.dtype == np.bool_ and forced.ndim == 2 and forced.shape[1] == len(targets)
    ticks, columns = np.nonzero(forced)
    requests = np.column_stack([ticks, targets[columns]]).astype(np.int32)
    if len(requests):
        requests = requests[np.lexsort((requests[:, 1], requests[:, 0]))]
        assert np.all(np.diff(requests[:, 0].astype(np.int64) * 166700 + requests[:, 1]) > 0)
    return requests


def make_requests(targets, rates, seed, steps=1000):
    import torch
    assert rates.dtype == np.float32 and rates.shape == targets.shape
    assert np.isfinite(rates).all() and np.all(rates >= 0) and steps == 1000
    original = rates * .1 / 1000.0
    explicit = rates * np.float32(.1) / np.float32(1000.)
    assert original.dtype == np.float32 and np.array_equal(original.view(np.uint32), explicit.view(np.uint32))
    assert np.all(original <= 1)
    generator = torch.Generator(device='cuda')
    generator.manual_seed(int(seed))
    probabilities = torch.as_tensor(original, dtype=torch.float32, device='cuda')
    forced = (torch.rand((steps, len(targets)), generator=generator, device='cuda') < probabilities[None, :]).cpu().numpy()
    return requests_from_forced(forced, targets)


def write_npz(path, **arrays):
    with Path(path).open('xb') as stream:
        np.savez_compressed(stream, **arrays)