"""Write synthetic captures, so tests never need hardware or a recorded file.

Shared by test_capture.py, which uses it to break one guarantee at a time, and
test_compare.py, which uses it to plant a known answer -- a scale factor, a
rotation, a tone at a chosen fraction of the sample rate -- and check that the
analysis recovers it.

Defaults produce a clean, loadable capture; every argument exists to change one
thing about it.
"""

import numpy as np

import capture
import rm3100

CYCLE_COUNT = 100
DT = 1.0 / 250.0          # a deliberately non-nominal true period


def write_capture(path, rows=200, dt=DT, nominal_hz=282.0, header=None,
                  flags=None, drop_header_key=None, extra_lines=(),
                  index_from=0, index_step=1, start_time=1_700_000_000.0,
                  amplitude=1000.0, seed=0, cycle_count=CYCLE_COUNT,
                  counts=None):
    """Write a synthetic capture and return its path.

    `counts` overrides the generated signal with an (rows, 3) array of raw
    counts, which is how a test plants an exact answer. `cycle_count` moves the
    header's gain and is what makes a decimation pair possible: two captures
    whose cycle counts differ by an integer factor.
    """
    meta = {
        "rm3100_capture": 1,
        "nominal_rate_hz": nominal_hz,
        "tmrc_nominal_hz": 600.0,
        "tmrc": "0x92",
        "cycle_count": cycle_count,
        "tesla_per_count": repr(rm3100.tesla_per_count(cycle_count)),
        "i2c_address": "0x23",
        "bus_speed_khz": 750,
        "revid": "0x22",
        "calibrated_period_s": repr(dt),
    }
    meta.update(header or {})
    if drop_header_key:
        meta.pop(drop_header_key, None)

    rng = np.random.default_rng(seed)
    flags = flags or {}
    lines = [f"# {k}: {v}" for k, v in meta.items()]
    lines += list(extra_lines)
    lines.append("sample_index,system_time_unix,x_raw,y_raw,z_raw,warning")
    for i in range(rows):
        index = index_from + i * index_step
        warning = flags.get(i, "")
        if capture.WARN_MISSED in warning:
            x = y = z = 0
        elif counts is not None:
            x, y, z = (int(round(c)) for c in counts[i])
        else:
            x = int(amplitude + rng.normal(0, 3))
            y = int(2 * amplitude + rng.normal(0, 3))
            z = int(-amplitude + rng.normal(0, 3))
        lines.append(f"{index},{start_time + index * dt:.6f},{x},{y},{z},{warning}")
    path.write_text("\n".join(lines) + "\n")
    return str(path)


def field_counts(rows, mean_nt, cycle_count, noise_nt=0.0, seed=0, tones=()):
    """Raw counts for a field of a given mean, noise and planted tones.

    `mean_nt` is an (x, y, z) field in nanotesla; `tones` is a sequence of
    (axis_index, cycles_per_sample, amplitude_nT) added on top. Quantisation to
    integer counts is deliberate -- it is what the real file carries, and a test
    that skipped it would not exercise the dither the analysis relies on.
    """
    lsb = rm3100.tesla_per_count(cycle_count) * rm3100.NT_PER_TESLA
    rng = np.random.default_rng(seed)
    n = np.arange(rows)
    nt = np.tile(np.asarray(mean_nt, dtype=float), (rows, 1))
    if noise_nt:
        nt += rng.normal(0, noise_nt, size=(rows, 3))
    for axis, cycles_per_sample, amplitude in tones:
        nt[:, axis] += amplitude * np.cos(2 * np.pi * cycles_per_sample * n)
    return nt / lsb