import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, r2_score
import joblib

# Load the data
data = pd.read_csv('sugar_cane_data.csv')

# Data preprocessing
data['Farm'] = data['Farm'].astype(str)
data['Variety'] = data['Variety'].fillna('Unknown')
data['Age'] = data['Age'].fillna('Unknown')

# Split features and target variables
X = data[['Farm', 'Variety', 'Age']]
y = data[['Brix', 'Purity', 'Pol', 'RS']]

# Split the data into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Create preprocessing steps
preprocessor = ColumnTransformer(
    transformers=[
        ('num', StandardScaler(), []),
        ('cat', OneHotEncoder(handle_unknown='ignore'), ['Farm', 'Variety', 'Age'])
    ])

# Create a pipeline with preprocessor and random forest regressor
model = Pipeline([
    ('preprocessor', preprocessor),
    ('regressor', RandomForestRegressor(n_estimators=100, random_state=42))
])

# Fit the model
model.fit(X_train, y_train)

# Make predictions on the test set
y_pred = model.predict(X_test)

# Evaluate the model
mse = mean_squared_error(y_test, y_pred)
r2 = r2_score(y_test, y_pred)

print(f"Mean Squared Error: {mse}")
print(f"R-squared Score: {r2}")

# Save the model
joblib.dump(model, 'sugar_cane_model.joblib')

# Function to make predictions for new data
def predict_sugar_cane_properties(farm, variety, age):
    new_data = pd.DataFrame({'Farm': [farm], 'Variety': [variety], 'Age': [age]})
    prediction = model.predict(new_data)
    return {
        'Brix': prediction[0][0],
        'Purity': prediction[0][1],
        'Pol': prediction[0][2],
        'RS': prediction[0][3]
    }

# Example usage
print(predict_sugar_cane_properties('01-18', 'CP69', 'P'))