sejeteralo/backend/seed.py

"""Seed the database with Saoû data from Eau2018.xls."""

import asyncio
import sys
import os
import random
from datetime import datetime

sys.path.insert(0, os.path.dirname(__file__))

import xlrd
from sqlalchemy import select

from app.database import engine, async_session, init_db
from app.models import Commune, TariffParams, Household, AdminUser, Vote
from app.services.auth_service import hash_password
from app.services.import_service import generate_auth_code
from app.engine.pricing import HouseholdData, compute_p0

XLS_PATH = os.path.join(os.path.dirname(__file__), "Eau2018.xls")
if not os.path.exists(XLS_PATH):
    XLS_PATH = os.path.join(os.path.dirname(__file__), "..", "Eau2018.xls")


# Codes fixes — identiques dans le dev hint frontend
DEV_FIXTURES = [
    {"identifier": "[DEV] Foyer RS 60m³",   "status": "RS",  "volume_m3": 60.0,  "price_paid_eur": 140.0, "auth_code": "DEVTEST2"},
    {"identifier": "[DEV] Foyer RP 120m³",  "status": "RP",  "volume_m3": 120.0, "price_paid_eur": 265.0, "auth_code": "DEVTEST3"},
    {"identifier": "[DEV] Foyer PRO 350m³", "status": "PRO", "volume_m3": 350.0, "price_paid_eur": 680.0, "auth_code": "DEVTEST4"},
]


async def seed():
    await init_db()

    async with async_session() as db:
        # ── Commune Saoû (idempotent) ──────────────────────────────────────────
        result = await db.execute(select(Commune).where(Commune.slug == "saou"))
        commune = result.scalar_one_or_none()

        if commune is None:
            commune = Commune(
                name="Saoû",
                slug="saou",
                description="Commune de Saoû - Tarification progressive de l'eau",
            )
            db.add(commune)
            await db.flush()

            # Create tariff params
            params = TariffParams(
                commune_id=commune.id,
                abop=100,
                abos=100,
                recettes=75000,
                pmax=20,
                vmax=2100,
                differentiated_tariff=False,
                data_year=2018,
                data_imported_at=datetime.utcnow(),
            )
            db.add(params)

            # Create super admin (manages all communes)
            super_admin = AdminUser(
                email="superadmin@sejeteralo.fr",
                hashed_password=hash_password("superadmin"),
                full_name="Super Admin",
                role="super_admin",
            )
            db.add(super_admin)

            # Create commune admin for Saoû (manages only this commune)
            commune_admin = AdminUser(
                email="saou@sejeteralo.fr",
                hashed_password=hash_password("saou2024"),
                full_name="Admin Saoû",
                role="commune_admin",
            )
            commune_admin.communes.append(commune)
            db.add(commune_admin)

            # Import households from Eau2018.xls
            book = xlrd.open_workbook(XLS_PATH)
            sheet = book.sheet_by_name("CALCULS")
            nb_hab = 363

            existing_codes = {f["auth_code"] for f in DEV_FIXTURES}
            for r in range(1, nb_hab + 1):
                name = sheet.cell_value(r, 0)
                status = sheet.cell_value(r, 3)
                volume = sheet.cell_value(r, 4)
                price = sheet.cell_value(r, 33)

                code = generate_auth_code()
                while code in existing_codes:
                    code = generate_auth_code()
                existing_codes.add(code)

                db.add(Household(
                    commune_id=commune.id,
                    identifier=str(name).strip(),
                    status=str(status).strip().upper(),
                    volume_m3=float(volume),
                    price_paid_eur=float(price) if price else 0.0,
                    auth_code=code,
                ))

            await db.flush()

            # ── Publish a reference curve ──
            ref_vinf, ref_a, ref_b, ref_c, ref_d, ref_e = 400, 0.5, 0.5, 0.5, 0.5, 0.5

            hh_result = await db.execute(
                select(Household).where(Household.commune_id == commune.id)
            )
            all_households = hh_result.scalars().all()
            hh_data = [
                HouseholdData(volume_m3=h.volume_m3, status=h.status, price_paid_eur=h.price_paid_eur)
                for h in all_households
            ]

            ref_p0 = compute_p0(
                hh_data,
                recettes=params.recettes, abop=params.abop, abos=params.abos,
                vinf=ref_vinf, vmax=params.vmax, pmax=params.pmax,
                a=ref_a, b=ref_b, c=ref_c, d=ref_d, e=ref_e,
            )
            params.published_vinf = ref_vinf
            params.published_a = ref_a
            params.published_b = ref_b
            params.published_c = ref_c
            params.published_d = ref_d
            params.published_e = ref_e
            params.published_p0 = ref_p0
            params.published_at = datetime.utcnow()

            # ── Generate 10 votes, small variations around the reference ──
            random.seed(42)

            vote_profiles = [
                # 5 votes slightly below reference (eco-leaning)
                {"vinf": 350, "a": 0.45, "b": 0.48, "c": 0.40, "d": 0.52, "e": 0.55},
                {"vinf": 370, "a": 0.42, "b": 0.50, "c": 0.45, "d": 0.48, "e": 0.52},
                {"vinf": 380, "a": 0.48, "b": 0.45, "c": 0.42, "d": 0.50, "e": 0.58},
                {"vinf": 360, "a": 0.50, "b": 0.52, "c": 0.38, "d": 0.55, "e": 0.50},
                {"vinf": 390, "a": 0.47, "b": 0.47, "c": 0.48, "d": 0.46, "e": 0.53},
                # 5 votes slightly above reference (lax-leaning)
                {"vinf": 420, "a": 0.52, "b": 0.50, "c": 0.55, "d": 0.48, "e": 0.45},
                {"vinf": 440, "a": 0.55, "b": 0.53, "c": 0.52, "d": 0.50, "e": 0.42},
                {"vinf": 430, "a": 0.50, "b": 0.55, "c": 0.58, "d": 0.45, "e": 0.48},
                {"vinf": 410, "a": 0.53, "b": 0.48, "c": 0.50, "d": 0.52, "e": 0.47},
                {"vinf": 450, "a": 0.48, "b": 0.52, "c": 0.60, "d": 0.42, "e": 0.40},
            ]

            used_households = set()
            vote_count = 0
            for prof in vote_profiles:
                hh_pick = random.choice(all_households)
                while hh_pick.id in used_households:
                    hh_pick = random.choice(all_households)
                used_households.add(hh_pick.id)

                vp0 = compute_p0(
                    hh_data,
                    recettes=params.recettes, abop=params.abop, abos=params.abos,
                    vinf=prof["vinf"], vmax=params.vmax, pmax=params.pmax,
                    a=prof["a"], b=prof["b"], c=prof["c"], d=prof["d"], e=prof["e"],
                )
                db.add(Vote(
                    commune_id=commune.id,
                    household_id=hh_pick.id,
                    vinf=prof["vinf"],
                    a=prof["a"], b=prof["b"], c=prof["c"], d=prof["d"], e=prof["e"],
                    computed_p0=vp0,
                ))
                hh_pick.has_voted = True
                vote_count += 1

            await db.commit()
            print(f"Seeded: commune 'saou', {nb_hab} households, {vote_count} votes")
            print(f"  Published curve: vinf={ref_vinf}, p0={ref_p0:.3f}")
            print(f"  Super admin: superadmin@sejeteralo.fr / superadmin")
            print(f"  Commune admin Saou: saou@sejeteralo.fr / saou2024")
        else:
            print("Saoû already seeded.")

        # ── Dev fixtures (idempotent — insérés à chaque run si absents) ────────
        fixture_added = 0
        for fixture in DEV_FIXTURES:
            res = await db.execute(
                select(Household).where(Household.auth_code == fixture["auth_code"])
            )
            if res.scalar_one_or_none() is None:
                db.add(Household(commune_id=commune.id, **fixture))
                fixture_added += 1
        if fixture_added:
            await db.commit()
        print(f"  Dev fixtures: {fixture_added} ajoutés — DEVTEST2 (RS 60m³) · DEVTEST3 (RP 120m³) · DEVTEST4 (PRO 350m³)")


if __name__ == "__main__":
    asyncio.run(seed())