# Troca as amostras de instrumento dos modulos Impulse Tracker (.it) do Solar Guard por ondas sinteticas nossas.
# Os .it do jogo foram feitos a partir de MIDI no OpenMPT, e as amostras vem da biblioteca do OpenMPT e do GM.DLS
# do Windows, de licenca desconhecida (README do jogo: pasta "Original Soundtrack"). As notas, o andamento e a
# estrutura da musica (do autor, GPL-3.0) ficam como estao: so os dados de amostra sao regerados aqui, com
# formas de onda simples (serra, pulso, triangulo, ruido) escolhidas pelo nome do instrumento.
#
# Uso: python3 -I replace-samples.py <pasta com os .it> (reescreve os arquivos no lugar)
import math
import struct
import sys
from pathlib import Path

C5 = 523.25
RATE = 16744  # 32 amostras por ciclo no C-5 (523.25 Hz x 32)
P = 32


def wave(kind, ph):
    """Forma de onda de periodo 1, de -1 a 1."""
    ph -= math.floor(ph)
    if kind == "saw":
        return 2.0 * ph - 1.0
    if kind == "pulse":
        return 1.0 if ph < 0.25 else -1.0
    if kind == "square":
        return 1.0 if ph < 0.5 else -1.0
    if kind == "tri":
        return 4.0 * abs(ph - 0.5) - 1.0
    # sine
    return math.sin(2 * math.pi * ph)


def kind_for(name):
    n = name.lower()
    if any(k in n for k in ("snare", "tambourine", "standard")):
        return "drum"
    if any(k in n for k in ("pulse", "gate")):
        return "pulse"
    if any(k in n for k in ("saw", "lead", "rez", "pop")):
        return "saw"
    if any(k in n for k in ("pad", "atmosphere", "regen")):
        return "tri"
    return "sine"  # pianos, flauta


def lfsr_noise(n, seed=0xACE1):
    out = []
    x = seed
    for _ in range(n):
        bit = ((x >> 0) ^ (x >> 2) ^ (x >> 3) ^ (x >> 5)) & 1
        x = (x >> 1) | (bit << 15)
        out.append((x & 0xFF) / 127.5 - 1.0)
    return out


GAIN = {"drum": 0.7}


def make_sample(name, looped, orig_len, orig_c5):
    kind = kind_for(name)
    if kind == "drum":
        n = max(600, min(orig_len, int(orig_c5 * 0.16)))
        rate = orig_c5
        noise = lfsr_noise(n)
        tau = n / 5.0
        if "standard" in name.lower():  # tom: seno que desce
            data = [math.sin(2 * math.pi * (140 * t / rate - 25 * (t / rate) ** 2 * 3)) * math.exp(-t / tau) for t in range(n)]
        else:
            data = [noise[t] * math.exp(-t / tau) for t in range(n)]
        return data, rate, False
    if looped:
        data = [wave(kind, i / P) for i in range(P * 2)]
        return data, RATE, True
    n = int(RATE * 1.2)
    tau = RATE * 0.45
    data = [wave(kind, i / P) * math.exp(-i / tau) for i in range(n)]
    return data, RATE, False


def process(path):
    d = bytearray(path.read_bytes())
    assert d[:4] == b"IMPM"
    ordn, insn, smpn, patn = struct.unpack("<4H", d[0x20:0x28])
    base = 0xC0 + ordn
    ins_offs = [struct.unpack("<I", d[base + 4 * i:base + 4 * i + 4])[0] for i in range(insn)]
    smp_offs = [struct.unpack("<I", d[base + 4 * insn + 4 * i:base + 4 * insn + 4 * i + 4])[0] for i in range(smpn)]
    pat_offs = [struct.unpack("<I", d[base + 4 * (insn + smpn) + 4 * i:base + 4 * (insn + smpn) + 4 * i + 4])[0] for i in range(patn)]
    msglen, msgoff = struct.unpack("<HI", d[0x36:0x3C])
    data_ptrs = []
    for off in smp_offs:
        ln = struct.unpack("<I", d[off + 0x30:off + 0x34])[0]
        ptr = struct.unpack("<I", d[off + 0x48:off + 0x4C])[0]
        if ln and ptr:
            data_ptrs.append(ptr)
    chunk_starts = sorted({o for o in ins_offs + smp_offs + pat_offs + data_ptrs if o} | ({msgoff} if msglen else set()))
    # zera os dados antigos das amostras (ate o proximo bloco do arquivo)
    for ptr in data_ptrs:
        nxt = next((c for c in chunk_starts if c > ptr), len(d))
        d[ptr:nxt] = bytes(nxt - ptr)
    # gera as novas, no fim do arquivo
    for off in smp_offs:
        ln, lb, le, c5 = struct.unpack("<4I", d[off + 0x30:off + 0x40])
        if not ln:
            continue
        name = d[off + 0x14:off + 0x2E].split(b"\0")[0].decode("latin1")
        flags = d[off + 0x12]
        looped = bool(flags & 0x10) and le > lb
        data, rate, loop = make_sample(name, looped, ln, c5)
        gain = GAIN.get(kind_for(name), 0.55)
        pcm = bytes(((int(round(max(-1.0, min(1.0, v * gain)) * 127)) & 0xFF) for v in data))
        ptr = len(d)
        d += pcm
        new_flags = 0x01 | (0x10 if loop else 0)  # amostra associada, 8 bits, mono, sem compressao
        d[off + 0x12] = new_flags
        d[off + 0x2E] = 0x01  # com sinal
        struct.pack_into("<IIII", d, off + 0x30, len(data), 0, len(data) if loop else 0, rate)
        struct.pack_into("<II", d, off + 0x40, 0, 0)  # sem loop de sustain
        struct.pack_into("<I", d, off + 0x48, ptr)
    path.write_bytes(bytes(d))


if __name__ == "__main__":
    for f in sorted(Path(sys.argv[1]).glob("*.it")):
        process(f)
        print(f.name, f.stat().st_size)