import logging, os
import time
from metaflow import FlowSpec, step, S3, conda_base

try:
    import joblib

    print("joblib is installed")
except ImportError:
    print("joblib is not installed")
    pass


@conda_base(python="3.9.13", packages={"scikit-learn": "1.3.1"})
class ModelTrainingFlow(FlowSpec):
    @step
    def start(self):
        # Load and preprocess data
        # os.system("pip install scikit-learn")
        from sklearn.datasets import load_iris

        self.data = load_iris()
        self.next(self.train_model)

    @step
    def train_model(self):
        # os.system("pip install scikit-learn")
        from sklearn.ensemble import RandomForestClassifier

        # Train the model
        self.model = RandomForestClassifier()
        self.model.fit(self.data.data, self.data.target)
        self.next(self.evaluate_model)

    @step
    def evaluate_model(self):
        # Evaluate the model
        print("Evaluating model...")
        time.sleep(5)
        print("Model Evaluated successfully!")
        self.next(self.save_to_s3)

    @step
    def save_to_s3(self):
        # Save model to S3
        # os.system("pip install joblib")
        import joblib

        with open("model.pkl", "wb") as f:
            joblib.dump(self.model, f)
            with open(f.name, "rb") as in_file:
                data = in_file.read()
                with S3(run=self) as s3:
                    url = s3.put(f.name, data)
                    # print it out for debug purposes
                    print("Model saved at: {}".format(url))
                    # save this path for downstream reference!
                    self.s3_path = url
        self.next(self.end)

    @step
    def end(self):
        # End of the flow
        print("Model training and uploading to S3 completed.")


if __name__ == "__main__":
    ModelTrainingFlow()