def make_dataset(text, seq_len):
    chars = sorted(set(text))
    stoi = {ch: i for i, ch in enumerate(chars)}
    itos = {i: ch for ch, i in stoi.items()}
    data = torch.tensor([stoi[ch] for ch in text])
    inputs, targets = [], []
    for i in range(0, len(data) - seq_len - 1, seq_len):
        inputs.append(data[i:i + seq_len])
        targets.append(data[i + 1:i + seq_len + 1])
    return torch.stack(inputs), torch.stack(targets), stoi, itos


@torch.no_grad()
def generate(model, stoi, itos, prime="To ", length=200):
    model.eval()
    h = torch.zeros(1, model.hidden_size)
    c = torch.zeros(1, model.hidden_size)
    out_chars = list(prime)
    x = torch.tensor([[stoi[ch] for ch in prime]])
    logits, (h, c) = model(x, None)
    last_char = out_chars[-1]

    for _ in range(length):
        x = torch.tensor([[stoi[last_char]]])
        logits, (h, c) = model(x, (h, c))
        probs = F.softmax(logits[0, -1], dim=0)
        idx = torch.multinomial(probs, 1).item()
        last_char = itos[idx]
        out_chars.append(last_char)
    model.train()
    return "".join(out_chars)


if __name__ == "__main__":
    torch.manual_seed(0)
    # A short public-domain excerpt (paraphrased opening rhythm,
    # public-domain style text) stands in for a larger Gutenberg
    # corpus in this minimal example.
    text = (
        "To be or not to be that is the question whether tis "
        "nobler in the mind to suffer the slings and arrows of "
        "outrageous fortune " * 20
    )
    seq_len = 40
    X, Y, stoi, itos = make_dataset(text, seq_len)

    model = CharLSTM(vocab_size=len(stoi), hidden_size=128)
    optimizer = torch.optim.Adam(model.parameters(), lr=2e-3)

    for epoch in range(60):
        logits, _ = model(X)
        loss = F.cross_entropy(
            logits.reshape(-1, len(stoi)), Y.reshape(-1)
        )
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if epoch % 20 == 0:
            print(f"epoch {epoch}  loss {loss.item():.4f}")

    print("\nGenerated text:")
    print(generate(model, stoi, itos, prime="To ", length=100))
