# PDC Project 2025 - Peters Stefan, Pelletier Antoine, Jeong Jaeyi, Cohen Adam

import argparse
import numpy as np
import encoder, decoder, channel
import subprocess, pathlib, os, tempfile, sys


def transmit(msg, C):
    return encoder.encode_message(msg, C)

def receive_local(x):
    return channel.channel(x)

def receive_server(x, host, port, input_file=None, output_file=None):
    in_f = input_file or tempfile.NamedTemporaryFile(suffix='.txt', delete=False).name
    np.savetxt(in_f, x)

    out_f = output_file or tempfile.mkstemp(suffix='.txt')[1]

    cmd = [
        sys.executable, str(pathlib.Path(__file__).parent / 'client.py'),
        '--input_file', in_f, '--output_file', out_f,
        '--srv_hostname', host, '--srv_port', str(port)
    ]
    try:
        subprocess.run(cmd, check=True)
        Y = np.loadtxt(out_f)
    finally:
        if not input_file and os.path.exists(in_f):
            os.remove(in_f)
        if not output_file and os.path.exists(out_f):
            os.remove(out_f)
    return Y


def receive(x, mode, host, port, input_file=None, output_file=None):
    if mode == 'local':
        return receive_local(x)
    return receive_server(x, host, port, input_file, output_file)


def test(msg, n_trials, mode, host, port, input_file=None, output_file=None):
    C, _ = encoder.make_codebook()
    print(f"Using codebook with {C.shape[0]} codewords, {C.shape[1]} symbols each.")
    ok = 0
    for _ in range(n_trials):
        x = transmit(msg, C)
        print(f"Transmitted {len(x):,} samples (energy={np.dot(x, x):.2f})")
        y = receive(x, mode, host, port, input_file, output_file)
        print(f"Received {len(y):,} samples (energy={np.dot(y, y):.2f})")
        est, _ = decoder.decode_blocks_with_state(y, C)
        if est == msg:
            ok += 1
    print(f"{ok}/{n_trials} perfect decodes ({100 * ok / n_trials:.2f}%)")


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument('-m', '--message', required=True, help='exactly 40 chars')
    p.add_argument('-n', '--trials', type=int, default=200)
    p.add_argument('--mode', choices=['local', 'server'], default='local')
    p.add_argument('--hostname', default='iscsrv72.epfl.ch')
    p.add_argument('--port', type=int, default=80)
    p.add_argument('-i', '--input_file')
    p.add_argument('-o', '--output_file')
    return p.parse_args()


if __name__ == '__main__':
    a = parse_args()
    test(a.message, a.trials, a.mode, a.hostname, a.port, a.input_file, a.output_file)