import numpy as np
import pandas as pd
from typing import List, Dict
from datetime import datetime, timedelta
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import PolynomialFeatures

class APMAgent:
    """AI agent for EV battery asset performance management."""

    def __init__(self):
        self.soh_threshold_low = 75.0
        self.temp_threshold = 45.0
        self.max_cycles = 2500

    def compute_soh(self, df: pd.DataFrame) -> float:
        """Estimate current SOH from BMS history using empirical degradation model."""
        if df.empty:
            return 100.0
        latest = df.iloc[-1]
        cycles = latest['cycle_count']
        avg_temp = df['temperature'].mean()
        # Empirical Arrhenius-like degradation: faster at high temp + cycle aging
        temp_factor = 1 + max(0, avg_temp - 25) * 0.008
        soh = 100 * np.exp(-cycles / (self.max_cycles / temp_factor))
        # Penalize high C-rate / fast charging events (current > 200A)
        fast_charge_events = (df['current'] > 200).sum()
        soh -= fast_charge_events * 0.08
        return max(60.0, round(soh, 2))

    def forecast_rul(self, df: pd.DataFrame, vehicle_id: str) -> Dict:
        """Predict remaining useful life in cycles and days."""
        if len(df) < 5:
            return self._fallback(df, vehicle_id)

        soh_now = self.compute_soh(df)

        # Build polynomial regression on cycle_count vs SOH
        X = df[['cycle_count']].values
        y = [self.compute_soh(df.iloc[:i+1]) for i in range(len(df))]

        poly = PolynomialFeatures(degree=2)
        X_poly = poly.fit_transform(X)
        model = LinearRegression()
        model.fit(X_poly, y)

        future_cycles = np.arange(int(df['cycle_count'].max()), int(df['cycle_count'].max()) + 1000, 10)
        future_X = poly.transform(future_cycles.reshape(-1, 1))
        predicted_soh = model.predict(future_X)

        # Find when SOH hits 70% end-of-life
        below_eol = np.where(predicted_soh < 70.0)[0]
        if len(below_eol) > 0:
            eol_cycle = int(future_cycles[below_eol[0]])
        else:
            eol_cycle = int(df['cycle_count'].max()) + 500

        rul_cycles = max(0, eol_cycle - int(df['cycle_count'].max()))
        # Assume ~1.5 cycles per day for industrial EV
        rul_days = int(rul_cycles / 1.5)

        risk_level = self._risk_level(soh_now, rul_days)
        recommendations = self._recommendations(df, soh_now, rul_days)

        return {
            "vehicle_id": vehicle_id,
            "current_soh": soh_now,
            "predicted_rul_cycles": rul_cycles,
            "predicted_rul_days": rul_days,
            "risk_level": risk_level,
            "recommendations": recommendations,
            "soh_trajectory": predicted_soh[:50].tolist()
        }

    def _fallback(self, df: pd.DataFrame, vehicle_id: str) -> Dict:
        soh = self.compute_soh(df) if not df.empty else 95.0
        return {
            "vehicle_id": vehicle_id,
            "current_soh": soh,
            "predicted_rul_cycles": 1200,
            "predicted_rul_days": 800,
            "risk_level": "LOW",
            "recommendations": ["Collect more telemetry for higher-confidence RUL prediction."],
            "soh_trajectory": []
        }

    def _risk_level(self, soh: float, rul_days: int) -> str:
        if soh < self.soh_threshold_low or rul_days < 90:
            return "CRITICAL"
        if soh < 85 or rul_days < 180:
            return "HIGH"
        if soh < 92 or rul_days < 365:
            return "MEDIUM"
        return "LOW"

    def _recommendations(self, df: pd.DataFrame, soh: float, rul_days: int) -> List[str]:
        recs = []
        if soh < self.soh_threshold_low:
            recs.append("Schedule battery pack inspection or replacement within 30 days.")
        if rul_days < 180:
            recs.append("High priority: plan secondary powertrain or procurement.")
        avg_temp = df['temperature'].mean()
        if avg_temp > self.temp_threshold:
            recs.append("Thermal stress detected — review cooling system and charging schedule.")
        fast_charges = (df['current'] > 200).sum()
        if fast_charges > len(df) * 0.15:
            recs.append("Reduce frequent fast charging to slow degradation.")
        if not recs:
            recs.append("Asset healthy. Continue scheduled maintenance.")
        return recs

    def detect_anomalies(self, df: pd.DataFrame) -> List[Dict]:
        """Detect thermal and electrical anomalies."""
        alerts = []
        for _, row in df.iterrows():
            if row['temperature'] > 50:
                alerts.append({
                    "timestamp": row['timestamp'].isoformat(),
                    "vehicle_id": row['vehicle_id'],
                    "type": "THERMAL_CRITICAL",
                    "value": row['temperature'],
                    "message": f"Battery temperature {row['temperature']}°C exceeds critical threshold."
                })
            elif row['temperature'] > self.temp_threshold:
                alerts.append({
                    "timestamp": row['timestamp'].isoformat(),
                    "vehicle_id": row['vehicle_id'],
                    "type": "THERMAL_WARNING",
                    "value": row['temperature'],
                    "message": f"Battery temperature {row['temperature']}°C above recommended range."
                })
            if row['voltage'] < 2.8 * 96:
                alerts.append({
                    "timestamp": row['timestamp'].isoformat(),
                    "vehicle_id": row['vehicle_id'],
                    "type": "LOW_VOLTAGE",
                    "value": row['voltage'],
                    "message": "Pack voltage below safe operating limit."
                })
        return alerts[-20:]  # return latest 20