"""
GridGuard AI — Synthetic Data Generation
=========================================
Generates a synthetic distribution-grid dataset that mimics a Nigerian DISCO
feeder structure:

    Substation -> Transformers -> Neighbourhoods -> Households -> 12 months of readings

A subset of households are seeded as "bypass/fraud" cases: from a random
month onward, their RECORDED (metered) consumption drops sharply while their
TRUE (actual) consumption stays close to normal — i.e. they keep using power
but stop paying for most of it. This creates the energy-balance gap at the
transformer level that the model is designed to detect.

Run this cell first in Colab, then run 02_train_model.py.
Outputs (written to ./data/):
    households.csv        - static household attributes
    readings_monthly.csv  - long-format household x month consumption
    transformers.csv      - transformer roll-up (injected vs metered energy)
"""

import numpy as np
import pandas as pd
import os

RNG_SEED = 42
rng = np.random.default_rng(RNG_SEED)

N_TRANSFORMERS = 12
HOUSEHOLDS_PER_TRANSFORMER = (40, 90)   # min, max households served per transformer
N_MONTHS = 12
FRAUD_RATE = 0.07                       # ~7% of households are bypassing
TECH_LOSS_RATE = 0.06                   # normal technical loss baseline (6%)

OUT_DIR = os.path.join(os.path.dirname(__file__), "data")
os.makedirs(OUT_DIR, exist_ok=True)

INCOME_CLASSES = ["low", "mid", "high"]
INCOME_WEIGHTS = [0.45, 0.40, 0.15]
INCOME_BASE_KWH = {"low": 90, "mid": 220, "high": 480}   # baseline monthly kWh
TARIFF_BAND = {"low": "B", "mid": "B", "high": "A"}

months = pd.date_range("2025-01-01", periods=N_MONTHS, freq="MS").strftime("%Y-%m")

households = []
readings = []
hh_id_counter = 1

for t_id in range(1, N_TRANSFORMERS + 1):
    n_hh = rng.integers(*HOUSEHOLDS_PER_TRANSFORMER)
    # Each transformer feeds 1-3 neighbourhoods (clusters of similar socioeconomic status)
    n_neighbourhoods = rng.integers(1, 4)
    neighbourhood_ids = [f"T{t_id:02d}-N{n+1}" for n in range(n_neighbourhoods)]
    # bias each neighbourhood toward a dominant income class (drives peer similarity)
    neighbourhood_income_bias = {
        nb: rng.choice(INCOME_CLASSES, p=INCOME_WEIGHTS) for nb in neighbourhood_ids
    }

    for _ in range(n_hh):
        nb = rng.choice(neighbourhood_ids)
        dominant_class = neighbourhood_income_bias[nb]
        # 80% of households in a neighbourhood match its dominant class, 20% differ
        income_class = dominant_class if rng.random() < 0.8 else rng.choice(INCOME_CLASSES, p=INCOME_WEIGHTS)

        house_size = {
            "low": rng.normal(45, 8), "mid": rng.normal(90, 15), "high": rng.normal(200, 35)
        }[income_class]
        house_size = max(20, house_size)
        occupants = max(1, int(rng.normal({"low": 5, "mid": 4, "high": 4}[income_class], 1.3)))
        appliance_score = np.clip(
            rng.normal({"low": 2.5, "mid": 5.0, "high": 8.0}[income_class], 1.2), 1, 10
        )

        hh_id = f"HH{hh_id_counter:05d}"
        hh_id_counter += 1

        is_fraud = rng.random() < FRAUD_RATE
        fraud_start_month = int(rng.integers(3, N_MONTHS - 1)) if is_fraud else None
        bypass_severity = float(rng.uniform(0.45, 0.85)) if is_fraud else 0.0  # fraction of true usage hidden

        # Genuine low consumers: real households (small family, frequently away,
        # energy-efficient) whose usage is low FROM THE START — no sudden drop,
        # no transformer-level loss contribution. These exist to make the
        # detection problem realistic: low consumption alone must NOT be enough
        # to flag fraud (this is the nuance the model has to learn).
        is_genuine_low = (not is_fraud) and (rng.random() < 0.05)
        # Lifestyle-change households: a real drop in TRUE consumption (e.g. fewer
        # occupants, moved out for part of year) with NO theft — metered tracks
        # true exactly, transformer loss stays normal. These overlap with fraud's
        # drop_ratio/trend_slope signature on purpose, so flagging can't rely on
        # consumption drop alone and must lean on the transformer/peer signals too.
        is_lifestyle_change = (not is_fraud) and (not is_genuine_low) and (rng.random() < 0.04)
        lifestyle_change_month = int(rng.integers(3, N_MONTHS - 1)) if is_lifestyle_change else None
        lifestyle_drop_severity = float(rng.uniform(0.3, 0.6)) if is_lifestyle_change else 0.0

        households.append({
            "household_id": hh_id,
            "transformer_id": f"T{t_id:02d}",
            "neighbourhood_id": nb,
            "income_class": income_class,
            "tariff_band": TARIFF_BAND[income_class],
            "house_size_sqm": round(house_size, 1),
            "occupants": occupants,
            "appliance_score": round(appliance_score, 2),
            "is_fraud_ground_truth": is_fraud,         # hidden label, used only for evaluation
            "is_genuine_low_consumer": is_genuine_low,  # hidden label, used only for evaluation
            "is_lifestyle_change": is_lifestyle_change,  # hidden label, used only for evaluation
            "fraud_start_month_idx": fraud_start_month,
            "bypass_severity": round(bypass_severity, 3),
        })

        # baseline monthly true consumption (kWh) with seasonality + noise
        base = INCOME_BASE_KWH[income_class] * (1 + 0.15 * (appliance_score - 5) / 5)
        base *= (0.85 + 0.05 * occupants)
        if is_genuine_low:
            base *= rng.uniform(0.3, 0.5)  # genuinely low usage, consistent across all months
        seasonal = 1 + 0.12 * np.sin(np.linspace(0, 2 * np.pi, N_MONTHS))  # dry-season AC/fan bump
        for m_idx in range(N_MONTHS):
            true_kwh = max(15, base * seasonal[m_idx] * rng.normal(1.0, 0.07))
            if is_lifestyle_change and m_idx >= lifestyle_change_month:
                true_kwh *= (1 - lifestyle_drop_severity)
            if is_fraud and m_idx >= fraud_start_month:
                metered_kwh = true_kwh * (1 - bypass_severity)
            else:
                metered_kwh = true_kwh
            readings.append({
                "household_id": hh_id,
                "month": months[m_idx],
                "month_idx": m_idx,
                "true_kwh": round(true_kwh, 2),
                "metered_kwh": round(metered_kwh, 2),
            })

households_df = pd.DataFrame(households)
readings_df = pd.DataFrame(readings)

# ---- Transformer-level rollup: injected energy vs sum of metered readings ----
roll = readings_df.merge(households_df[["household_id", "transformer_id"]], on="household_id")
tx_monthly = roll.groupby(["transformer_id", "month", "month_idx"]).agg(
    sum_true_kwh=("true_kwh", "sum"),
    sum_metered_kwh=("metered_kwh", "sum"),
).reset_index()
# injected energy = true consumption + normal technical losses (line/transformer losses)
tx_monthly["technical_loss_kwh"] = tx_monthly["sum_true_kwh"] * TECH_LOSS_RATE * rng.normal(1, 0.1, len(tx_monthly)).clip(0.7, 1.3)
tx_monthly["energy_injected_kwh"] = tx_monthly["sum_true_kwh"] + tx_monthly["technical_loss_kwh"]
tx_monthly["unaccounted_kwh"] = tx_monthly["energy_injected_kwh"] - tx_monthly["sum_metered_kwh"]
tx_monthly["loss_pct"] = tx_monthly["unaccounted_kwh"] / tx_monthly["energy_injected_kwh"]

households_df.to_csv(os.path.join(OUT_DIR, "households.csv"), index=False)
readings_df.to_csv(os.path.join(OUT_DIR, "readings_monthly.csv"), index=False)
tx_monthly.to_csv(os.path.join(OUT_DIR, "transformers.csv"), index=False)

print(f"Households generated : {len(households_df)}")
print(f"Fraud (ground truth) : {households_df['is_fraud_ground_truth'].sum()} "
      f"({households_df['is_fraud_ground_truth'].mean()*100:.1f}%)")
print(f"Transformers          : {N_TRANSFORMERS}")
print(f"Avg transformer loss% (latest month): "
      f"{tx_monthly[tx_monthly.month_idx==N_MONTHS-1]['loss_pct'].mean()*100:.1f}%")
print(f"Files written to: {OUT_DIR}")