import sys
import pandas as pd
import logging


from src.exception import CustomException
from src.logger import logging
from src.utils import load_object



class PredictPipeline:
    def __init__(self):
        pass

    def predict(self, features):
        try:
            model_path = 'artifacts/model.pkl'
            preprocessor_path = 'artifacts/preprocessor.pkl'

            # Load model and preprocessor
            model = load_object(file_path=model_path)
            preprocessor = load_object(file_path=preprocessor_path)

            # Check if model and preprocessor are loaded
            if model is None:
                raise Exception("Model failed to load")
            if preprocessor is None:
                raise Exception("Preprocessor failed to load")

            # Transform features and make prediction
            data_scaled = preprocessor.transform(features)
            preds = model.predict(data_scaled)
            
            return preds
            
        except Exception as e:
            logging.error(f"Error in prediction: {str(e)}")
            raise CustomException(e, sys)


class CustomData:
    def __init__(self, gender: str, race: str,
                 parental_level_of_education: str,
                 lunch: str,
                 test_preparation_course: str,
                 reading_score: float,
                 writing_score: float):
        self.gender = gender
        self.race = race
        self.parental_level_of_education = parental_level_of_education
        self.lunch = lunch
        self.test_preparation_course = test_preparation_course
        self.reading_score = reading_score
        self.writing_score = writing_score

    def get_data_as_data_frame(self):
        try:
            # Note: Column names must match your training data exactly
            custom_data_input_dict = {
                "gender": [self.gender],
                "race/ethnicity": [self.race],
                "parental level of education": [self.parental_level_of_education],
                "lunch": [self.lunch],
                "test preparation course": [self.test_preparation_course],
                "reading score": [self.reading_score],
                "writing score": [self.writing_score]
            }
            
            df = pd.DataFrame(custom_data_input_dict)
            logging.info("Dataframe created successfully")
            return df
            
        except Exception as e:
            logging.error(f"Error in Dataframe Creation: {str(e)}")
            raise CustomException(e, sys)