import numpy as np
import plotly.graph_objs as go
import plotly.io as pio

def calculate_fat_free_mass(weight_kg, body_fat_percentage):
    """
    Calculate fat-free mass in kilograms.
    
    :param weight_kg: Weight in kilograms [kg].
    :param body_fat_percentage: Body fat percentage [%].
    :return: Fat-free mass in kilograms [kg].
    """
    fat_free_mass = weight_kg * (1 - (body_fat_percentage / 100))
    return fat_free_mass

def calculate_ffmi(fat_free_mass_kg, height_m):
    """
    Calculate FFMI (Fat-Free Mass Index).
    
    :param fat_free_mass_kg: Fat-free mass in kilograms [kg].
    :param height_m: Height in meters [m].
    :return: FFMI in kg/m².
    """
    ffmi = fat_free_mass_kg / (height_m ** 2)
    return ffmi

def plot_ffmi_curves(height_m, gender):
    # Define weight range (X-axis)
    weight_kg_range = np.linspace(45, 160, 100)  # weight in kg

    # Define body fat percentages (Y-axis)
    body_fat_percentage_range = np.linspace(1, 50, 100)  # body fat percentage
    if gender == "Male":
        # Define FFMI ranges and their descriptions for mans
        ffmi_ranges = [
            {"min": 0, "max": 18, "description": "Below average", "color": "blue"},
            {"min": 18, "max": 20, "description": "Average", "color": "green"},
            {"min": 20, "max": 22, "description": "Above average", "color": "yellow"},
            {"min": 22, "max": 23, "description": "Excellent", "color": "orange"},
            {"min": 23, "max": 26, "description": "Superior", "color": "red"},
            {"min": 26, "max": 28, "description": "Suspicion of steroid use", "color": "purple"},
            {"min": 28, "max": float("inf"), "description": "Steroid usage likely", "color": "black"}
        ]
    else:
        # Define FFMI ranges and their descriptions for womens
        ffmi_ranges = [
            {"min": 0, "max": 15, "description": "Below average", "color": "blue"},
            {"min": 15, "max": 17, "description": "Average", "color": "green"},
            {"min": 17, "max": 18, "description": "Above average", "color": "yellow"},
            {"min": 18, "max": 19, "description": "Excellent", "color": "orange"},
            {"min": 19, "max": 21.5, "description": "Superior", "color": "red"},
            {"min": 21.5, "max": 25, "description": "Suspicion of steroid use", "color": "purple"},
            {"min": 25, "max": float("inf"), "description": "Steroid usage likely", "color": "black"}
        ]



    # Create plot
    fig = go.Figure()

    # Calculate and plot curves for each FFMI range
    for ffmi_range in ffmi_ranges:
        weight_kg_values = []
        body_fat_values = []

        for weight in weight_kg_range:
            for body_fat_percentage in body_fat_percentage_range:
                fat_free_mass = calculate_fat_free_mass(weight, body_fat_percentage)
                ffmi = calculate_ffmi(fat_free_mass, height_m)

                # Check if the calculated FFMI falls within the current FFMI range
                if ffmi_range["min"] <= ffmi < ffmi_range["max"]:
                    weight_kg_values.append(weight)
                    body_fat_values.append(body_fat_percentage)

        # Add curve for the current FFMI range
        fig.add_trace(go.Scatter(
            x=weight_kg_values,
            y=body_fat_values,
            mode='markers',
            name=ffmi_range["description"],
            marker=dict(color=ffmi_range["color"], size=5, opacity=0.6),
        ))

    # Update layout
    fig.update_layout(
        title='FFMI Curves for Different Body Fat Percentages and Weight',
        xaxis_title='Weight (kg)',
        yaxis_title='Body Fat Percentage (%)',
        legend_title='FFMI Range',
        hovermode='closest',
        template='plotly_dark'
    )

    # Show the interactive plot
    return fig