import json
import warnings
import numpy as np
import pandas as pd
import joblib
from datetime import datetime, timedelta

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional
import os

warnings.filterwarnings("ignore")

# ─── App ─────────────────────────────────────────────────────────────────────
app = FastAPI(
    title="EcoSmart Energy Prediction API",
    description="AI-powered energy consumption forecasting API",
    version="2.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ─── Constants ────────────────────────────────────────────────────────────────
RATE_PER_KWH  = 1.55
FORECAST_TEMP = 20
DATA_FILE     = "EcoSmart PT3.xlsx"
MODEL_FILE    = "ecosmart_model.pkl"
OUTPUT_JSON   = "ecosmart_predictions.json"

AVG_POWER_KW = {
    "Air Conditioner": 1.5,
    "Smart Fridge":    0.15,
    "Smart Lights":    0.05,
    "Water Heater":    1.2,
    "Electronics":     0.08,
    "Other":           0.1,
}

FEATURES = [
    "hour", "is_weekend", "temperature",
    "lag_1h", "lag_24h",
    "hour_sin", "hour_cos",
    "temp_sq", "rolling_3h",
]

# ─── Global state ─────────────────────────────────────────────────────────────
_model       = None
_df          = None
_last_output = None


# ─── Schemas ──────────────────────────────────────────────────────────────────
class PredictRequest(BaseModel):
    user_id:    int
    temperature: float = FORECAST_TEMP
    hour:       Optional[int]   = None
    is_weekend: Optional[int]   = None
    lag_1h:     Optional[float] = None
    lag_24h:    Optional[float] = None


# ─── Helpers ──────────────────────────────────────────────────────────────────
def get_weekday_name(date):
    return date.strftime("%a")


def _load_model_and_data():
    """Load pre-trained model from disk + data file, then build output JSON."""
    global _model, _df, _last_output

    # 1. Load model
    if not os.path.exists(MODEL_FILE):
        raise FileNotFoundError(f"Model file '{MODEL_FILE}' not found.")
    _model = joblib.load(MODEL_FILE)

    # 2. Load data
    if not os.path.exists(DATA_FILE):
        raise FileNotFoundError(f"Data file '{DATA_FILE}' not found.")
    df = pd.read_excel(DATA_FILE, parse_dates=["timestamp"])
    df = df.sort_values("timestamp").reset_index(drop=True)

    # Feature engineering (needed for lag/rolling lookups)
    df["hour_sin"]   = np.sin(2 * np.pi * df["hour"] / 24)
    df["hour_cos"]   = np.cos(2 * np.pi * df["hour"] / 24)
    df["temp_sq"]    = df["temperature"] ** 2
    df["rolling_3h"] = df["consumption_kwh"].rolling(3, min_periods=1).mean().shift(1)
    df.bfill(inplace=True)
    df.ffill(inplace=True)
    _df = df

    # 3. Build full output
    _last_output = _build_output(user_id=1)
    with open(OUTPUT_JSON, "w") as f:
        json.dump(_last_output, f, indent=2, default=str)

    print(f"✓ Model loaded from {MODEL_FILE}")
    print(f"✓ Data loaded — {len(df)} rows")
    print(f"✓ Output JSON ready")


def _build_output(user_id: int) -> dict:
    global _model, _df

    df      = _df
    model   = _model
    now     = datetime.now()
    now_str = now.strftime("%Y-%m-%dT%H:%M:%S")
    last_time = df["timestamp"].max()
    today_str = str(last_time.date())
    avg_rate  = RATE_PER_KWH

    # ── Forecast next 24 h ────────────────────────────────────────────────────
    last_rows    = df.tail(24).copy()
    forecast_rows = []
    for h in range(24):
        ts         = last_time + timedelta(hours=h + 1)
        hour       = ts.hour
        is_weekend = int(ts.weekday() >= 5)
        lag_1h     = last_rows["consumption_kwh"].iloc[-1]
        lag_24h    = last_rows["consumption_kwh"].iloc[-24] if len(last_rows) >= 24 else lag_1h
        rolling3h  = last_rows["consumption_kwh"].iloc[-3:].mean()
        row = {
            "hour": hour, "is_weekend": is_weekend,
            "temperature": FORECAST_TEMP, "lag_1h": lag_1h, "lag_24h": lag_24h,
            "hour_sin": np.sin(2 * np.pi * hour / 24),
            "hour_cos": np.cos(2 * np.pi * hour / 24),
            "temp_sq": FORECAST_TEMP ** 2, "rolling_3h": rolling3h,
        }
        x = np.array([[row[f] for f in FEATURES]])
        predicted_kwh = round(float(model.predict(x)[0]), 3)
        forecast_rows.append({
            "timestamp": ts.strftime("%Y-%m-%dT%H:%M:%S"),
            "hour": hour, "predicted_kwh": predicted_kwh,
        })
        new_row = pd.DataFrame([{
            "timestamp": ts, "hour": hour, "consumption_kwh": predicted_kwh,
            "temperature": FORECAST_TEMP, "is_weekend": is_weekend,
            "lag_1h": lag_1h, "lag_24h": lag_24h,
            "hour_sin": row["hour_sin"], "hour_cos": row["hour_cos"],
            "temp_sq": row["temp_sq"], "rolling_3h": rolling3h,
        }])
        last_rows = pd.concat([last_rows, new_row], ignore_index=True)

    forecast_kwh = [r["predicted_kwh"] for r in forecast_rows]
    peak_idx     = int(np.argmax(forecast_kwh))

    # ── Users ─────────────────────────────────────────────────────────────────
    users = [{
        "id": user_id, "name": "EcoSmart User",
        "email": f"user{user_id}@ecosmart.ai",
        "monthly_budget": 200.0,
        "daily_target_kwh": round(df["consumption_kwh"].mean() * 24, 2),
        "created_at": now_str,
    }]

    # ── Devices ───────────────────────────────────────────────────────────────
    APP_INFO = [
        {"name": "Air Conditioner", "category": "HVAC"},
        {"name": "Smart Fridge",    "category": "Kitchen"},
        {"name": "Smart Lights",    "category": "Lighting"},
        {"name": "Water Heater",    "category": "Heating"},
        {"name": "Electronics",     "category": "Entertainment"},
    ]
    devices = [
        {"id": i+1, "user_id": user_id, "name": a["name"],
         "category": a["category"], "is_active": True, "created_at": now_str}
        for i, a in enumerate(APP_INFO)
    ]

    # ── Hourly usage (last 24 h) ──────────────────────────────────────────────
    last_24h     = df.tail(24)
    hourly_usage = [
        {"id": i+1, "user_id": user_id,
         "recorded_at": row["timestamp"].strftime("%Y-%m-%dT%H:%M:%S"),
         "hour": int(row["hour"]), "kwh": round(float(row["consumption_kwh"]), 3)}
        for i, (_, row) in enumerate(last_24h.iterrows())
    ]

    # ── Weekly / monthly ──────────────────────────────────────────────────────
    last_date = last_time.date()

    def _usage_records(days):
        start  = last_date - timedelta(days=days - 1)
        sub    = df[df["timestamp"].dt.date >= start]
        daily  = sub.groupby(sub["timestamp"].dt.date)["consumption_kwh"].sum()
        records, idx = [], 1
        for d in pd.date_range(start=start, end=last_date, freq="D"):
            kwh = round(float(daily.get(d.date(), 0.0)), 2)
            records.append({"id": idx, "user_id": user_id,
                            "date": str(d.date()), "kwh": kwh,
                            "cost": round(kwh * avg_rate, 2)})
            idx += 1
        return records

    weekly_usage  = _usage_records(7)
    monthly_usage = _usage_records(30)

    # ── Rate schedules ────────────────────────────────────────────────────────
    rate_schedules = [
        {"id": 1, "user_id": user_id, "name": "Off-Peak", "start_hour":  0, "end_hour": 12, "rate_per_kwh": 1.20},
        {"id": 2, "user_id": user_id, "name": "Peak",     "start_hour": 12, "end_hour": 20, "rate_per_kwh": 1.90},
        {"id": 3, "user_id": user_id, "name": "Evening",  "start_hour": 20, "end_hour": 24, "rate_per_kwh": 1.40},
    ]

    # ── Energy targets ────────────────────────────────────────────────────────
    today_sub        = df[df["timestamp"].dt.date == last_date]
    today_total      = round(float(today_sub["consumption_kwh"].sum()), 3)
    daily_avg_target = round(float(df["consumption_kwh"].mean() * 24), 2)
    energy_targets   = [{
        "id": 1, "user_id": user_id,
        "daily_kwh_target":   daily_avg_target,
        "monthly_kwh_target": round(daily_avg_target * 30, 2),
        "current_daily_kwh":  today_total,
        "created_at": now_str,
    }]

    # ── Alerts & Notifications ────────────────────────────────────────────────
    high_thr = float(df["consumption_kwh"].quantile(0.90))
    alerts, notifs = [], []
    for aid, (_, row) in enumerate(df[df["consumption_kwh"] >= high_thr].tail(5).iterrows(), 1):
        alerts.append({
            "id": aid, "user_id": user_id, "type": "high_usage",
            "message": f"High consumption: {row['consumption_kwh']:.2f} kWh at {row['timestamp']}",
            "kwh_value": round(float(row["consumption_kwh"]), 3),
            "timestamp": str(row["timestamp"]), "is_read": False,
        })
        notifs.append({
            "id": aid, "user_id": user_id, "title": "⚡ High Usage Alert",
            "body": f"Usage spike: {row['consumption_kwh']:.2f} kWh detected.",
            "type": "alert", "is_read": False, "created_at": str(row["timestamp"]),
        })

    # ── Appliances ────────────────────────────────────────────────────────────
    appliances = []
    portions   = [0.38, 0.18, 0.12, 0.20, 0.12]
    remaining  = today_total * 0.85
    other_kwh  = today_total * 0.15
    for i, (a, portion) in enumerate(zip(APP_INFO, portions)):
        daily_kwh = round(remaining * portion, 3)
        power_kw  = AVG_POWER_KW.get(a["name"], 0.1)
        appliances.append({
            "id": i+1, "user_id": user_id, "name": a["name"], "category": a["category"],
            "daily_kwh": daily_kwh, "daily_cost": round(daily_kwh * avg_rate, 2),
            "runtime_hours": round(daily_kwh / power_kw, 2) if power_kw > 0 else 0.0,
            "monthly_cost": round(daily_kwh * 30 * avg_rate, 2),
            "percentage": round(daily_kwh / today_total * 100, 1) if today_total > 0 else 0,
            "created_at": now_str,
        })
    if other_kwh > 0.01:
        appliances.append({
            "id": len(APP_INFO)+1, "user_id": user_id, "name": "Other Devices", "category": "Other",
            "daily_kwh": round(other_kwh, 3), "daily_cost": round(other_kwh * avg_rate, 2),
            "runtime_hours": 0.0, "monthly_cost": round(other_kwh * 30 * avg_rate, 2),
            "percentage": round(other_kwh / today_total * 100, 1) if today_total > 0 else 0,
            "created_at": now_str,
        })
    appliances.sort(key=lambda x: x["daily_kwh"], reverse=True)
    top_consumer     = appliances[0]
    appliance_summary = {
        "total_usage_today": round(sum(a["daily_kwh"] for a in appliances), 2),
        "total_cost_today":  round(sum(a["daily_cost"] for a in appliances), 2),
        "top_consumer": {
            "name": top_consumer["name"], "percentage": top_consumer["percentage"],
            "daily_kwh": top_consumer["daily_kwh"], "daily_cost": top_consumer["daily_cost"],
        },
    }

    # ── Energy savings ────────────────────────────────────────────────────────
    yest_df    = df[df["timestamp"].dt.date == (last_time.date() - timedelta(days=1))]
    yest_total = float(yest_df["consumption_kwh"].sum()) if not yest_df.empty else today_total
    saved_kwh  = max(0, yest_total - today_total)
    energy_savings = [{
        "id": 1, "user_id": user_id,
        "money_saved":        round(saved_kwh * avg_rate, 2),
        "efficiency_percent": round(saved_kwh / yest_total * 100, 1) if yest_total > 0 else 0,
        "co2_reduced_kg":     round(saved_kwh * 0.233, 2),
        "period_start": str((last_time - timedelta(days=7)).date()),
        "period_end":   today_str, "created_at": now_str,
    }]

    # ── Peak forecast & recommendations ──────────────────────────────────────
    hourly_avg_dict = df.groupby("hour")["consumption_kwh"].mean().to_dict()
    peak_hour       = forecast_rows[peak_idx]["hour"]
    peak_value      = forecast_kwh[peak_idx]
    peak_rate       = next(
        (r["rate_per_kwh"] for r in rate_schedules if r["start_hour"] <= peak_hour < r["end_hour"]),
        avg_rate,
    )
    avg_for_peak = hourly_avg_dict.get(peak_hour, peak_value)
    pct_above    = round((peak_value - avg_for_peak) / avg_for_peak * 100, 1) if avg_for_peak > 0 else 0

    recommendations = []
    if 12 <= peak_hour <= 20:
        recommendations.append({
            "id": 1, "user_id": user_id, "title": "Shift usage away from peak",
            "description": (f"High usage predicted at {peak_hour}:00 ({peak_value} kWh). "
                            "Run dishwasher, laundry, or water heater before 12 PM or after 8 PM."),
            "type": "peak_avoidance",
            "potential_savings_kwh": round(peak_value * 0.3, 2),
            "affected_device": "Water Heater, Dishwasher, Laundry", "created_at": now_str,
        })
    ac = next((a for a in appliances if a["name"] == "Air Conditioner"), None)
    if ac and ac["daily_kwh"] > 2.0:
        recommendations.append({
            "id": 2, "user_id": user_id, "title": "Pre-cool before peak hours",
            "description": (f"Your AC consumes {ac['daily_kwh']:.1f} kWh/day. "
                            "Run it 1 hour earlier to reduce peak cost."),
            "type": "load_shifting",
            "potential_savings_kwh": round(ac["daily_kwh"] * 0.15, 2),
            "affected_device": "Air Conditioner", "created_at": now_str,
        })
    if today_total > daily_avg_target:
        recommendations.append({
            "id": 3, "user_id": user_id, "title": "Reduce standby power",
            "description": (f"Today's usage ({today_total} kWh) exceeds your daily target ({daily_avg_target} kWh). "
                            "Turn off electronics when not in use."),
            "type": "conservation",
            "potential_savings_kwh": round((today_total - daily_avg_target) * 0.2, 2),
            "affected_device": "Electronics, Lights", "created_at": now_str,
        })

    peak_forecast = {
        "expected_peak_hour":   peak_hour,
        "expected_peak_kwh":    peak_value,
        "expected_peak_cost":   round(peak_value * peak_rate, 2),
        "percent_above_average": pct_above,
        "peak_rate_period":     "Peak (12-20)" if 12 <= peak_hour <= 20 else "Off-peak",
        "recommended_action":   recommendations[0]["description"] if recommendations else "",
    }

    hourly_averages = [
        {"hour": h, "avg_kwh": round(hourly_avg_dict.get(h, 0), 3)} for h in range(24)
    ]

    # ── Chart data ────────────────────────────────────────────────────────────
    def _alert_trend():
        days_order   = ["Mon","Tue","Wed","Thu","Fri","Sat","Sun"]
        high_thr_    = float(df["consumption_kwh"].quantile(0.85))
        high_counts  = {d: 0 for d in days_order}
        spike_counts = {d: 0 for d in days_order}
        for idx in range(1, len(df)):
            curr = df.iloc[idx]; prev = df.iloc[idx-1]
            dow  = get_weekday_name(curr["timestamp"])
            if curr["consumption_kwh"] >= high_thr_:
                high_counts[dow] += 1
            if prev["consumption_kwh"] > 0 and curr["consumption_kwh"] / prev["consumption_kwh"] >= 1.4:
                spike_counts[dow] += 1
        return {"x_axis": days_order, "series": [
            {"name": "high_usage", "data": [high_counts[d]  for d in days_order]},
            {"name": "spike",      "data": [spike_counts[d] for d in days_order]},
        ]}

    def _daily_chart():
        day_data = df[df["timestamp"].dt.date == last_date]
        hourly   = [round(float(day_data[day_data["hour"]==h]["consumption_kwh"].sum()), 3) for h in range(24)]
        total    = round(sum(hourly), 2)
        return {"hourly_kwh": hourly, "total_kwh": total, "cost": round(total * avg_rate, 2)}

    def _weekly_chart():
        start  = last_date - timedelta(days=6)
        sub    = df[df["timestamp"].dt.date >= start]
        daily  = sub.groupby(sub["timestamp"].dt.date)["consumption_kwh"].sum()
        drange = pd.date_range(start=start, end=last_date, freq="D")
        vals   = [round(float(daily.get(d.date(), 0.0)), 2) for d in drange]
        return {"x_axis": [d.strftime("%a") for d in drange], "kwh": vals,
                "cost": [round(v * avg_rate, 2) for v in vals],
                "dates": [d.strftime("%Y-%m-%d") for d in drange]}

    def _monthly_chart():
        start  = last_date - timedelta(days=29)
        sub    = df[df["timestamp"].dt.date >= start]
        daily  = sub.groupby(sub["timestamp"].dt.date)["consumption_kwh"].sum()
        drange = pd.date_range(start=start, end=last_date, freq="D")
        vals   = [round(float(daily.get(d.date(), 0.0)), 2) for d in drange]
        return {"x_axis": [d.strftime("%Y-%m-%d") for d in drange], "kwh": vals,
                "cost": [round(v * avg_rate, 2) for v in vals]}

    chart_data = {
        "alert_volume_trend": _alert_trend(),
        "daily_usage":        _daily_chart(),
        "weekly_usage":       _weekly_chart(),
        "monthly_usage":      _monthly_chart(),
    }

    # ── Final assembly ────────────────────────────────────────────────────────
    return {
        "_meta": {"generated_at": now_str, "model_version": "2.0.0", "user_id": user_id},
        "users":             users,
        "peak_forecast":     peak_forecast,
        "hourly_averages":   hourly_averages,
        "recommendations":   recommendations,
        "appliance_summary": appliance_summary,
        "devices":           devices,
        "hourly_usage":      hourly_usage,
        "daily_usage":       weekly_usage,
        "monthly_usage":     monthly_usage,
        "rate_schedules":    rate_schedules,
        "energy_targets":    energy_targets,
        "energy_savings":    energy_savings,
        "alerts":            alerts,
        "notifications":     notifs,
        "appliances":        appliances,
        "forecast_next_24h": [
            {"id": i+1, "user_id": user_id, "recorded_at": r["timestamp"],
             "hour": r["hour"], "kwh": r["predicted_kwh"]}
            for i, r in enumerate(forecast_rows)
        ],
        "chart_data": chart_data,
    }


# ─── Startup ──────────────────────────────────────────────────────────────────
@app.on_event("startup")
async def startup_event():
    try:
        _load_model_and_data()
    except Exception as e:
        print(f"⚠ Startup error: {e}")


# ─── Routes ───────────────────────────────────────────────────────────────────

@app.get("/")
def root():
    return {
        "status":       "online",
        "service":      "EcoSmart Energy Prediction API",
        "version":      "2.0.0",
        "model_loaded": _model is not None,
        "docs":         "/docs",
    }


@app.get("/health")
def health():
    return {"status": "ok", "model_ready": _model is not None}


@app.post("/predict")
def predict(req: PredictRequest):
    """Single-row prediction for a user."""
    if _model is None or _df is None:
        raise HTTPException(503, "Model not ready.")

    now        = datetime.now()
    hour       = req.hour       if req.hour       is not None else now.hour
    is_wk      = req.is_weekend if req.is_weekend is not None else int(now.weekday() >= 5)
    last_rows  = _df.tail(24)
    lag_1h     = req.lag_1h  if req.lag_1h  is not None else float(last_rows["consumption_kwh"].iloc[-1])
    lag_24h    = req.lag_24h if req.lag_24h is not None else float(last_rows["consumption_kwh"].iloc[0])
    rolling3h  = float(last_rows["consumption_kwh"].iloc[-3:].mean())

    x = np.array([[
        hour, is_wk, req.temperature,
        lag_1h, lag_24h,
        np.sin(2 * np.pi * hour / 24),
        np.cos(2 * np.pi * hour / 24),
        req.temperature ** 2,
        rolling3h,
    ]])

    predicted_kwh = round(float(_model.predict(x)[0]), 4)
    rate_period   = "Peak" if 12 <= hour < 20 else ("Off-Peak" if hour < 12 else "Evening")
    rate_map      = {"Peak": 1.90, "Off-Peak": 1.20, "Evening": 1.40}

    return {
        "user_id":       req.user_id,
        "predicted_kwh": predicted_kwh,
        "hour":          hour,
        "timestamp":     now.strftime("%Y-%m-%dT%H:%M:%S"),
        "cost_estimate": round(predicted_kwh * rate_map[rate_period], 4),
        "rate_period":   rate_period,
    }


@app.get("/forecast/{user_id}")
def forecast(user_id: int):
    """24-hour energy forecast for a user."""
    if _last_output is None:
        raise HTTPException(503, "Model not ready.")
    return [{**item, "user_id": user_id} for item in _last_output["forecast_next_24h"]]


@app.get("/full-output/{user_id}")
def full_output(user_id: int):
    """Full JSON payload — all tables with user_id injected."""
    if _model is None or _df is None:
        raise HTTPException(503, "Model not ready.")
    return _build_output(user_id=user_id)


@app.get("/summary/{user_id}")
def summary(user_id: int):
    """Lightweight dashboard summary for a user."""
    if _last_output is None:
        raise HTTPException(503, "Model not ready.")
    o  = _last_output
    et = o["energy_targets"][0] if o["energy_targets"] else {}
    pf = o["peak_forecast"]
    es = o["energy_savings"][0] if o["energy_savings"] else {}
    return {
        "user_id":               user_id,
        "today_kwh":             et.get("current_daily_kwh", 0),
        "daily_target_kwh":      et.get("daily_kwh_target", 0),
        "today_cost":            o["appliance_summary"].get("total_cost_today", 0),
        "peak_hour":             pf.get("expected_peak_hour"),
        "peak_kwh":              pf.get("expected_peak_kwh"),
        "top_appliance":         o["appliance_summary"].get("top_consumer", {}).get("name"),
        "recommendations_count": len(o["recommendations"]),
        "money_saved":           es.get("money_saved", 0),
        "efficiency_percent":    es.get("efficiency_percent", 0),
        "co2_reduced_kg":        es.get("co2_reduced_kg", 0),
        "alerts_count":          len(o["alerts"]),
    }


@app.get("/devices/{user_id}")
def devices(user_id: int):
    if _last_output is None:
        raise HTTPException(503, "Model not ready.")
    return [{**d, "user_id": user_id} for d in _last_output["devices"]]


@app.get("/appliances/{user_id}")
def appliances(user_id: int):
    if _last_output is None:
        raise HTTPException(503, "Model not ready.")
    return [{**a, "user_id": user_id} for a in _last_output["appliances"]]


@app.get("/alerts/{user_id}")
def alerts(user_id: int):
    if _last_output is None:
        raise HTTPException(503, "Model not ready.")
    return [{**a, "user_id": user_id} for a in _last_output["alerts"]]


@app.get("/notifications/{user_id}")
def notifications(user_id: int):
    if _last_output is None:
        raise HTTPException(503, "Model not ready.")
    return [{**n, "user_id": user_id} for n in _last_output["notifications"]]


@app.get("/chart-data/{user_id}")
def chart_data(user_id: int):
    if _last_output is None:
        raise HTTPException(503, "Model not ready.")
    return {**_last_output["chart_data"], "user_id": user_id}


if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=False)