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


def huff_model_probability(df_distances, df_attractiveness, alpha, beta, df_population=None, distance_threshold=None, decay_factor=0.1):
    """
    Calculates the probability of choosing among destinations and the adjustment factors for willingness to travel.
    """
    if df_population is None:
        df_population = pd.Series(np.ones(df_distances.shape[0]), index=df_distances.index)

    adjustment_factors = pd.DataFrame(index=df_distances.index, columns=df_distances.columns)
    if distance_threshold is not None and distance_threshold is not "None":
        # Calculate adjustment factors for each origin-destination pair
        for destination in df_distances.columns:
            adjustment_factors[destination] = df_distances[destination].apply(
                lambda x: np.exp(-(max(0, x - distance_threshold)) * decay_factor))
    else:
        adjustment_factors[:] = 1
    # deactivate adjustment factor !!!!!!!!!!!!!!!!!!!!
    adjustment_factors[:] = 1

    adjusted_population = df_population.repeat(df_distances.shape[1]).values.reshape(df_distances.shape) * adjustment_factors
    attractiveness_term = df_attractiveness ** alpha
    distance_term =np.exp(df_distances * -beta)

    numerator = (attractiveness_term * distance_term).multiply(adjusted_population, axis=0)
    denominator = numerator.sum(axis=1)
    probabilities = numerator.div(denominator, axis=0).fillna(0)

    return probabilities, adjustment_factors

def app_function(input_json):
    print("Received input")
    # Parse the input JSON string
    try:
        inputs = json.loads(input_json)
    except json.JSONDecodeError:
        inputs = json.loads(input_json.replace("'", '"'))
    print("Parsed input keys:", inputs.keys())

    # Assuming the input structure is correctly formatted to include the necessary parameters
    inputs = inputs["input"]

    # Convert 'df_distances' from a list of lists into a DataFrame
    df_distances = pd.DataFrame(inputs["df_distances"])
    print("df_distances shape:", df_distances.shape)
    
    # Convert 'df_attractiveness' into a Series
    df_attractiveness = pd.Series(inputs["df_attractiveness"])
    print("df_attractiveness shape:", df_attractiveness.shape)
    
    # Extract alpha and beta parameters
    alpha = inputs["alpha"]
    beta = inputs["beta"]
    
    # Check if 'df_population' is provided and convert to Series, else default to None
    df_population = pd.Series(inputs["df_population"]) if "df_population" in inputs else None

    # Additional parameters for the updated Huff model
    distance_threshold = inputs.get("distance_threshold", None)
    decay_factor = inputs.get("decay_factor", 0.1)  # Default decay factor if not provided
    
    # Call the updated Huff model function
    probabilities, adjustment_factors = huff_model_probability(
        df_distances=df_distances,
        df_attractiveness=df_attractiveness,
        alpha=alpha,
        beta=beta,
        df_population=df_population,
        distance_threshold=distance_threshold,
        decay_factor=decay_factor
    )

    # Prepare the output
    output = {
        "probabilities": probabilities.to_dict(orient='split'),
        "adjustment_factors": adjustment_factors.to_dict(orient='split')
    }

    return json.dumps(output)

# 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()