model.py

import os
import sys
import psycopg2
import pandas as pd
from sklearn.ensemble import RandomForestRegressor
import joblib
import numpy as np
from datetime import datetime, timedelta

def fetch_user_data(user_id: str) -> pd.DataFrame:
    conn = psycopg2.connect(
        dbname=os.getenv("DB_NAME", "energy_data"),
        user=os.getenv("DB_USER", "divyanshagrawal"),
        password=os.getenv("DB_PASS", ""),
        host=os.getenv("DB_HOST", "localhost"),
        port=os.getenv("DB_PORT", "5432"),
    )
    query = """
        SELECT 
            ts AS timestamp,
            consumption_kwh AS consumption,
            EXTRACT(MONTH FROM ts) AS month,
            EXTRACT(DOW FROM ts) AS day_of_week,
            EXTRACT(HOUR FROM ts) AS hour
        FROM meter_readings
        WHERE user_id = %s
        ORDER BY ts
    """
    df = pd.read_sql(query, conn, params=(user_id,))
    conn.close()
    return df

def train_user_model(df, user_id):
    X = df[["month", "day_of_week", "hour"]]
    y = df["consumption"]

    model = RandomForestRegressor(n_estimators=100, random_state=42)
    model.fit(X, y)

    filename = f"user_model_{user_id}.pkl"
    joblib.dump(model, filename)
    print(f"User model saved as {filename}")
    return model

def hybrid_predict(user_id, month, day_of_week, hour):
    # Load base model
    base_model = joblib.load("base_model.pkl")

    # Try loading user model
    try:
        user_model = joblib.load(f"user_model_{user_id}.pkl")
        user_data_available = True
    except FileNotFoundError:
        print("No user model found, using only base model")
        user_model = None
        user_data_available = False

    # Build base-model input using the model's own expected feature names and order
    # Handles odd spellings like 'day_of _week' and arbitrary ordering
    def _normalize(name: str) -> str:
        return ''.join(ch for ch in name.lower() if ch.isalnum() or ch == '_').replace('__', '_')

    # Defaults for features we don't have live values for
    defaults = {
        'temperature_c': 20.0,
        'lag1': 2.0,
        'lag4': 2.0,
        'lag96': 2.0,
    }

    # Canonical inputs we do have
    available = {
        'month': float(month),
        'day_of_week': float(day_of_week),
        'day_of_week_with_space': float(day_of_week),  # for 'day_of _week'
        'hour': float(hour),
    }

    # Try to read expected feature names from the model
    expected_features = None
    if hasattr(base_model, 'feature_names_in_'):
        try:
            expected_features = list(base_model.feature_names_in_)
        except Exception:
            expected_features = None
    # XGBoost legacy path
    if expected_features is None and hasattr(base_model, 'get_booster'):
        try:
            booster = base_model.get_booster()
            expected_features = booster.feature_names
        except Exception:
            expected_features = None

    # Fallback to a sensible default order if the model doesn't expose names
    if expected_features is None:
        expected_features = ["month", "day_of_week", "hour", "temperature_c", "lag1", "lag4", "lag96"]

    # Map normalized names to values
    value_by_norm = {}
    for key, val in available.items():
        value_by_norm[_normalize(key)] = val
    for key, val in defaults.items():
        value_by_norm[_normalize(key)] = val

    # Construct the row in the exact expected order and with the exact column names
    row_values = []
    for feat in expected_features:
        norm = _normalize(feat)
        # Special-case common typos
        if norm == _normalize('day_of _week'):
            norm = _normalize('day_of_week')
        val = value_by_norm.get(norm)
        if val is None:
            # Unknown feature → use 0.0 as safe fallback
            val = 0.0
        row_values.append(val)

    X_base = pd.DataFrame([row_values], columns=expected_features)

    # Create DataFrame for user model using its expected feature names/order
    X_user = None
    if user_data_available:
        user_expected = None
        if hasattr(user_model, 'feature_names_in_'):
            try:
                user_expected = list(user_model.feature_names_in_)
            except Exception:
                user_expected = None
        if user_expected is None:
            user_expected = ["month", "day_of_week", "hour"]

        def _normalize(name: str) -> str:
            return ''.join(ch for ch in name.lower() if ch.isalnum() or ch == '_').replace('__', '_')

        provided = {
            'month': float(month),
            'day_of_week': float(day_of_week),
            'day_of_week_with_space': float(day_of_week),  # handles 'day_of _week'
            'hour': float(hour),
        }
        norm_to_val = { _normalize(k): v for k, v in provided.items() }

        user_row = []
        for feat in user_expected:
            norm = _normalize(feat)
            if norm == _normalize('day_of _week'):
                norm = _normalize('day_of_week')
            val = norm_to_val.get(norm, 0.0)
            user_row.append(val)
        X_user = pd.DataFrame([user_row], columns=user_expected)

    # Get base model prediction
    base_pred = base_model.predict(X_base)[0]

    if user_data_available:
        user_pred = user_model.predict(X_user)[0]

        # weight based on how much user data exists
        weight = min(0.8, len(fetch_user_data(user_id)) / 1000.0)  
        final_pred = (1 - weight) * base_pred + weight * user_pred
        return final_pred
    else:
        return base_pred


def forecast_next24_series(user_id: str, start_dt: datetime | None = None) -> list:
    """Generate 24 hourly forecasts using the base model's expected features.
    Maintains lag features across steps and adds gentle noise to temperature.
    Returns list of dicts: {timestamp, hour, kwh}
    """
    base_model = joblib.load("base_model.pkl")

    # Helper to normalize names (match hybrid_predict logic)
    def _normalize(name: str) -> str:
        return ''.join(ch for ch in name.lower() if ch.isalnum() or ch == '_').replace('__', '_')

    # Discover expected features from the base model
    expected_features = None
    if hasattr(base_model, 'feature_names_in_'):
        try:
            expected_features = list(base_model.feature_names_in_)
        except Exception:
            expected_features = None
    if expected_features is None and hasattr(base_model, 'get_booster'):
        try:
            expected_features = base_model.get_booster().feature_names
        except Exception:
            expected_features = None
    if expected_features is None:
        expected_features = ["month", "day_of_week", "hour", "temperature_c", "lag1", "lag4", "lag96"]

    # Seed state from latest DB row if available
    last_ts = None
    last_temp = 20.0
    lag1 = 2.0
    lag4 = 2.0
    lag96 = 2.0
    try:
        conn = psycopg2.connect(
            dbname=os.getenv("DB_NAME", "energy_data"),
            user=os.getenv("DB_USER", "admin"),
            password=os.getenv("DB_PASS", "admin123"),
            host=os.getenv("DB_HOST", "localhost"),
            port=os.getenv("DB_PORT", "5432"),
        )
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT ts, consumption_kwh
                FROM meter_readings
                WHERE user_id = %s
                ORDER BY ts DESC
                LIMIT 1
                """,
                (user_id,),
            )
            row = cur.fetchone()
            if row:
                last_ts = row[0]
                lag1 = float(row[1])
        conn.close()
    except Exception:
        pass

    # Use provided start time or derive from last_ts or now
    if start_dt is None:
        if last_ts is not None:
            try:
                start_dt = last_ts.replace(minute=0, second=0, microsecond=0) + timedelta(hours=1)
            except Exception:
                start_dt = datetime.utcnow().replace(minute=0, second=0, microsecond=0)
        else:
            start_dt = datetime.utcnow().replace(minute=0, second=0, microsecond=0)

    series = []
    for i in range(24):
        ts = start_dt + timedelta(hours=i)
        month = ts.month
        dow = ts.weekday()  # 0=Mon
        hour = ts.hour

        # Build value map for all potentially needed features
        values = {
            'month': float(month),
            'day_of_week': float(dow),
            'day_of_week_with_space': float(dow),
            'hour': float(hour),
            'temperature_c': float(last_temp),
            'lag1': float(lag1),
            'lag4': float(lag4),
            'lag96': float(lag96),
        }

        # Assemble in expected order
        row_vals = []
        for feat in expected_features:
            norm = _normalize(feat)
            if norm == _normalize('day_of _week'):
                norm = _normalize('day_of_week')
            row_vals.append(float(values.get(norm, 0.0)))
        X = pd.DataFrame([row_vals], columns=expected_features)

        y = float(base_model.predict(X)[0])
        series.append({"timestamp": ts.isoformat(), "hour": hour, "kwh": y})

        # Update lags for next step
        lag96 = lag4
        lag4 = lag1
        lag1 = y
        # Drift temperature slightly
        try:
            last_temp = float(last_temp) + float(np.random.normal(0, 0.2))
        except Exception:
            last_temp = float(last_temp)

    return series


if __name__ == "__main__":
    # Entrypoint for retraining job
    user_id = os.getenv("USER_ID", "user_001")
    print(f"[MODEL] Starting retrain for user_id={user_id}")
    try:
        df = fetch_user_data(user_id)
        if df.empty:
            print(f"[MODEL] No data for {user_id}. Skipping training.")
            sys.exit(0)
        model = train_user_model(df, user_id)
        print(f"[MODEL] Training complete for {user_id}")
        sys.exit(0)
    except Exception as e:
        print(f"[MODEL] Training failed: {e}")
        sys.exit(1)