# model.py
import joblib
import pandas as pd
from typing import List, Dict, Union
from transformers import Pipeline

class LinearRegressionModel:
    def __init__(self):
        # Load the trained linear regression model
        self.model = joblib.load("linear_regression_model.joblib")

    def get_coefficients(self, columns: List[str]) -> Dict[str, float]:
        # Return coefficients with column names as keys
        coefficients = dict(zip(columns, self.model.coef_.tolist()))
        intercept = self.model.intercept_
        return {"coefficients": coefficients, "intercept": intercept}

    def predict_from_csv(self, columns: List[str], csv_file_path: str) -> List[float]:
        # Read the CSV file and select specified columns
        data = pd.read_csv(csv_file_path)[columns]
        return self.model.predict(data).tolist()

# Instantiate the model
model = LinearRegressionModel()

# Define a custom pipeline class for inference
class CustomLinearRegressionPipeline(Pipeline):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def __call__(self, inputs: Dict[str, Union[List[str], bytes]]) -> Dict[str, Union[Dict[str, float], List[float]]]:
        columns = inputs.get("columns", [])
        csv_file = inputs.get("csv_file", None)  # Expect a CSV file in binary format
        
        if not csv_file:  # If no CSV file, return coefficients and intercept
            return self.model.get_coefficients(columns)
        else:  # If CSV file is provided, save and process it for predictions
            csv_file_path = "/tmp/input_data.csv"  # Temporary path to save the CSV file
            with open(csv_file_path, "wb") as f:
                f.write(csv_file)

            # Make predictions from CSV file
            predictions = self.model.predict_from_csv(columns, csv_file_path)
            return {"predictions": predictions}

# Instantiate the custom pipeline
pipeline = CustomLinearRegressionPipeline(model)