import pandas as pd
import torch
from transformers import TimeSeriesTransformerModel, TimeSeriesTransformerConfig

class AIEngine:
    def __init__(self):
        config = TimeSeriesTransformerConfig(
            input_size=4,
            output_size=1,
            context_length=10,
            prediction_length=1
        )
        self.model = TimeSeriesTransformerModel(config)
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)

    def preprocess(self, df: pd.DataFrame) -> torch.Tensor:
        data = df[["SolarGen(kWh)", "WindGen(kWh)", "Tilt(°)", "Vibration(g)"]].values
        return torch.tensor(data, dtype=torch.float32).unsqueeze(0).to(self.device)

    def predict_health(self, df: pd.DataFrame) -> pd.DataFrame:
        input_tensor = self.preprocess(df)
        with torch.no_grad():
            predictions = self.model(input_tensor).logits.squeeze().cpu().numpy()
        df["HealthScore"] = predictions
        df["ML_Anomaly"] = df["HealthScore"].apply(lambda x: "Risk" if x < 0.5 else "Normal")
        return df