import os
import pandas as pd
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
from joblib import dump, load

# Function to train models
def train_models():
    # Synthetic dataset
    data = {
        "Thickness": [10, 15, 20, 25],
        "Hole_Diameter": [5, 10, 15, 20],
        "Force": [5000, 7000, 10000, 12000],
        "Max_Stress": [300, 250, 200, 150],
        "Max_Deformation": [0.5, 0.4, 0.3, 0.2],
        "Pass_Fail": [1, 1, 0, 0]
    }
    df = pd.DataFrame(data)
    X = df[["Thickness", "Hole_Diameter", "Force"]]

    # Train regression models
    stress_model = RandomForestRegressor().fit(X, df["Max_Stress"])
    deformation_model = RandomForestRegressor().fit(X, df["Max_Deformation"])

    # Train classification model
    pass_fail_model = RandomForestClassifier().fit(X, df["Pass_Fail"])

    # Save models
    dump(stress_model, "stress_model.pkl")
    dump(deformation_model, "deformation_model.pkl")
    dump(pass_fail_model, "pass_fail_model.pkl")
    print("Models have been trained and saved.")

# Function to load models with error handling
def load_models():
    try:
        # Check if model files exist
        if not os.path.exists("stress_model.pkl") or not os.path.exists("deformation_model.pkl") or not os.path.exists("pass_fail_model.pkl"):
            print("Model files not found. Training models...")
            train_models()
        
        # Load models
        stress_model = load("stress_model.pkl")
        deformation_model = load("deformation_model.pkl")
        pass_fail_model = load("pass_fail_model.pkl")
        print("Models have been loaded.")
        
        return stress_model, deformation_model, pass_fail_model
    
    except Exception as e:
        print(f"Error loading models: {e}")
        return None, None, None

# Main program
if __name__ == "__main__":
    stress_model, deformation_model, pass_fail_model = load_models()

    if stress_model and deformation_model and pass_fail_model:
        print("Models are ready to use!")
    else:
        print("Failed to load models.")