"""
Pricing computation for Bézier tariff model.

Ported from eau.py:120-167 (NewModel.updateComputation).
Pure Python + numpy, no matplotlib.
"""

import numpy as np
from dataclasses import dataclass

from app.engine.integrals import compute_integrals


@dataclass
class HouseholdData:
    """Minimal household data needed for computation."""
    volume_m3: float
    status: str  # "RS", "RP", or "PRO"
    price_paid_eur: float = 0.0


@dataclass
class TariffResult:
    """Result of a full tariff computation."""
    p0: float
    curve_volumes: list[float]
    curve_prices_m3: list[float]
    curve_bills_rp: list[float]
    curve_bills_rs: list[float]
    household_bills: list[float]  # projected bill for each household


@dataclass
class ImpactRow:
    """Price impact for a specific volume level."""
    volume: float
    old_price: float
    new_price_rp: float
    new_price_rs: float


def compute_p0(
    households: list[HouseholdData],
    recettes: float,
    abop: float,
    abos: float,
    vinf: float,
    vmax: float,
    pmax: float,
    a: float,
    b: float,
    c: float,
    d: float,
    e: float,
) -> float:
    """
    Compute p0 (inflection price) that balances total revenue.

    p0 = (R - Σ(abo + β₂)) / Σ(α₁ + α₂)
    """
    total_abo = 0.0
    total_alpha = 0.0
    total_beta = 0.0

    for h in households:
        abo = abos if h.status == "RS" else abop
        total_abo += abo

        vol = max(h.volume_m3, 1e-5)  # avoid div by 0
        alpha1, alpha2, beta2 = compute_integrals(vol, vinf, vmax, pmax, a, b, c, d, e)
        total_alpha += alpha1 + alpha2
        total_beta += beta2

    if total_abo >= recettes:
        return 0.0

    if total_alpha == 0:
        return 0.0

    return (recettes - total_abo - total_beta) / total_alpha


def compute_tariff(
    households: list[HouseholdData],
    recettes: float,
    abop: float,
    abos: float,
    vinf: float,
    vmax: float,
    pmax: float,
    a: float,
    b: float,
    c: float,
    d: float,
    e: float,
    nbpts: int = 200,
) -> TariffResult:
    """
    Full tariff computation: p0, price curves, and per-household bills.
    """
    p0 = compute_p0(households, recettes, abop, abos, vinf, vmax, pmax, a, b, c, d, e)

    # Generate curve points
    tt = np.linspace(0, 1 - 1e-6, nbpts)

    # Tier 1 volumes and prices
    vv1 = vinf * ((1 - 3 * b) * tt**3 + 3 * b * tt**2)
    prix_m3_1 = p0 * ((3 * a - 2) * tt**3 + (-6 * a + 3) * tt**2 + 3 * a * tt)

    # Tier 2 volumes and prices
    vv2 = vinf + (vmax - vinf) * (
        (3 * (c + d - c * d) - 2) * tt**3
        + 3 * (1 - 2 * c - d + c * d) * tt**2
        + 3 * c * tt
    )
    prix_m3_2 = p0 + (pmax - p0) * ((1 - 3 * e) * tt**3 + 3 * e * tt**2)

    vv = np.concatenate([vv1, vv2])
    prix_m3 = np.concatenate([prix_m3_1, prix_m3_2])

    # Compute full bills (integral) for each curve point
    alpha1_arr = np.zeros(len(vv))
    alpha2_arr = np.zeros(len(vv))
    beta2_arr = np.zeros(len(vv))
    for iv, v in enumerate(vv):
        alpha1_arr[iv], alpha2_arr[iv], beta2_arr[iv] = compute_integrals(
            v, vinf, vmax, pmax, a, b, c, d, e
        )

    bills_rp = abop + (alpha1_arr + alpha2_arr) * p0 + beta2_arr
    bills_rs = abos + (alpha1_arr + alpha2_arr) * p0 + beta2_arr

    # Per-household projected bills
    household_bills = []
    for h in households:
        vol = max(h.volume_m3, 1e-5)
        abo = abos if h.status == "RS" else abop
        a1, a2, b2 = compute_integrals(vol, vinf, vmax, pmax, a, b, c, d, e)
        household_bills.append(abo + (a1 + a2) * p0 + b2)

    def sanitize(arr):
        return [0.0 if (x != x or x == float('inf') or x == float('-inf')) else float(x) for x in arr]

    return TariffResult(
        p0=p0,
        curve_volumes=sanitize(vv),
        curve_prices_m3=sanitize(prix_m3),
        curve_bills_rp=sanitize(bills_rp),
        curve_bills_rs=sanitize(bills_rs),
        household_bills=sanitize(household_bills),
    )


def compute_impacts(
    households: list[HouseholdData],
    recettes: float,
    abop: float,
    abos: float,
    vinf: float,
    vmax: float,
    pmax: float,
    a: float,
    b: float,
    c: float,
    d: float,
    e: float,
    reference_volumes: list[float] | None = None,
) -> tuple[float, list[ImpactRow]]:
    """
    Compute p0 and price impacts for reference volume levels.

    Returns (p0, list of ImpactRow).
    """
    if reference_volumes is None:
        reference_volumes = [30, 60, 90, 150, 300, 600]

    p0 = compute_p0(households, recettes, abop, abos, vinf, vmax, pmax, a, b, c, d, e)

    # Compute average 2018 price per m³ for a rough "old price" baseline
    total_vol = sum(max(h.volume_m3, 1e-5) for h in households)
    total_abo_old = sum(abos if h.status == "RS" else abop for h in households)
    old_p_m3 = (recettes - total_abo_old) / total_vol if total_vol > 0 else 0

    impacts = []
    for vol in reference_volumes:
        # Old price (linear model)
        old_price_rp = abop + old_p_m3 * vol
        # New price
        a1, a2, b2 = compute_integrals(vol, vinf, vmax, pmax, a, b, c, d, e)
        new_price_rp = abop + (a1 + a2) * p0 + b2
        new_price_rs = abos + (a1 + a2) * p0 + b2
        impacts.append(ImpactRow(
            volume=vol,
            old_price=old_price_rp,
            new_price_rp=new_price_rp,
            new_price_rs=new_price_rs,
        ))

    return p0, impacts