"""Generate SSTV, APT and HF fax audio, so the decoders can be tested.

Written from the mode specifications rather than from the decoders, so that a
decoder agreeing with one of these is evidence rather than a tautology.
"""
import numpy as np

from bandsaunter.sstv import (BIT_ONE_HZ, BIT_ZERO_HZ, BLACK_HZ, BREAK_HZ,
                              LEADER_HZ, SYNC_HZ, WHITE_HZ, SSTVMode)


def fm(segments, rate: float) -> np.ndarray:
    """Turn ``(hertz, seconds)`` pairs into one continuous-phase tone."""
    phase = 0.0
    out = []
    for hz, seconds in segments:
        n = max(1, int(round(seconds * rate)))
        freq = np.full(n, float(hz)) if np.isscalar(hz) else np.asarray(hz)
        if freq.size != n:
            freq = np.interp(np.linspace(0, 1, n),
                             np.linspace(0, 1, freq.size), freq)
        step = 2.0 * np.pi * freq / rate
        angles = phase + np.cumsum(step)
        phase = float(angles[-1] % (2.0 * np.pi))
        out.append(np.sin(angles))
    return np.concatenate(out) if out else np.zeros(0)


def _levels(row: np.ndarray) -> np.ndarray:
    """A row of 0-255 as the frequencies that carry it."""
    return BLACK_HZ + np.asarray(row, dtype=np.float64) / 255.0 * (
        WHITE_HZ - BLACK_HZ)


def vis_header(code: int) -> list:
    bits = [(code >> i) & 1 for i in range(7)]
    bits.append(sum(bits) % 2)                 # even parity
    out = [(LEADER_HZ, 0.300), (BREAK_HZ, 0.010), (LEADER_HZ, 0.300),
           (SYNC_HZ, 0.030)]
    out += [(BIT_ONE_HZ if b else BIT_ZERO_HZ, 0.030) for b in bits]
    out.append((SYNC_HZ, 0.030))               # stop bit
    return out


def sstv_audio(mode: SSTVMode, image: np.ndarray, rate: float = 16000.0,
               lines: int | None = None, lead: float = 0.2,
               snr_db: float = 40.0, seed: int = 0) -> np.ndarray:
    """A whole SSTV transmission: header, then the picture, line by line.

    Built from the mode's own segment list, so that a decoder agreeing with
    this is agreeing about the published line structure rather than about a
    second copy of the same guess.
    """
    rng = np.random.default_rng(seed)
    height = lines if lines is not None else min(mode.height, image.shape[0])
    segments = [(0.0, lead)] + vis_header(mode.vis)
    if not mode.sync_first:
        # Scottie sends one sync pulse before the first line, because its
        # per-line pulse comes two thirds of the way through.
        segments.append((SYNC_HZ, mode.sync))
    for y in range(height):
        row = image[y % image.shape[0]].astype(np.float64)
        red, green, blue = row[:, 0], row[:, 1], row[:, 2]
        luma = 0.299 * red + 0.587 * green + 0.114 * blue
        chroma_u = 0.564 * (blue - luma) + 128.0
        chroma_v = 0.713 * (red - luma) + 128.0
        scans = {"R": red, "G": green, "B": blue, "Y": luma,
                 "U": chroma_u, "V": chroma_v,
                 "C": chroma_v if y % 2 == 0 else chroma_u}
        for what, seconds in mode.segments():
            if seconds <= 0:
                continue
            if what == "sync":
                segments.append((SYNC_HZ, seconds))
            elif len(what) == 1:
                segments.append((_levels(scans[what]), seconds))
            else:
                segments.append((BLACK_HZ, seconds))
    audio = fm(segments, rate)
    if snr_db < 60:
        noise = np.sqrt(np.mean(audio ** 2) / (10 ** (snr_db / 10.0)))
        audio = audio + noise * rng.standard_normal(audio.size)
    return audio


def colour_card(width: int = 320, height: int = 256) -> np.ndarray:
    """A picture with structure a decoder can be checked against.

    Colour bars across the top, a grey wedge down the middle and a border, so
    that a picture decoded with the colours swapped, the lines out of order or
    the geometry wrong looks wrong in a way a test can measure.
    """
    image = np.zeros((height, width, 3), dtype=np.uint8)
    bars = [(255, 255, 255), (255, 255, 0), (0, 255, 255), (0, 255, 0),
            (255, 0, 255), (255, 0, 0), (0, 0, 255), (0, 0, 0)]
    band = height // 2
    for i, colour in enumerate(bars):
        lo = i * width // len(bars)
        hi = (i + 1) * width // len(bars)
        image[:band, lo:hi] = colour
    wedge = np.linspace(0, 255, width).astype(np.uint8)
    image[band:, :] = wedge[None, :, None]
    image[:, :2] = 255
    image[:, -2:] = 255
    return image


# ---------------------------------------------------------------------------
# APT
# ---------------------------------------------------------------------------

APT_SYNC_A = "0" + "1100" * 7 + "0" * 8          # 1040 Hz square, 39 words
APT_SYNC_B = "0" + "11100" * 7 + "0" * 3


def apt_audio(image: np.ndarray, rate: float = 16000.0,
              snr_db: float = 40.0, seed: int = 0,
              subcarrier: float = 2400.0) -> np.ndarray:
    """A NOAA APT transmission: two lines a second, 2080 words each.

    ``image`` is greyscale and becomes channel A; channel B is the same
    picture inverted, which is what a real pass looks like when one channel is
    visible and the other infrared.
    """
    rng = np.random.default_rng(seed)
    words_per_line = 2080
    lines = image.shape[0]
    frame = np.zeros((lines, words_per_line), dtype=np.float64)
    for y in range(lines):
        row = np.interp(np.linspace(0, 1, 909),
                        np.linspace(0, 1, image.shape[1]),
                        image[y].astype(np.float64))
        line = np.zeros(words_per_line)
        line[0:39] = np.array([255.0 if c == "1" else 11.0
                               for c in APT_SYNC_A.ljust(39, "0")])
        line[39:86] = 11.0                       # space A
        line[86:995] = row
        line[995:1040] = 128.0                   # telemetry A
        line[1040:1079] = np.array([255.0 if c == "1" else 11.0
                                    for c in APT_SYNC_B.ljust(39, "0")])
        line[1079:1126] = 11.0
        line[1126:2035] = 255.0 - row
        line[2035:2080] = 128.0
        frame[y] = line

    words = frame.reshape(-1)
    word_rate = 4160.0
    n = int(round(words.size / word_rate * rate))
    envelope = np.interp(np.linspace(0, 1, n), np.linspace(0, 1, words.size),
                         words) / 255.0
    t = np.arange(n) / rate
    audio = (0.1 + 0.9 * envelope) * np.sin(2 * np.pi * subcarrier * t)
    if snr_db < 60:
        noise = np.sqrt(np.mean(audio ** 2) / (10 ** (snr_db / 10.0)))
        audio = audio + noise * rng.standard_normal(audio.size)
    return audio


def grey_card(width: int = 909, height: int = 40) -> np.ndarray:
    """A greyscale picture with a hard edge and a ramp, for APT and fax."""
    image = np.zeros((height, width), dtype=np.uint8)
    image[:, :] = np.linspace(0, 255, width).astype(np.uint8)[None, :]
    image[height // 3:2 * height // 3, width // 4:width // 2] = 255
    image[:, ::128] = 0
    return image


# ---------------------------------------------------------------------------
# HF fax
# ---------------------------------------------------------------------------

def fax_audio(image: np.ndarray, rate: float = 16000.0, lpm: float = 120.0,
              ioc: int = 576, black: float = 1500.0, white: float = 2300.0,
              start_seconds: float = 5.0, phasing_lines: int = 20,
              snr_db: float = 40.0, seed: int = 0) -> np.ndarray:
    """A weather fax transmission: start tone, phasing, then the chart."""
    rng = np.random.default_rng(seed)
    pixels = int(round(np.pi * ioc))
    line_seconds = 60.0 / lpm
    segments = [(black, 0.5)]
    # The start tone: 300 Hz for 120 lpm, sent as black/white alternation.
    tone_period = 1.0 / 300.0
    for _ in range(int(start_seconds / tone_period)):
        segments += [(white, tone_period / 2), (black, tone_period / 2)]
    # Phasing: a white pulse at the start of each line, black for the rest.
    for _ in range(phasing_lines):
        segments += [(white, line_seconds * 0.05),
                     (black, line_seconds * 0.95)]
    for y in range(image.shape[0]):
        row = np.interp(np.linspace(0, 1, pixels),
                        np.linspace(0, 1, image.shape[1]),
                        image[y].astype(np.float64))
        segments.append((black + row / 255.0 * (white - black), line_seconds))
    audio = fm(segments, rate)
    if snr_db < 60:
        noise = np.sqrt(np.mean(audio ** 2) / (10 ** (snr_db / 10.0)))
        audio = audio + noise * rng.standard_normal(audio.size)
    return audio