import numpy as np
import pandas as pd
from typing import List, Dict

class FleetReadinessAgent:
    """AI agent for fleet electrification readiness & procurement intelligence."""

    def __init__(self):
        # OEM catalogue: recommended EVs for industrial/commercial segments
        self.ev_catalogue = [
            {"model": "Tata Ultra EV 7", "oem": "Tata Motors", "category": "Freight", "range_km": 200, "battery_kwh": 168, "price_inr_lakh": 45, "payload_tons": 7},
            {"model": "Ashok Leyland AVTR EV", "oem": "Ashok Leyland", "category": "Freight", "range_km": 300, "battery_kwh": 250, "price_inr_lakh": 75, "payload_tons": 16},
            {"model": "Mahindra Treo Zor", "oem": "Mahindra Electric", "category": "Last Mile", "range_km": 130, "battery_kwh": 48, "price_inr_lakh": 12, "payload_tons": 0.5},
            {"model": "Piaggio Ape E-City", "oem": "Piaggio", "category": "Last Mile", "range_km": 90, "battery_kwh": 26, "price_inr_lakh": 6, "payload_tons": 0.3},
            {"model": "JBM ECO-LIFE", "oem": "JBM Auto", "category": "Bus", "range_km": 250, "battery_kwh": 200, "price_inr_lakh": 120, "payload_tons": 12},
            {"model": " Volvo FM Electric", "oem": "Volvo Group", "category": "Construction", "range_km": 320, "battery_kwh": 450, "price_inr_lakh": 250, "payload_tons": 25},
        ]

    def score_asset(self, asset: Dict) -> Dict:
        """Compute Electrification Readiness Index (ERI) 0-100."""
        route_fit = self._route_fit(asset['route_distance_km'])
        payload_fit = self._payload_fit(asset['payload_tons'], asset['category'])
        duty_fit = self._duty_fit(asset['duty_cycle_hrs'], asset['dwell_time_hrs'])
        terrain_fit = self._terrain_fit(asset['terrain_score'])
        age_penalty = max(0, 1 - asset['age_years'] / 15)

        # Weighted score
        score = (
            route_fit * 0.30 +
            payload_fit * 0.20 +
            duty_fit * 0.25 +
            terrain_fit * 0.15 +
            age_penalty * 0.10
        ) * 100

        confidence = min(0.95, 0.6 + (len([route_fit, payload_fit, duty_fit, terrain_fit]) * 0.08))

        recommended = self._match_ev(asset)
        tco_savings = self._estimate_tco_savings(asset, recommended)
        co2_reduction = self._estimate_co2_reduction(asset)

        reasoning = [
            f"Route distance {asset['route_distance_km']} km fits EV range with {route_fit:.0%} confidence.",
            f"Payload/category compatibility: {payload_fit:.0%}.",
            f"Duty cycle {asset['duty_cycle_hrs']}h vs dwell {asset['dwell_time_hrs']}h gives {duty_fit:.0%} charging feasibility.",
            f"Terrain difficulty score {asset['terrain_score']} maps to {terrain_fit:.0%} energy efficiency."
        ]

        return {
            "asset_id": asset['asset_id'],
            "readiness_score": round(score, 1),
            "confidence": round(confidence, 2),
            "recommended_ev": recommended['model'],
            "oem": recommended['oem'],
            "battery_capacity_kwh": recommended['battery_kwh'],
            "range_km": recommended['range_km'],
            "tco_savings_inr_lakh": round(tco_savings, 2),
            "payback_months": round(recommended['price_inr_lakh'] / (tco_savings / 12 + 0.01), 1),
            "co2_reduction_tons_yr": round(co2_reduction, 2),
            "reasoning": reasoning
        }

    def _route_fit(self, distance_km: float) -> float:
        # Ideal route < 150 km per shift
        if distance_km <= 100:
            return 1.0
        if distance_km <= 200:
            return 0.85
        if distance_km <= 300:
            return 0.60
        return 0.35

    def _payload_fit(self, payload_tons: float, category: str) -> float:
        if category.lower() in ["last mile", "intra-plant"]:
            return 1.0 if payload_tons <= 1 else 0.8
        if payload_tons <= 7:
            return 0.9
        if payload_tons <= 16:
            return 0.75
        if payload_tons <= 25:
            return 0.55
        return 0.35

    def _duty_fit(self, duty_hrs: float, dwell_hrs: float) -> float:
        if duty_hrs <= 8 and dwell_hrs >= 8:
            return 1.0
        if duty_hrs <= 12 and dwell_hrs >= 4:
            return 0.8
        if duty_hrs <= 16 and dwell_hrs >= 2:
            return 0.55
        return 0.3

    def _terrain_fit(self, terrain_score: float) -> float:
        # terrain_score 1=flat, 5=extreme hills/mining
        return max(0.2, 1 - (terrain_score - 1) * 0.2)

    def _match_ev(self, asset: Dict) -> Dict:
        candidates = [ev for ev in self.ev_catalogue if ev['payload_tons'] >= asset['payload_tons']]
        if not candidates:
            candidates = self.ev_catalogue
        # Prefer lower price with sufficient range
        best = min(candidates, key=lambda x: x['price_inr_lakh'] / (x['range_km'] + 1))
        return best

    def _estimate_tco_savings(self, asset: Dict, ev: Dict) -> float:
        diesel_cost_per_km = asset['diesel_l_per_100km'] * 90 / 100  # INR/km at Rs 90/l
        electricity_cost_per_km = ev['battery_kwh'] * 9 / ev['range_km']  # INR/km at Rs 9/kWh
        annual_km = asset['route_distance_km'] * 250  # ~250 operating days
        annual_fuel_savings = (diesel_cost_per_km - electricity_cost_per_km) * annual_km / 100000  # INR lakh
        maintenance_savings = 0.5  # INR lakh/year rough estimate
        return annual_fuel_savings + maintenance_savings

    def _estimate_co2_reduction(self, asset: Dict) -> float:
        annual_km = asset['route_distance_km'] * 250
        diesel_co2_kg_per_km = asset['diesel_l_per_100km'] * 2.68 / 100
        ev_co2_kg_per_km = 0.05  # India grid factor approx
        return (diesel_co2_kg_per_km - ev_co2_kg_per_km) * annual_km / 1000

    def score_fleet(self, assets: List[Dict]) -> Dict:
        results = [self.score_asset(a) for a in assets]
        ready = [r for r in results if r['readiness_score'] >= 70]
        return {
            "total_assets": len(results),
            "ready_assets": len(ready),
            "avg_readiness": round(np.mean([r['readiness_score'] for r in results]), 1),
            "total_potential_savings_inr_lakh": round(sum(r['tco_savings_inr_lakh'] for r in ready), 2),
            "total_co2_reduction_tons_yr": round(sum(r['co2_reduction_tons_yr'] for r in ready), 2),
            "asset_scores": results
        }