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):
    """
    Generate a plot of FFMI (Fat-Free Mass Index) curves based on weight 
    and body fat percentage for a specified gender.

    :param height_m: Height of the individual in meters [m].
    :param gender: Gender of the individual ('Male' or 'Female'), which determines 
                   the FFMI ranges used for categorization.
    :return: A Plotly Figure object containing the scatter plot of FFMI curves.
    """
    weight_kg_range = np.linspace(45, 160, 100, endpoint=True)
    body_fat_percentage_range = np.linspace(1, 50, 100, endpoint=True)

    if gender == "Male":
        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 FAKEE", "color": "purple"},
            {"min": 28, "max": float("inf"), "description": "FAKEE", "color": "black"}
        ]
    else:
        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 FAKEE", "color": "purple"},
            {"min": 25, "max": float("inf"), "description": "FAKEE", "color": "black"}
        ]

    fig = go.Figure()

    for ffmi_range in ffmi_ranges:
        weight_kg_values = []
        body_fat_values = []
        ffmi_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)

                if ffmi_range["min"] <= ffmi < ffmi_range["max"]:
                    weight_kg_values.append(weight)
                    body_fat_values.append(body_fat_percentage)
                    ffmi_values.append(ffmi)

        fig.add_trace(go.Scatter(
            x=weight_kg_values,
            y=body_fat_values,
            mode='markers',
            name=ffmi_range["description"],
            marker=dict(
                color=ffmi_range.get("marker_color", ffmi_range["color"]),
                size=5, opacity=0.6,
                line=dict(color=ffmi_range.get("line_color", ffmi_range["color"]))
            ),
            hovertext=[
                f'FFMI: {ffmi:.2f}<br>Weight: {weight_kg:.2f}<br>BodyFat: {body_fat:.2f}<br>{ffmi_range["description"]}' for ffmi,weight_kg, body_fat in 
                       
                       zip(ffmi_values, weight_kg_values, body_fat_values)],
            hoverinfo='text'
        ))

    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'
    )

    return fig