"""The deliberately small model catalogue, and the release-boundary panel.

Two things live here and they must never be mixed.

**Release assessment** is a cradle-to-release figure about a *model release*:
research, data preparation, failed runs, training, tuning, evaluation,
checkpoints, storage, networks, packaging, release and allocated hardware.  It
explicitly excludes inference.  Values may only be entered here when they come
from a published, model-specific assessment, used as published.

**Use-phase energy** is what this network actually measures: the joules a
single inference drew on one community worker.  It lives in
``distinct_server.ui`` because it is per-run and live, and it is outside the
release boundary entirely.

The two are different functional units.  They are never added, never averaged,
never graded against each other, and never rendered in the same total.
``combined_total`` exists solely to raise if some future caller tries.
"""

from __future__ import annotations

from collections.abc import Iterable, Mapping, Sequence
from dataclasses import dataclass
from enum import Enum
from types import MappingProxyType
from typing import Optional

from distinct_protocol import ModelManifest

RELEASE_BOUNDARY = (
    "Cradle to release. Counts the whole programme: research, data preparation, failed "
    "runs, training, tuning, evaluation, checkpoints, storage, networks, packaging, "
    "release and allocated hardware. Excludes inference and downstream use."
)

USE_PHASE_BOUNDARY = (
    "Use phase only: one inference on one community worker. Outside the cradle-to-release "
    "boundary of the environmental rubric. Not comparable with, and never added to, a "
    "release result."
)

# The rubric's six areas, in its order.  A model block is understandable in
# isolation, so every area is always present -- absence is stated, not omitted.
RUBRIC_AREAS: tuple[str, ...] = (
    "energy",
    "climate",
    "water",
    "land",
    "materials",
    "pollution",
)

AREA_LABELS: Mapping[str, str] = MappingProxyType(
    {
        "energy": "Energy",
        "climate": "Climate",
        "water": "Water",
        "land": "Land / biodiversity",
        "materials": "Materials / hardware",
        "pollution": "Pollution / waste",
    }
)


class AreaStatus(str, Enum):
    """The rubric's status vocabulary, verbatim.

    ``MISSING`` means no usable evidence.  It is not zero, not an F, and not a
    grade.  Rendering it as a number is a bug, and ``AreaResult`` refuses to
    hold one.
    """

    MEASURED = "Measured"
    MODELLED = "Modelled"
    PROXY_BASED = "Proxy-based"
    MISSING = "Missing"


@dataclass(frozen=True)
class AreaResult:
    """One area of one release assessment.

    The constructor enforces the rubric's two hard rules structurally rather
    than by convention: a ``Missing`` area cannot carry a number, and a
    non-``Missing`` area cannot carry a number without a primary source.
    """

    status: AreaStatus
    value: Optional[str] = None
    unit: Optional[str] = None
    note: str = ""
    source_url: Optional[str] = None

    def __post_init__(self) -> None:
        if not isinstance(self.status, AreaStatus):
            object.__setattr__(self, "status", AreaStatus(self.status))
        if self.status is AreaStatus.MISSING:
            if self.value is not None or self.unit is not None:
                raise ValueError("a Missing area cannot carry a value; Missing is not zero")
        else:
            if not self.value:
                raise ValueError(f"a {self.status.value} area requires a value")
            if not self.source_url:
                raise ValueError(
                    f"a {self.status.value} area requires a source_url; every displayed "
                    "number carries its primary source"
                )

    @property
    def display(self) -> str:
        """The exact text to render. Never returns '0' for a Missing area."""

        if self.status is AreaStatus.MISSING:
            return "Missing"
        unit = f" {self.unit}" if self.unit else ""
        return f"{self.value}{unit} ({self.status.value})"

    def to_dict(self) -> dict:
        return {
            "status": self.status.value,
            "value": self.value,
            "unit": self.unit,
            "note": self.note,
            "source_url": self.source_url,
            "display": self.display,
        }


def _all_missing(reason: str) -> Mapping[str, AreaResult]:
    return MappingProxyType(
        {area: AreaResult(AreaStatus.MISSING, note=reason) for area in RUBRIC_AREAS}
    )


@dataclass(frozen=True)
class ReleaseAssessment:
    """A cradle-to-release assessment of one model release.

    This used to carry an ``evidence_grade`` letter alongside the areas, and a
    sibling ordinal for impact.  Both are gone.  A letter is a ranking this
    catalogue invented on top of numbers somebody else published, and once the
    published numbers are shown in full -- with their units, their method and
    their primary source -- the letter adds no information and one more way to
    be wrong.  What replaced it is ``coverage``, which is not a score: it counts
    how many of the six areas have a published result and says nothing about
    whether those results are large or small.
    """

    model_id: str
    coverage: str
    areas: Mapping[str, AreaResult]
    summary: str = ""
    covers: str = ""
    weights_sha256: str = ""
    lineage: str = ""
    boundary: str = RELEASE_BOUNDARY
    reason: str = ""
    links: Sequence[Mapping[str, str]] = ()

    def __post_init__(self) -> None:
        missing_areas = set(RUBRIC_AREAS) - set(self.areas)
        if missing_areas:
            raise ValueError(f"release assessment omits areas: {sorted(missing_areas)}")
        if self.coverage not in {"Complete", "Partial", "Insufficient", "No data"}:
            raise ValueError("coverage must be Complete, Partial, Insufficient or No data")
        if self.coverage == "No data" and self.has_any_result:
            raise ValueError("'No data' cannot carry a published result")
        if self.coverage != "No data" and not self.has_any_result:
            raise ValueError("a coverage other than 'No data' requires at least one result")
        for link in self.links:
            if not link.get("label") or not link.get("url"):
                raise ValueError("every link needs both a label and a url")
        object.__setattr__(self, "areas", MappingProxyType(dict(self.areas)))
        object.__setattr__(self, "links", tuple(MappingProxyType(dict(link)) for link in self.links))

    @property
    def has_any_result(self) -> bool:
        return any(item.status is not AreaStatus.MISSING for item in self.areas.values())

    @property
    def assessed_areas(self) -> int:
        return sum(1 for item in self.areas.values() if item.status is not AreaStatus.MISSING)

    @property
    def coverage_display(self) -> str:
        """Coverage and its boundary, inseparable.

        Coverage on its own invites the reading that four of six is a middling
        score.  It is not a score at all, so the count always travels with the
        word that says what it counts.
        """

        if self.coverage == "No data":
            return "no published result"
        return f"{self.coverage.lower()} · {self.assessed_areas} of {len(RUBRIC_AREAS)} areas published"

    def to_dict(self) -> dict:
        return {
            "model_id": self.model_id,
            "coverage": self.coverage,
            "coverage_display": self.coverage_display,
            "assessed_areas": self.assessed_areas,
            "total_areas": len(RUBRIC_AREAS),
            "summary": self.summary,
            "covers": self.covers,
            "boundary": self.boundary,
            "weights_sha256": self.weights_sha256,
            "lineage": self.lineage,
            "reason": self.reason,
            "links": [dict(link) for link in self.links],
            "areas": {area: item.to_dict() for area, item in self.areas.items()},
        }


# LABELS NAME THE PUBLICATION, NOT THE SIZE.
#
# These used to read "smallest", "balanced test model" and "largest test
# model". Size is not something this catalogue evaluates, and putting it in the
# name of the thing being chosen invited exactly the reasoning this surface
# exists to refuse: that a smaller file is a better environmental choice. It
# might be, and it might not, and the only way to know is a published result.
# What belongs in a label is who published the weights, because that is what
# the reader has to be able to check.
#
# WHY THESE THREE. Every model here has a published, model-specific assessment
# of its training. Models with none were removed rather than listed with an
# empty record; see the note in ``distinct_agent.models`` for what went and why.
MODEL_CATALOG: dict[str, ModelManifest] = {
    # DEFAULT of the worker catalogue: the lowest published figures of any
    # assessed release at a size this network can run. Kept in step with
    # distinct_agent.models.KNOWN_MODEL_MANIFESTS; the drift test in
    # tests/test_catalog.py pins filename, RAM and context length.
    "olmoe-1b-7b-0924-instruct": ModelManifest(
        id="olmoe-1b-7b-0924-instruct",
        label="OLMoE 1B-7B 0924 Instruct · Ai2 first-party GGUF",
        filename="olmoe-1b-7b-0924-instruct-q4_k_m.gguf",
        min_ram_gb=8.0,
        context_length=4_096,
    ),
    "olmo-2-1124-7b-instruct": ModelManifest(
        id="olmo-2-1124-7b-instruct",
        label="OLMo 2 1124 7B Instruct · Ai2 first-party GGUF",
        filename="olmo-2-1124-7B-instruct-Q4_K_M.gguf",
        min_ram_gb=8.0,
        context_length=4_096,
    ),
}


# THE ONE ROUTE, AND WHAT IT NOW LETS THROUGH.
#
# The rule has not changed: a published, model-specific result, used exactly as
# published. GPU-hours, FLOPs, tokens, parameters, TDP, PUE, grid factors and
# another model's assessment must not be converted here into invented MWh,
# tonnes or litres. The contextual route that once allowed an ordinal judgement
# from "disclosed scale, hardware, location and controls" stays deleted, because
# it was caught reasoning "a small 0.49B checkpoint moderates demand" into a
# grade, which is the parameters-to-impact conversion wearing a letter.
#
# What changed is that the catalogue now holds models this rule has something to
# say about. Every entry below carries real published figures instead of the
# uniform "No data" this structure used to render, and the detail -- the method
# behind each number, the primary source, the plain-language summary -- lives in
# ``model_assessments.json`` and is rendered from there by ``distinct_server.ui``.
# This structure is the in-catalogue record: enough to be correct on its own if
# that file is absent, never the fuller surface.
#
# WHAT IS STILL FORBIDDEN. GPU-hours are not energy. A publisher who reports
# hours on hardware of a stated wattage has not reported a measurement, and
# multiplying the two here would be this surface inventing the number it
# exists to report. An area with no published figure stays Missing.

#: The document each Ai2 figure is quoted from. Named once so a citation can
#: never drift from the number beside it.
_AI2_ASSESSMENT = "https://arxiv.org/abs/2503.05804"

_NO_LAND = AreaResult(
    AreaStatus.MISSING,
    note=(
        "No published per-model land or biodiversity figure exists for any model, from "
        "any lab. ITU-T L.1801 lists the area and states the methodology is developing."
    ),
)

RELEASE_ASSESSMENT: dict[str, ReleaseAssessment] = {
    "olmoe-1b-7b-0924-instruct": ReleaseAssessment(
        model_id="olmoe-1b-7b-0924-instruct",
        coverage="Partial",
        covers=(
            "The 0924 pretraining run. The tuning that produced this Instruct checkpoint "
            "is not separately costed in the source."
        ),
        summary=(
            "The lowest published footprint of any assessed release at a size this network "
            "can run. Ai2 sampled real GPU power at sub-second intervals rather than "
            "assuming chips draw their rated wattage, which makes the energy figure one of "
            "very few genuinely measured training numbers in print. Carbon and water follow "
            "from that measured energy using the site's own grid factor and cooling "
            "efficiency, so they describe Texas as much as they describe the model."
        ),
        areas={
            "energy": AreaResult(
                AreaStatus.MEASURED,
                value="54",
                unit="MWh",
                note="Final pretraining run, Jupiter cluster, Texas.",
                source_url=_AI2_ASSESSMENT,
            ),
            "climate": AreaResult(
                AreaStatus.MODELLED,
                value="18",
                unit="tCO2e",
                note=(
                    "Measured energy at 0.332 kg CO2 per kWh, PUE 1.2. Location-based; "
                    "no offset applied."
                ),
                source_url=_AI2_ASSESSMENT,
            ),
            "water": AreaResult(
                AreaStatus.MODELLED,
                value="70",
                unit="kL",
                note=(
                    "Measured energy at a water usage effectiveness of 1.29 L per kWh. "
                    "On-site cooling only; upstream power-station water is not counted."
                ),
                source_url=_AI2_ASSESSMENT,
            ),
            "land": _NO_LAND,
            "materials": AreaResult(
                AreaStatus.MISSING,
                note=(
                    "Ai2 report 22 tCO2e and 4.8 kL of embodied hardware impact across their "
                    "whole programme. That is a programme total, not this release's share."
                ),
            ),
            "pollution": AreaResult(AreaStatus.MISSING, note="Not assessed in the source."),
        },
        links=(
            {
                "label": "Model card",
                "url": "https://huggingface.co/allenai/OLMoE-1B-7B-0924-Instruct",
            },
            {
                "label": "Pinned weights",
                "url": "https://huggingface.co/allenai/OLMoE-1B-7B-0924-Instruct-GGUF",
            },
            {"label": "Published assessment", "url": _AI2_ASSESSMENT},
        ),
    ),
    "olmo-2-1124-7b-instruct": ReleaseAssessment(
        model_id="olmo-2-1124-7b-instruct",
        coverage="Partial",
        covers=(
            "The 1124 pretraining run. The tuning that produced this Instruct checkpoint "
            "is not separately costed in the source."
        ),
        summary=(
            "A dense 7B measured by the same team, in the same data centre, in the same "
            "year, by the same method as OLMoE. That makes the pair one of the few honest "
            "like-for-like comparisons in this field: about three times the energy, carbon "
            "and water for comparable capability. Offered as the fallback if the "
            "mixture-of-experts architecture gives trouble locally, not because it is lighter."
        ),
        areas={
            "energy": AreaResult(
                AreaStatus.MEASURED,
                value="157",
                unit="MWh",
                note="Final pretraining run, Jupiter cluster, Texas.",
                source_url=_AI2_ASSESSMENT,
            ),
            "climate": AreaResult(
                AreaStatus.MODELLED,
                value="52",
                unit="tCO2e",
                note=(
                    "Measured energy at 0.332 kg CO2 per kWh, PUE 1.2. Location-based; "
                    "no offset applied."
                ),
                source_url=_AI2_ASSESSMENT,
            ),
            "water": AreaResult(
                AreaStatus.MODELLED,
                value="202",
                unit="kL",
                note=(
                    "Measured energy at 1.29 L per kWh. The 13B sibling drank four times as "
                    "much on the same architecture because it trained in Iowa, where the "
                    "figure is 3.1 L per kWh: water tracks the site more than the model."
                ),
                source_url=_AI2_ASSESSMENT,
            ),
            "land": _NO_LAND,
            "materials": AreaResult(
                AreaStatus.MISSING,
                note=(
                    "Ai2 report 22 tCO2e and 4.8 kL of embodied hardware impact across their "
                    "whole programme. That is a programme total, not this release's share."
                ),
            ),
            "pollution": AreaResult(AreaStatus.MISSING, note="Not assessed in the source."),
        },
        links=(
            {
                "label": "Model card",
                "url": "https://huggingface.co/allenai/OLMo-2-1124-7B-Instruct",
            },
            {
                "label": "Pinned weights",
                "url": "https://huggingface.co/allenai/OLMo-2-1124-7B-Instruct-GGUF",
            },
            {"label": "Published assessment", "url": _AI2_ASSESSMENT},
        ),
    ),
}


class BoundaryError(ValueError):
    """Raised when two incompatible measurement boundaries would be combined."""


def combined_total(*, release_joules: float, use_phase_joules: float) -> float:
    """Refuse, always.

    This function exists to be called by a future refactor that has forgotten
    why the two panels are separate, and to fail loudly at that moment. A
    release figure is cradle-to-release and excludes inference; a per-run
    figure is inference and nothing else. There is no sum of the two that
    means anything.
    """

    raise BoundaryError(
        "release and use-phase energy are different functional units on different "
        "boundaries and cannot be summed; render them in separate panels"
    )


def model_choices() -> list[tuple[str, str]]:
    """Plain model labels, with no figure attached.

    This used to append ``needs N GB RAM``, which was a category error rather
    than a formatting choice. The person choosing a model on this server is not
    the person whose machine runs it: the worker operator already decided at
    setup which models their hardware can host, and the server only ever offers
    what a worker has actually granted. Telling a user that a model "needs 8 GB"
    invited them to reason about someone else's RAM, which they cannot see and
    are not responsible for.

    What belongs beside a model here is its assessment, and that is added by
    :func:`distinct_server.ui.model_choices_with_grades`, which has access to
    the assessment file. This function stays the plain source of truth for
    ``label -> id`` so the drift test has something stable to pin.
    """

    return [(model.label, model.id) for model in MODEL_CATALOG.values()]


def release_assessment(model_id: str) -> ReleaseAssessment:
    return RELEASE_ASSESSMENT[model_id]


def public_catalog() -> Iterable[dict]:
    for model in MODEL_CATALOG.values():
        value = model.to_dict()
        value["weights_identity"] = model.identity_status
        value["release_assessment"] = RELEASE_ASSESSMENT[model.id].to_dict()
        value["release_boundary"] = RELEASE_BOUNDARY
        value["use_phase_boundary"] = USE_PHASE_BOUNDARY
        yield value