import gradio as gr
import pandas as pd
import numpy as np
import json
from io import StringIO


def dynamic_huff_model(df_distances, df_attractiveness, alpha, beta, df_capacity, df_population=None, iterations=5, crowding_threshold=1.0):
    """
    Iteratively calculates the distribution of people/visitors to destinations considering capacity and crowding based on an extended Huff model with linear decay function.

    Parameters:
    - df_distances, df_attractiveness, alpha, beta, df_capacity, df_population are the same as before.
    - iterations (int): The number of iterations to distribute the population.
    - crowding_threshold (float): The ratio of current visitors to capacity at which the decay of attractiveness starts.

    Returns:
    - pd.DataFrame: A DataFrame with the final distribution of visitors to each destination.
    """
    if df_population is None:
        df_population = pd.Series(np.ones(df_distances.shape[0]), index=df_distances.index)

    # Initialize the visitors DataFrame
    df_visitors = pd.DataFrame(0, index=df_distances.index, columns=df_distances.columns)
    
    # Distribute the population over the iterations
    df_population_per_iteration = df_population / iterations

    # Run the iterative distribution process
    for i in range(int(iterations)):
        print("iteration " + str(i) + "/"+str(int(iterations)))
        attractiveness = df_attractiveness.copy()
        current_visitors = df_visitors.sum(axis=0)
        
        # Calculate the decay based on the relative share of free capacity
        print(" Calculate the decay based on the relative share of free capacity")
        relative_crowding = current_visitors / df_capacity
        decay_factor = np.where(relative_crowding < crowding_threshold, 1, 1 - (relative_crowding - crowding_threshold) / (1 - crowding_threshold))
        attractiveness *= decay_factor
        
        # Calculate Huff model probabilities
        print("Calculate Huff model probabilities")
        distance_term = df_distances ** -beta
        # If df_distances is a DataFrame and df_attractiveness is a Series, you might need an operation like:
        numerator = df_distances.multiply(df_attractiveness, axis=0)  # Adjust based on actual intent

        denominator = numerator.sum(axis='columns')
        probabilities = numerator.div(denominator, axis='index').fillna(0)

        print("Distribute visitors based on probabilities and population")
        # Distribute visitors based on probabilities and population
        visitors_this_iteration = probabilities.multiply(df_population_per_iteration, axis='index')
        
        # Adjust for excess visitors beyond capacity
        potential_new_visitors = df_visitors + visitors_this_iteration
        excess_visitors = potential_new_visitors.sum(axis=0) - df_capacity
        excess_visitors[excess_visitors < 0] = 0
        visitors_this_iteration -= visitors_this_iteration.multiply(excess_visitors, axis='columns') / visitors_this_iteration.sum(axis=0)
        
        df_visitors += visitors_this_iteration

    # Return the final distribution of visitors
    return df_visitors

def app_function(input_json):
    print("something happend")
    # Parse the input JSON string
    print(input_json)

    try:
        inputs = json.loads(input_json)
    except:
        inputs = json.loads(input_json.replace("'", '"'))
    print(inputs.keys())
    # navigate to payload
    inputs = inputs["input"]

    print(inputs.keys())
    # Convert 'df_distances' from a list of lists into a DataFrame directly
    df_distances = pd.DataFrame(inputs["df_distances"])
    
    # 'inputs["df_attractiveness"]' and others are directly usable as lists
    df_attractiveness = pd.Series(inputs["df_attractiveness"])
    alpha = inputs["alpha"]
    beta = inputs["beta"]
    df_capacity = pd.Series(inputs["df_capacity"])
    
    # Check if 'df_population' is provided and convert to Series
    df_population = pd.Series(inputs["df_population"]) if "df_population" in inputs else None
    
    iterations = int(inputs.get("iterations", 5))
    crowding_threshold = inputs.get("crowding_threshold", 1.0)
    
    # Assuming dynamic_huff_model function is correctly defined to accept these parameters
    result = dynamic_huff_model(df_distances, df_attractiveness, alpha, beta, df_capacity, df_population, iterations, crowding_threshold)
    print(result)
    # Convert the result DataFrame to a JSON string for output
    return result.to_json(orient='split')

# Define the Gradio interface with a single JSON input
iface = gr.Interface(
    fn=app_function,
    inputs=gr.Textbox(label="Input JSON", lines=20, placeholder="Enter JSON with all parameters here..."),
    outputs=gr.JSON(label="Output JSON"),
    title="Dynamic Huff Model"
)

iface.launch()