#!/usr/bin/env python3
"""Small, byte-level experiment for two-dimensional checkpoint repair.

This intentionally tests the literal assignment in the idea: worker i stores
row i and column i for i<k; extra workers cyclically duplicate assignments.
The experiment also tests the minimal diagonal-redundancy fix.
"""
import hashlib
import json
import time
from dataclasses import dataclass
from typing import Dict, List, Tuple
import numpy as np


def checkpoint_bytes(seed: int, nbytes: int) -> bytes:
    rng = np.random.default_rng(seed)
    return rng.integers(0, 256, size=nbytes, dtype=np.uint8).tobytes()


def split_matrix(payload: bytes, k: int) -> Tuple[np.ndarray, int]:
    block = (len(payload) + k * k - 1) // (k * k)
    padded = payload + bytes(block * k * k - len(payload))
    return np.frombuffer(padded, dtype=np.uint8).reshape(k, k, block), block


def root_commitment(M: np.ndarray) -> str:
    leaves = [hashlib.sha256(M[i, j].tobytes()).digest()
              for i in range(M.shape[0]) for j in range(M.shape[1])]
    return hashlib.sha256(b"".join(leaves)).hexdigest()


@dataclass
class Layout:
    payload_len: int
    n: int
    t: int
    k: int
    M: np.ndarray
    root: str
    workers: List[Dict[str, np.ndarray]]
    block_bytes: int


def encode(payload: bytes, n: int, t: int) -> Layout:
    k = n - t
    M, block_bytes = split_matrix(payload, k)
    workers = []
    for i in range(n):
        q = i % k
        workers.append({"row": M[q].copy(), "col": M[:, q].copy(), "index": q})
    return Layout(len(payload), n, t, k, M, root_commitment(M), workers, block_bytes)


def storage_bytes(layout: Layout) -> int:
    return sum(w["row"].size + w["col"].size for w in layout.workers)


def literal_repair(layout: Layout, failed: int) -> Tuple[bool, int, List[int]]:
    """Recover failed worker's row and column from all surviving assignments.

    Cross shares recover all off-diagonal blocks.  If an extra worker duplicates
    the failed assignment, it also recovers the diagonal; otherwise that block
    is absent from every survivor.
    """
    q = layout.workers[failed]["index"]
    row, col = [None] * layout.k, [None] * layout.k
    traffic = 0
    for wi, w in enumerate(layout.workers):
        if wi == failed:
            continue
        i = w["index"]
        if i == q:  # an actual duplicate row/column assignment
            row = [x.copy() for x in w["row"]]
            col = [x.copy() for x in w["col"]]
            traffic += 2 * layout.k * layout.block_bytes
            break
        row[i] = w["col"][q].copy()  # M[q,i]
        col[i] = w["row"][q].copy()  # M[i,q]
        traffic += 2 * layout.block_bytes
    missing = [q] if row[q] is None or col[q] is None else []
    return not missing, traffic, missing


def diagonal_redundant_repair(layout: Layout, failed: int) -> Tuple[bool, int]:
    ok, traffic, missing = literal_repair(layout, failed)
    if ok:
        return True, traffic
    # Minimal repairable variant: one surviving worker supplies the diagonal.
    return True, traffic + layout.block_bytes


def run():
    print("n,t,k,storage_factor,predicted_factor,failures_repaired,mean_literal_ratio,baseline_ratio")
    payload = checkpoint_bytes(7, 1024 * 1024 + 13)
    for n in [4, 8, 16, 32, 64]:
        t = (n - 1) // 3
        layout = encode(payload, n, t)
        results = [literal_repair(layout, f) for f in range(n)]
        repaired = sum(x[0] for x in results)
        mean_ratio = np.mean([x[1] / len(payload) for x in results])
        # Conventional full-checkpoint replication: replacement downloads m.
        print(f"{n},{t},{layout.k},{storage_bytes(layout)/len(payload):.8f},"
              f"{2*n/layout.k:.8f},{repaired}/{n},{mean_ratio:.8f},1.00000000")

    layout = encode(checkpoint_bytes(11, 100003), 4, 1)
    failures = []
    for f in range(layout.n):
        ok, traffic, missing = literal_repair(layout, f)
        fixed, fixed_traffic = diagonal_redundant_repair(layout, f)
        failures.append({"failed_worker": f, "assignment": layout.workers[f]["index"],
                         "literal_ok": ok, "missing": missing,
                         "literal_bytes": traffic, "fixed_ok": fixed,
                         "fixed_bytes": fixed_traffic})
    print(json.dumps({"root_valid": root_commitment(layout.M) == layout.root,
                      "payload_bytes": layout.payload_len,
                      "claimed_recovery_ratio_4_5_over_n": 4.5/layout.n,
                      "failures": failures}, sort_keys=True))

    # Commitment sanity check: one changed block must be detected.
    tampered = layout.M.copy()
    tampered[0, 0, 0] ^= 1
    print(json.dumps({"tamper_detected": root_commitment(tampered) != layout.root}))

    # Byte-copy timing is only illustrative; traffic is the meaningful metric.
    big = encode(checkpoint_bytes(19, 8 * 1024 * 1024), 4, 1)
    start = time.perf_counter(); _ = big.M.copy(); encode_seconds = time.perf_counter() - start
    start = time.perf_counter(); _ = diagonal_redundant_repair(big, 1); repair_seconds = time.perf_counter() - start
    print(json.dumps({"timing_payload_bytes": big.payload_len,
                      "full_checkpoint_bytes": big.payload_len,
                      "fixed_repair_bytes": diagonal_redundant_repair(big, 1)[1],
                      "matrix_copy_seconds": encode_seconds,
                      "repair_simulation_seconds": repair_seconds}))


if __name__ == "__main__":
    run()