import os

import pandas as pd
import numpy as np

from sklearn.preprocessing import StandardScaler
from sklearn import tree

"""
Model used to predict the space size to provide to `scipy.interpolate.splev`
in the function `gui.getSpline`. This function is used to generate the coordinates 
for `skimage.draw.polygon` when converting a closed spline to an object 
(`gui.splineToObj` function).

The training data for the model is generated by the script 
`cellacdc._profile.spline_to_obj.profile_skimage_draw_polygon.py`
"""

pwd_path = os.path.dirname(os.path.abspath(__file__))

class Model:
    def __init__(self):
        pass

    def fit(self):
        # Read data
        filename = '1_exec_time_space_size_step_10.csv'
        df_path = os.path.join(pwd_path, filename)
        df = pd.read_csv(df_path)
        
        # Define predictor and response variables
        X_train = df[['bbox_area', 'exec_time']]
        y_train = df['space_size']

        # Scale the data
        scaler = StandardScaler().fit(X_train)
        X_scaled = pd.DataFrame(
            scaler.transform(X_train), columns=X_train.columns
        )
        self.scaler = scaler

        # Define regression model and fit
        model = tree.DecisionTreeRegressor() # LinearRegression()
        reg = model.fit(X_scaled, y_train)

        self.model = model

    def predict(self, bbox_area, max_exec_time=150):
        X_pred = pd.DataFrame({
            'bbox_area': [bbox_area], 'exec_time': [max_exec_time]
        })
        X_scaled = pd.DataFrame(
            self.scaler.transform(X_pred), columns=X_pred.columns
        )
        y_pred = self.model.predict(X_scaled)
        pred_space_size = y_pred[0]
        return pred_space_size
