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")


async def seed():
    await init_db()

    async with async_session() as db:
        # Check if already seeded
        result = await db.execute(select(Commune).where(Commune.slug == "saou"))
        if result.scalar_one_or_none():
            print("Saoû already seeded.")
            return

        # Create commune
        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 = set()
        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)

            household = 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,
            )
            db.add(household)

        await db.flush()

        # ── Publish a reference curve ──
        # Reference: vinf=400, all params=0.5
        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:
            # Pick a unique household
            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"],
            )
            vote = 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,
            )
            db.add(vote)
            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")


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