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 _ in range(iterations):
        attractiveness = df_attractiveness.copy()
        current_visitors = df_visitors.sum(axis=0)
        
        # 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
        distance_term = df_distances ** -beta
        numerator = (attractiveness ** alpha).multiply(distance_term, axis='columns')
        denominator = numerator.sum(axis='columns')
        probabilities = numerator.div(denominator, axis='index').fillna(0)
        
        # 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):
    # Parse the input JSON string
    inputs = json.loads(input_json)
    
    # Convert stringified CSVs in JSON to DataFrames
    df_distances = pd.read_csv(StringIO(inputs["df_distances"]))
    df_attractiveness = pd.Series(json.loads(inputs["df_attractiveness"]))
    alpha = inputs["alpha"]
    beta = inputs["beta"]
    df_capacity = pd.Series(json.loads(inputs["df_capacity"]))
    df_population = pd.Series(json.loads(inputs["df_population"])) if "df_population" in inputs else None
    iterations = inputs.get("iterations", 5)
    crowding_threshold = inputs.get("crowding_threshold", 1.0)
    
    # Call the dynamic Huff model function with these parameters
    result = dynamic_huff_model(df_distances, df_attractiveness, alpha, beta, df_capacity, df_population, iterations, crowding_threshold)
    
    # Convert the result DataFrame to a CSV string for Gradio output
    return result.to_csv(index=False)

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

iface.launch()