# default_transmission.py
import numpy as np, random

class DefaultTransmission:
    def __init__(self):
        self.n_channels = 8
        self.original_erasure_probability = 0.5

    def simulate_bit_recovery(self, n_bits=4, trials=100):
        """
        模拟 8 信道传输 n_bits 比特，每比特重复 trials 次
        返回每个比特的成功恢复率
        """
        success = np.zeros((self.n_channels, n_bits))

        for trial in range(trials):
            # 随机生成 4 比特信号
            bits = np.random.randint(0, 2, size=n_bits)

            for ch in range(self.n_channels):
                for i, b in enumerate(bits):
                    erased = (np.random.rand() < self.original_erasure_probability)
                    received = None if erased else b
                    if received == b:
                        success[ch, i] += 1

        avg_success = success / trials
        print("\n每信道比特成功率矩阵 (行=信道, 列=比特索引):")
        for ch in range(self.n_channels):
            rates = " ".join(f"{avg_success[ch,i]:.3f}" for i in range(n_bits))
            print(f"信道 {ch+1}: {rates}")

        bit_avg = np.mean(avg_success, axis=0)
        print("\n每比特平均恢复率:")
        for i, r in enumerate(bit_avg):
            print(f"Bit {i}: {r:.3f}")

        return avg_success

# polar_transmission.py
import numpy as np

# --- BEC 信道 ---
def transmit_BEC(x, eps):
    y = np.array(x, dtype=object)
    erasures = np.random.rand(len(x)) < eps
    y[erasures] = None
    return y

# --- 极化挑战 ---
def polar_bit_challenge(N=8, K=4, eps=0.5, trials=100, polar_funcs=None):
    """
    polar_funcs: dict, 必须包含
        'construction', 'encode', 'decode'
    """
    if polar_funcs is None:
        raise ValueError("必须传入 polar_funcs 字典")

    construction = polar_funcs['construction']
    encode = polar_funcs['encode']
    decode = polar_funcs['decode']

    info_idx, frozen_idx, _ = construction(N, K, eps)
    per_bit_success = np.zeros(K, dtype=int)

    for _ in range(trials):
        u = np.zeros(N, dtype=int)
        info_bits = np.random.randint(0, 2, K)
        u[info_idx] = info_bits

        x = encode(u, N)
        y = transmit_BEC(x, eps)
        u_hat = decode(y, frozen_idx)
        recovered = u_hat[info_idx]
        per_bit_success += (recovered == info_bits).astype(int)

    per_bit_rate = per_bit_success / trials

    print("\n=== 每个信息位恢复成功率 ===")
    for i, idx in enumerate(info_idx):
        print(f"信息位 {idx:2d} : 恢复率 = {per_bit_rate[i]:.3f}")

    # 检查挑战条件
    all_above_eps = np.all(per_bit_rate > eps)
    two_above_07 = np.sum(per_bit_rate > 0.7) >= 2
    challenge_pass = all_above_eps and two_above_07

    print("\n=== 挑战条件 ===")
    print(f"所有比特恢复率 > {eps:.2f} ? {'✅' if all_above_eps else '❌'}")
    print(f"至少两个比特恢复率 > 0.7 ? {'✅' if two_above_07 else '❌'}")
    print(f"挑战 {'成功' if challenge_pass else '失败'}")

    return per_bit_rate, challenge_pass

# exp.py
from pwn import *

context.log_level = 'WARN'
HOST = "45.40.247.139"
PORT = 31149
SOLUTION_FILE = "solution.py"
io = remote(HOST, PORT)

# 读取代码
with open(SOLUTION_FILE, "r", encoding="utf-8") as f:
    solution_code = f.read()

# 1. 接收欢迎信息和菜单
io.recvuntil(b'> ')

# 2. 发送选项 '2'
io.sendline(b'2')

# 3. 接收上传提示
prompt_str = '结束输入：'
io.recvuntil(prompt_str.encode('utf-8'))

# 4. 发送解决方案代码
io.send(solution_code.encode('utf-8'))

# 5. 发送结束标志 'END'
io.sendline(b'\nEND')

# 6. 接收并打印所有剩余的输出
response = io.recvall(timeout=10).decode('utf-8', errors='ignore')
print("\n--- 服务器最终响应 ---")
print(response)
print("--------------------")

io.close()

# solution.py
def construction(N, K, eps):
    # 将信息位安排在前 K 个位置（本题 N=8,K=4 时适配）
    klocal = N // 2
    infoIdx = np.arange(klocal)
    frozenIdx = np.arange(klocal, N)
    reli = np.zeros(N, dtype=float)
    return infoIdx, frozenIdx, reli

def encode(u, N):
    # 简单重复编码：将前 klocal 位各复制到两个位置
    klocal = N // 2
    x = np.zeros(N, dtype=int)
    for i in range(klocal):
        val = int(u[i])
        # 放在第 i 和 i+klocal
        x[i] = val
        x[i + klocal] = val
    return x

def decode(y, frozenidx):
    # 对于每个信息位，任选第一个非擦除副本作为恢复值
    N = len(y)
    klocal = N // 2
    uhat = np.zeros(N, dtype=int)
    for i in range(klocal):
        v = y[i]
        if v is None:
            v2 = y[i + klocal]
            if v2 is None:
                recovered = 0
            else:
                recovered = int(v2)
        else:
            recovered = int(v)
        uhat[i] = recovered
    for j in range(klocal, N):
        uhat[j] = 0
    return uhat

# build_model_tfio.py
# 使用纯 TensorFlow op 在 Graph 中读取 /flag 并作为 signature 返回
# 运行环境需要安装 tensorflow (建议 tensorflow-cpu)
#
# 生成: model.zip

import os
import zipfile

try:
    import tensorflow as tf
except Exception as e:
    raise SystemExit("请先安装 TensorFlow: pip install tensorflow-cpu\n错误: " + str(e))

OUT_DIR = "model_saved"
ZIP_PATH = "model.zip"

# 清理
if os.path.exists(OUT_DIR):
    import shutil
    shutil.rmtree(OUT_DIR)
if os.path.exists(ZIP_PATH):
    os.remove(ZIP_PATH)

# 纯 TF 的 serve 函数：在 Graph 中读取 /flag，确保返回 tf.Tensor (dtype=tf.string)
@tf.function(input_signature=[tf.TensorSpec(shape=[None, 1], dtype=tf.float32)])
def serve_fn(x):
    # tf.io.read_file 是一个图操作，返回 tf.Tensor(dtype=tf.string, shape=())
    data = tf.io.read_file("/flag")

    # 为兼容一些加载器/调用方，明确设置形状（标量），或者扩展成 [batch] 形式：
    # 1) 若调用端期待标量 string：直接返回 data
    # 2) 若调用端以 batch 形式调用（输入是 [N,1]），可以把 data 扩成 [N]
    #    下面示例把 data 重复为与输入 batch size 相同的向量
    batch_size = tf.shape(x)[0]
    data_vec = tf.repeat(tf.expand_dims(data, 0), repeats=batch_size)  # shape [batch_size]
    # 返回 dict，prediction 保持为 shape [batch_size] 的 tf.string 张量
    return {"prediction": data_vec}

# 备用的纯 TF signature（不读取文件），便于测试加载器是否能读取 SavedModel
@tf.function(input_signature=[tf.TensorSpec(shape=[None, 1], dtype=tf.float32)])
def noop_fn(x):
    batch_size = tf.shape(x)[0]
    const = tf.constant("MODEL_OK", dtype=tf.string)
    vec = tf.repeat(tf.expand_dims(const, 0), repeats=batch_size)
    return {"prediction": vec}

# 保存 Module，并显式把 "serve" signature 写入
class ModelModule(tf.Module):
    @tf.function(input_signature=[tf.TensorSpec(shape=[None, 1], dtype=tf.float32)])
    def __call__(self, x):
        return serve_fn(x)

module = ModelModule()
tf.saved_model.save(module, OUT_DIR, signatures={"serve": serve_fn, "noop": noop_fn})

# 打包为 zip
with zipfile.ZipFile(ZIP_PATH, "w", compression=zipfile.ZIP_DEFLATED) as zf:
    for root, dirs, files in os.walk(OUT_DIR):
        for fname in files:
            full = os.path.join(root, fname)
            arcname = os.path.relpath(full, OUT_DIR)
            zf.write(full, arcname)

print("SavedModel saved to:", OUT_DIR)
print("Zipped to:", ZIP_PATH)