"""
Example: Recurrent Bit Sequence Prediction (with LFSR)

This example demonstrates how an EvoNet with recurrent connections
can be evolved to predict the next bit in a pseudo-random binary sequence
generated by a 5-bit LFSR.

Key details:
- The sequence has period length 31 (maximal length LFSR).
- The sequence is repeated twice: first pass = warm-up, second pass = evaluation.
- Fitness combines mean squared error with a penalty for misclassification.
- Visualization shows input, target, and predictions during the evaluation phase.
"""

from evolib import Individual, Population, plot_bit_prediction
from evolib.utils.benchmarks import lfsr_sequence

FRAME_FOLDER = "07_frames"
CONFIG_FILE = "configs/07_recurrent_bit_prediction.yaml"

SEQ_LENGTH = 31
warmup_steps = SEQ_LENGTH

input_seq = lfsr_sequence(SEQ_LENGTH) * 2
target_seq = input_seq[1:] + [input_seq[0]]


def fitness_bitseq(indiv: Individual) -> None:
    """
    Evaluate how well the EvoNet predicts the next bit in a fixed sequence.

    Additionally logs classification accuracy in indiv.extra_metrics.
    Note: Fitness combines MSE with penalty term (1-accuracy)/4
    """
    net = indiv.para["brain"].net
    net.reset(full=True)

    total_error = 0.0
    correct = 0
    count = 0

    for t in range(len(input_seq)):
        output = net.calc([input_seq[t]])[0]
        if t >= warmup_steps:
            target = target_seq[t]
            # --- MSE ---
            error = output - target
            total_error += error**2

            # --- Accuracy ---
            pred_bit = 1 if output > 0.5 else 0
            if pred_bit == target:
                correct += 1
            count += 1

    mse = total_error / count
    acc = correct / count if count > 0 else 0.0

    # Log extra metrics
    indiv.extra_metrics = {"accuracy": acc, "mse": mse}

    indiv.fitness = mse + (1 - acc) / 4


def save_plot(pop: Population) -> None:
    best = pop.best()
    net = best.para["brain"].net
    net.reset(full=True)

    # Warmup
    for bit in input_seq[:SEQ_LENGTH]:
        net.calc([bit])

    # Prediction
    y_preds = [net.calc([bit])[0] for bit in input_seq[SEQ_LENGTH:]]

    true_bits = target_seq[SEQ_LENGTH:]
    input_bits = input_seq[SEQ_LENGTH:]

    acc = best.extra_metrics.get("accuracy", 0.0)
    mse = best.extra_metrics.get("mse", best.fitness)

    print(
        f"[Gen {pop.generation_num}] best fitness={best.fitness:.5f}, "
        f"acc={acc:.3%}, mse={mse:.5f}"
    )

    plot_bit_prediction(
        true_bits=true_bits,
        pred_values=y_preds,
        input_bits=input_bits,
        title=(
            f"Bit Prediction (gen={pop.generation_num}, "
            f"MSE={mse:.5f}, Acc={acc:.3%})"
        ),
        save_path=f"{FRAME_FOLDER}/gen_{pop.generation_num:03d}.png",
    )


if __name__ == "__main__":
    pop = Population(CONFIG_FILE, fitness_function=fitness_bitseq)
    pop.run(verbosity=0, on_generation_end=save_plot)

    best = pop.best()
    best.para["brain"].net.plot("07_recurrent_bit_prediction", fillcolors_on=True)
