import gradio as gr
from transformers import CLIPProcessor, CLIPModel
from PIL import Image

# Load the model and processor
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

# Define the broad BMI classes for Model 1
bmi_classes_model1 = [
    "underweight (x < 18.5 BMI)", 
    "normal weight (18.5 < x < 25 BMI)", 
    "overweight (25 BMI < x < 30)",  
    "obesity (x > 30 BMI)"
]

# Define the finer BMI ranges for Model 2
bmi_ranges_model2 = {
    "underweight (x < 18.5 bmi)": [
        "BMI less than 16.0", 
        "BMI between 16.0 and 16.99",  
        "BMI between 17.0 and 18.49"
    ],
    "normal weight (18.5 < x < 25 bmi)": [
        "BMI between 18.5 and 20.4", 
        "BMI between 20.5 and 22.4",  
        "BMI between 22.5 and 24.9"
    ],
    "overweight (25 bmi < x < 30)": [
        "BMI between 25.0 and 26.9", 
        "BMI between 27.0 and 28.9",  
        "BMI between 29.0 and 29.9"
    ],
    "obesity (x > 30 bmi)": [
        "BMI between 30.0 and 34.9", 
        "BMI between 35.0 and 39.9",  
        "BMI 40.0 and above"
    ]
}
bmi_mapp = {
    "BMI less than 16.0":"BMI < 16.0",
        "BMI between 16.0 and 16.99":"16.0 ≤ BMI ≤ 16.99",
        "BMI between 17.0 and 18.49":"17.0 ≤ BMI ≤ 18.49",
    "BMI between 18.5 and 20.4":"18.5 ≤ BMI ≤ 20.4",
        "BMI between 20.5 and 22.4":"20.5 ≤ BMI ≤ 22.4",
        "BMI between 22.5 and 24.9":"22.5 ≤ BMI ≤ 24.9",
    "BMI between 25.0 and 26.9":"25.0 ≤ BMI ≤ 26.9",
        "BMI between 27.0 and 28.9":"27.0 ≤ BMI ≤ 28.9",
        "BMI between 29.0 and 29.9":"29.0 ≤ BMI ≤ 29.9",
    "BMI between 30.0 and 34.9": "30.0 ≤ BMI ≤ 34.9",
        "BMI between 35.0 and 39.9": "35.0 ≤ BMI ≤ 39.9",
        "BMI 40.0 and above":"BMI ≥ 40.0"
}

# Define BMI ranges with boundaries and midpoints for adjusted BMI calculation
bmi_ranges = {
    "underweight (x < 18.5 bmi)": {
        "BMI < 16.0": (0, 16.0, 16.0),  # Upper boundary 16.0
        "16.0 ≤ BMI ≤ 16.99": (16.0, 16.99, 16.5),  # Midpoint 16.5
        "17.0 ≤ BMI ≤ 18.49": (17.0, 18.49, 17.75)  # Midpoint 17.75
    },
    "normal weight (18.5 < x < 25 bmi)": {
        "18.5 ≤ BMI ≤ 20.4": (18.5, 20.4, 19.45),  # Midpoint 19.45
        "20.5 ≤ BMI ≤ 22.4": (20.5, 22.4, 21.45),  # Midpoint 21.45
        "22.5 ≤ BMI ≤ 24.9": (22.5, 24.9, 23.7)    # Midpoint 23.7
    },
    "overweight (25 bmi < x < 30)": {
        "25.0 ≤ BMI ≤ 26.9": (25.0, 26.9, 25.95),  # Midpoint 25.95
        "27.0 ≤ BMI ≤ 28.9": (27.0, 28.9, 27.95),  # Midpoint 27.95
        "29.0 ≤ BMI ≤ 29.9": (29.0, 29.9, 29.45)   # Midpoint 29.45
    },
    "obesity (x > 30 bmi)": {
        "30.0 ≤ BMI ≤ 34.9": (30.0, 34.9, 32.5),   # Midpoint 32.5
        "35.0 ≤ BMI ≤ 39.9": (35.0, 39.9, 37.45),  # Midpoint 37.45
        "BMI ≥ 40.0": (40.0, 100, 40.0)            # Lower boundary 40.0
    }
}

def predict_bmi(image, height_in_inches):
    # Prepare inputs for Model 1
    inputs_model1 = processor(text=bmi_classes_model1, images=image, return_tensors="pt", padding=True)
    outputs_model1 = model(**inputs_model1)
    probs_model1 = outputs_model1.logits_per_image.softmax(dim=1)
    print(probs_model1,'probs_model1')
    
    # Get the broad category prediction from Model 1
    max_prob_index_model1 = probs_model1.argmax().item()
    predicted_bmi_class = bmi_classes_model1[max_prob_index_model1]
    print(predicted_bmi_class,'predicted_bmi_class')

    # Select class names for Model 2 based on Model 1's prediction
    model2_classes = bmi_ranges_model2[predicted_bmi_class.lower()]
    print(model2_classes,'model2_classes')
    
    # Prepare inputs for Model 2
    inputs_model2 = processor(text=model2_classes, images=image, return_tensors="pt", padding=True)
    outputs_model2 = model(**inputs_model2)
    probs_model2 = outputs_model2.logits_per_image.softmax(dim=1)
    print(probs_model2,'probs_model2')

    # Get the finer range prediction from Model 2
    max_prob_index_model2 = probs_model2.argmax().item()
    finer_bmi_range = model2_classes[max_prob_index_model2]
    print(finer_bmi_range,'finer_bmi_range')

    # Determine the BMI prediction based on the range
    bmi_prediction = get_adjusted_bmi(predicted_bmi_class, finer_bmi_range)
    print(bmi_prediction)
    
    # Calculate weight using user-provided height
    predicted_weight = calculate_weight(bmi_prediction, height_in_inches)

    # Create the JSON output
    result = {
        "weightCategory": f"{predicted_bmi_class} - {finer_bmi_range}",
        "bmiPrediction": f"{bmi_prediction:.2f}",
        "height": str(height_in_inches),
        "predictedWeight": f"{predicted_weight:.2f} lbs"
    }

    return result

def get_adjusted_bmi(weight_category, finer_range):
    """Return the appropriate BMI value for the given finer range within the weight category."""
    category_ranges = bmi_ranges.get(weight_category.lower())
    print(category_ranges,'category_ranges')
    for range_label, (low, high, mid) in category_ranges.items():
        print(range_label,'range_label')
        if range_label == bmi_mapp[finer_range]:
            return mid
        """elif "BMI <" in range_label or "BMI ≥" in range_label:
            return high if "BMI <" in range_label else low"""
    return None

def calculate_weight(bmi, height_in_inches):
    """Calculate the weight from BMI and height (in inches)."""
    height_in_meters = height_in_inches * 0.0254
    weight_kg = bmi * (height_in_meters ** 2)
    weight_lbs = weight_kg * 2.20462
    return weight_lbs

# Create Gradio interface with updated input components
interface = gr.Interface(
    fn=predict_bmi,
    inputs=[
        gr.Image(type="pil"),
        gr.Number(label="Height in inches")  # Allow user to enter height
    ],
    outputs="json",
    title="BMI Prediction",
    description="Upload an image and enter your height to predict BMI category and receive a detailed prediction."
)

interface.launch()