import json
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import r2_score
import joblib # Import joblib for saving the model

try:
    with open('./data.json', 'r') as f:
        data = json.load(f)
except FileNotFoundError:
    print("Error: 'data.json' not found. Please ensure the file is in the same directory as this script.")
    exit()

df = pd.DataFrame.from_dict(data, orient='index')
df[['latitude', 'longitude']] = pd.DataFrame(df['middle_point'].tolist(), index=df.index)
df.drop('middle_point', axis=1, inplace=True)
df = df[df['price'] > 1.0]

features = ['area', 'dis', 'type', 'latitude', 'longitude']
target = 'price'
X = df[features]
y = df[target]

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
print(f"Data loaded. Training model on {len(X_train)} samples...")

model = RandomForestRegressor(n_estimators=100, random_state=42)
model.fit(X_train, y_train)
print("Model training complete.")

y_pred = model.predict(X_test)
r2 = r2_score(y_test, y_pred)
print(f"Model R-squared on test data: {r2:.4f}")


model_filename = 'random_forest_model.joblib'
joblib.dump(model, model_filename)
print(f"\nModel saved successfully as '{model_filename}'")