from transformers import Qwen3ForCausalLM
import torch
import torch.nn.functional as F


def bits_from_interval(min_val, max_val, numbits):
    df = 0.5

    tal = 0
    my_bits = []
    for i in range(numbits):
        if (tal < min_val and tal + df < max_val):
            my_bits.append(1)
            tal += df
        else:
            my_bits.append(0)
        df /= 2
    return my_bits


def bits_to_tal(bits):
    # print(f"Antal 0:or: {bits.count(0)}")
    # print(f"Antal 1:or: {bits.count(1)}")
    tal = 0
    for i in range(len(bits)):
        tal += bits[i] * 2 ** i
    return tal


def bits_from_tal(tal):
    Bits = []

    while True:
        rest = tal % 2
        Bits.append(rest)
        tal -= rest
        tal //= 2
        if (tal == 0):
            break
    return Bits


def bits_to_text(start, bits, min_length, model, device, do_prints=False, write=False):
    Bits = bits
    BitID = 0
    text_enc = list(start.encode())
    b = []
    output = ""
    while True:
        # print(text_enc)
        if (BitID >= len(Bits) and len(output) >= min_length):
            break

        current_intervals = []
        # print(text_enc)
        seed = torch.tensor(text_enc, dtype=torch.long)[None].to(device)
        with torch.no_grad():
            out = model(seed)

        logits = out.logits  # shape: (batch, seq_len, vocab_size)

        probs = F.softmax(logits[0, -1], dim=-1)  # sannolikheter för nästa byte
        topk = torch.topk(probs, k=256)  # de 10 troligaste

        value = 0

        min_val = None
        max_val = 0
        all_bits = []
        remove_bits = 0
        ones = 0
        for prob, idx in zip(topk.values, topk.indices):
            min_val = max_val
            max_val += prob.item()
            current_bits = bits_from_interval(min_val, max_val, 50)
            all_bits.append((current_bits, idx))

            prev_ones = ones
            ones = 0
            for bit in current_bits:
                if (bit == 1):
                    ones += 1
                else:
                    break
            if (ones - prev_ones >= 2):
                remove_bits += ones - prev_ones - 1
            for i in range(remove_bits):
                current_bits.pop(0)

            byte_val = idx.item()
            try:
                char = bytes([byte_val]).decode("utf8")
            except UnicodeDecodeError:
                char = f"<byte {byte_val}>"  # kan hända vid multi-byte UTF-8-tecken (t.ex. å, ä, ö)
            # print(f"{char!r}: {min_val} to {max_val}")

        expected_bits = []
        prediction = None
        if (do_prints):
            print("alla bitsekvenser:", all_bits[:10])
        for i in range(100):

            if (BitID < len(Bits)):
                bit = Bits[BitID]
            else:
                bit = 0
            BitID += 1
            expected_bits.append(bit)
            if (do_prints):
                print("Hittills extraherade bits:", expected_bits)
            my_bits = None
            potential = None

            matches = []

            for c in all_bits:
                other_bits = c[0]
                idx = c[1]
                if (len(other_bits) < len(expected_bits)):
                    continue
                match = True
                for j in range(len(expected_bits)):
                    if (expected_bits[j] != other_bits[j]):
                        match = False
                        break
                if (match):
                    matches.append(idx)
            if (len(matches) == 1):
                prediction = matches[0]
                # print(expected_bits)
                break
            elif (len(matches) == 0):
                prediction = idx

        text_enc.append(prediction)
        b.append(prediction)
        if (len(b) > 4):
            b.pop(0)
        try:
            if (write):
                print(bytes(b).decode("utf-8"), end='', flush=True)
            output += bytes(b).decode("utf-8")
            b = []
        except:
            continue
    return output


def text_to_bits(text, model, device, do_prints=False):
    combined = []
    text_enc = list(text.encode())
    intervals = []
    for i in range(1, len(text_enc)):
        # print("------")
        seed = torch.tensor(text_enc[:i], dtype=torch.long)[None].to(device)
        expected = torch.tensor(text_enc[i], dtype=torch.long)[None].to(device)
        # print(expected)

        with torch.no_grad():
            out = model(seed)

        logits = out.logits  # shape: (batch, seq_len, vocab_size)

        probs = F.softmax(logits[0, -1], dim=-1)  # sannolikheter för nästa byte
        topk = torch.topk(probs, k=256)  # de 10 troligaste
        if (do_prints):
            print(bytes([text_enc[i]]).decode("utf-8", errors="replace"), probs[expected].item())

        value = 0

        my_and_next = None
        all_bits = []

        min_val = None
        max_val = 0
        remove_bits = 0
        ones = 0
        for prob, idx in zip(topk.values, topk.indices):

            min_val = max_val
            max_val += prob.item()
            current_bits = bits_from_interval(min_val, max_val, 50)
            prev_ones = ones
            ones = 0
            for bit in current_bits:
                if (bit == 1):
                    ones += 1
                else:
                    break
            if (ones - prev_ones >= 2):
                remove_bits += ones - prev_ones - 1
            for i in range(remove_bits):
                current_bits.pop(0)

            byte_val = idx.item()
            try:
                char = bytes([byte_val]).decode("utf8")
            except UnicodeDecodeError:
                char = f"<byte {byte_val}>"  # kan hända vid multi-byte UTF-8-tecken (t.ex. å, ä, ö)
            # print(f"{char!r}: {min_val} to {max_val}")
            if (idx == expected):
                correct_bits = current_bits
            else:
                all_bits.append(current_bits)

        for k in range(1, len(correct_bits)):
            buffer = correct_bits[:k]
            any_similar = False
            for bits in all_bits:
                match = True
                for j in range(len(buffer)):
                    if (buffer[j] != bits[j]):
                        match = False
                if (match):
                    any_similar = True
                    break
            if (not any_similar):
                correct_bits = buffer
        for bit in correct_bits:
            combined.append(bit)

    return combined