import gradio as gr
import time
import math

from backend import test_optimized_sequence

def run_max(k_start, num_iterations):
    """
    1) Converts inputs to integers, handles limits, and runs the test.
    2) Returns:
       - summary_text: main results (iterations, prime counts, etc.)
       - enrichment_text: a step-by-step table for enrichment factors
       - results_state: data for show_all_primes
       - primes_found_state: list of found primes (N and d)
    """
    # Validation
    try:
        k_start = str(k_start).strip()
        if not k_start.isdigit():
            return "Error: k must be a positive integer.", "", "", ""

        k_length = len(k_start)
        if k_length > 400:
            return "Error: k must have at most 400 digits.", "", "", ""

        k_start = int(k_start)
        num_iterations = int(num_iterations)

        # Define max iterations based on digit length
        if k_length <= 10:
            max_iterations = 1_000_000
        elif k_length <= 100:
            max_iterations = 100_000
        else:  # 100 < k_length <= 400
            max_iterations = 10_000

        if num_iterations <= 0 or num_iterations > max_iterations:
            return f"Error: The number of iterations must be between 1 and {max_iterations}.", "", "", ""

    except ValueError:
        return "Error: Please enter valid integer values.", "", "", ""

    # Execution
    start_time = time.time()
    results, total_primes = test_optimized_sequence(k_start, num_iterations)
    end_time = time.time()

    total_iter = num_iterations
    execution_time = end_time - start_time
    success_percentage = (total_primes / total_iter) * 100 if total_iter else 0

    # Distribution of prime N/d
    prime_N_digits = {}
    prime_d_digits = {}
    primes_found = []

    prime_N_count = 0
    prime_d_count = 0
    sum_digits_N = 0
    sum_digits_d = 0

    for r in results:
        _, _, _, N, d, is_N_prime, is_d_prime = r
        if is_N_prime:
            prime_N_count += 1
            digits_N = len(str(N))
            sum_digits_N += digits_N
            prime_N_digits[digits_N] = prime_N_digits.get(digits_N, 0) + 1
            primes_found.append(str(N))

        if is_d_prime:
            prime_d_count += 1
            digits_d = len(str(d))
            sum_digits_d += digits_d
            prime_d_digits[digits_d] = prime_d_digits.get(digits_d, 0) + 1
            primes_found.append(str(d))

    # Ecco il cambio: elenco puntato per il count
    prime_N_summary = "\n".join([
        f"- {count} prime(s) of {digits} digits" for digits, count in sorted(prime_N_digits.items())
    ])
    prime_d_summary = "\n".join([
        f"- {count} prime(s) of {digits} digits" for digits, count in sorted(prime_d_digits.items())
    ])

    summary_text = (
        f"**Total iterations**: {total_iter}\n\n"
        f"**Total prime numbers found (N and d)**: {total_primes}\n\n"
        f"**Percentage of primes found per iteration**: {success_percentage:.2f}%\n\n"
        f"**Total execution time**: {execution_time:.2f} seconds\n\n"
        f"### Prime N digits count:\n{prime_N_summary}\n\n"
        f"### Prime d digits count:\n{prime_d_summary}\n"
    )

    # Enrichment factor steps

    # Step 1: average digit count (a_N, a_d)
    a_N = sum_digits_N / prime_N_count if prime_N_count else 0
    a_d = sum_digits_d / prime_d_count if prime_d_count else 0

    # Step 2: p = 1 / (a * ln(10)) (if a>0)
    def prime_probability(a):
        return 1.0 / (a * math.log(10)) if a > 0 else 0

    p_N = prime_probability(a_N)
    p_d = prime_probability(a_d)

    # Step 3: E = I * p
    E_N = total_iter * p_N
    E_d = total_iter * p_d
    E_tot = E_N + E_d

    # Step 4: O (observed)
    O_N = prime_N_count
    O_d = prime_d_count
    O_tot = O_N + O_d

    # Step 5: C = O / E
    C_N = (O_N / E_N) if E_N else 0
    C_d = (O_d / E_d) if E_d else 0
    C_tot = (O_tot / E_tot) if E_tot else 0

    enrichment_text = f"""
## Enrichment Factors - Step by Step

In this experiment, the enrichment factors (C) are computed in five steps:

1. a (average digits): average number of digits in the prime numbers found  
2. p = 1 / [a * ln(10)]: estimated probability of primality  
3. E = I * p: expected number of primes, given I total iterations  
4. O: observed number of primes  
5. C = O / E: enrichment factor

| Step  |       N       |       d       |     Total     |
|-------|---------------|---------------|---------------|
| 1. a  | {a_N:.2f}     | {a_d:.2f}     | -             |
| 2. p  | {p_N:.4f}     | {p_d:.4f}     | -             |
| 3. E  | {E_N:.2f}     | {E_d:.2f}     | {E_tot:.2f}   |
| 4. O  | {O_N}         | {O_d}         | {O_tot}       |
| 5. C  | {C_N:.2f}     | {C_d:.2f}     | {C_tot:.2f}   |

**Interpretation**:  
- If C > 1, there are more primes than expected.  
- If C < 1, there are fewer primes than expected.  
- If C = 1, observation matches the expectation.
"""

    return summary_text, enrichment_text, results, primes_found

def show_all_primes(results, primes_found, show_all):
    """
    Displays all found prime numbers if the user selects 'yes'.
    """
    if show_all.lower() == "yes":
        prime_list = "\n".join([f"{i+1}. {p}" for i, p in enumerate(primes_found)])
        return prime_list, "Displaying full prime list."
    else:
        return "", "Finished without displaying all prime numbers."

# Build the Gradio interface
with gr.Blocks() as iface:
    gr.Markdown("# 🔢 Prime Number Analysis - Sequence A (Modulo 3 & 7 Applied)")
    gr.Markdown("### Visit my website (https://max-russo.com) for a comprehensive overview of my research on MAX Prime Theory and MAX Key.")
    gr.Markdown("### Enter parameters to find prime numbers using Sequence A")

    with gr.Row():
        k_start_input = gr.Textbox(
            label="Starting value for k (e.g., 1234567890)",
            placeholder="Enter a positive integer with up to 400 digits"
        )
        num_iterations_input = gr.Textbox(
            label="Number of iterations (e.g., 50000)",
            placeholder="Enter an integer (max 1M for k ≤ 10, 100K for k ≤ 100, 10K for k ≤ 400)"
        )

    generate_button = gr.Button("Generate Prime Numbers")
    gr.Markdown("---")  # visual separator

    # Two outputs side by side:
    with gr.Row():
        summary_output = gr.Markdown(label="Prime Results")
        enrichment_output = gr.Markdown(label="Enrichment Table")

    # Hidden states to store results
    results_state = gr.State()
    primes_found_state = gr.State()

    gr.Markdown("### Do you want to see the full list of found prime numbers?")
    show_all_input = gr.Radio(["yes", "no"], label="Show full list of primes?", value="no")
    show_all_button = gr.Button("Show full results")

    # Output of the full prime list, if selected
    primes_output = gr.Textbox(
        label="Prime numbers found (if selected)", 
        lines=10, 
        max_lines=20, 
        interactive=False
    )
    message_output = gr.Textbox(label="Message", interactive=False)

    # Function wiring
    generate_button.click(
        fn=run_max,
        inputs=[k_start_input, num_iterations_input],
        outputs=[summary_output, enrichment_output, results_state, primes_found_state]
    )

    show_all_button.click(
        fn=show_all_primes,
        inputs=[results_state, primes_found_state, show_all_input],
        outputs=[primes_output, message_output]
    )

iface.launch()