import pandas as pd
import numpy as np
from datetime import datetime

# Adjust the data generation function as per the new conditions
def generate_enhanced_data_v2(num_samples=1000):
    data = []
    for _ in range(num_samples):
        # Randomly generate temperature and duration
        temp = np.random.randint(50, 201)  # Temperature between 50 and 200°C
        duration = np.random.randint(5, 120)  # Duration between 5 and 120 minutes

        # Assign risk level and alert based on conditions
        if temp <= 150 and duration <= 30:
            risk_level = "Low"
            risk_score = np.random.uniform(0, 40)  # Low risk
            alert = "Safe"
        elif 150 < temp <= 180 and 30 < duration <= 60:
            risk_level = "Moderate"
            risk_score = np.random.uniform(40, 70)  # Moderate risk
            alert = "Risk"
        else:
            risk_level = "High"
            risk_score = np.random.uniform(70, 100)  # High risk
            alert = "High Risk"
        
        # Add timestamp as current time
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

        # Append to data
        data.append([temp, duration, risk_level, risk_score, alert, timestamp])

    # Create DataFrame
    df = pd.DataFrame(data, columns=["temperature", "duration", "risk_level", "risk_score", "alert", "timestamp"])

    return df

# Generate the updated dataset with adjusted conditions
df_v2 = generate_enhanced_data_v2(1000)

# Save to CSV
df_v2.to_csv("enhanced_mantle_training.csv", index=False)

print("Data generation complete! Dataset saved as 'enhanced_mantle_training.csv'.")