import gradio as gr
import numpy as np
import plotly.graph_objects as go
import numba


# N sequences, each with n coinflips.
# for each sequence, determine the length of the longest constant 0 interval.
def histogram_vectorized(N, n):
    data = np.random.binomial(size=(N, n + 2), n=1, p=0.5)
    # putting 1s at both ends as barriers:
    data[:, 0] = 1
    data[:, -1] = 1

    ks = []
    for i in range(N):
        seq = data[i]
        ones = np.nonzero(seq)[0]
        jumps = np.diff(ones)
        k = np.max(jumps) - 1
        ks.append(k)
    return np.array(ks)


@numba.njit
def length_of_longest(seq):
    longest = 0
    current = 1
    for j in range(1, len(seq)):
        if seq[j - 1] == seq[j]:
            current += 1
        else:
            if current > longest:
                longest = current
            current = 1
    if current > longest:
        longest = current
    return longest


# N sequences, each with n coinflips.
# for each sequence, determine the length of the longest constant 0 interval.
@numba.njit
def histogram(N, n):
    data = np.random.binomial(size=(N, n), n=1, p=0.5)
    ks = np.empty(N)
    for i in range(N):
        ks[i] = length_of_longest(data[i])
    return ks


def visualize(x):
    if len(x) > 0:
        x_range = int(np.max(x) - np.min(x) + 1)
    else:
        x_range = 1
    fig = go.Figure(data=[go.Histogram(x=x, marker=dict(color='#7A88CC'), nbinsx=x_range)])
    fig.update_layout(title='Longest row of same flips', xaxis_title='length k', yaxis_title='count')
    return fig


def update(N, n):
    return visualize(histogram(N, n))


demo = gr.Interface(update,
    inputs = [
        gr.Slider(0, 10000, value=1000, label="N, the number of simulations"),
        gr.Slider(0, 1000, value=50, label="n, the number of coin flips in a simulation"),
    ],
    outputs = ["plot"],
    live=True)


demo.launch(show_api=False)