#!/usr/bin/env python3
"""Build the screening dataset used by the AI Cost Index paper.

The inputs are deliberately explicit. They are country-level proxies, not supplier
quotes. Replace them with site bids before an investment decision.
"""

from __future__ import annotations

import csv
import json
from pathlib import Path


OUT = Path(__file__).resolve().parent
SECONDS_YEAR = 365 * 24 * 60 * 60
EUR_USD = 1.1467  # ECB reference rate, 16 July 2026
DATA_VERSION = "2026-07"
RELEASE_DATE = "2026-07-27"

# Kimi K3 price-implied calibration, 22 July 2026. Kimi publishes $15 per
# million output tokens, $3 per million cache-miss input tokens and $0.30 per
# million cache-hit input tokens, with a reported cache-hit rate above 90% for
# coding workloads. We use exactly 90% and assume the API price covers at least
# 75% of production cost, corresponding to no more than a 25% subsidy.
KIMI_K3_API = {
    "output_usd_per_m": 15.0,
    "cache_miss_input_usd_per_m": 3.0,
    "cache_hit_input_usd_per_m": .30,
    "cache_hit_rate": .90,
    "general_input_output_ratio": 4.0,
    "maximum_subsidy_share": .25,
    "source": "https://www.kimi.com/blog/kimi-k3",
}

# Cost and output assumptions are for one reference 72-GPU rack-scale deployment.
# Production throughput is measured under the workload and service requirements.
WORKLOADS = {
    "general": {
        "label": "General inference",
        "unit": "$/million output tokens",
        "output_tps": 900,
        "it_kw": 98,
        "utilization": 0.70,
        "extra_opex_m": 0.02,
    },
    "long_context": {
        "label": "Long-context analysis",
        "unit": "$/million output tokens",
        "output_tps": 420,
        "it_kw": 124,
        "utilization": 0.70,
        "extra_opex_m": 0.06,
    },
    "agentic": {
        "label": "Agentic engineering",
        "unit": "$/million output tokens",
        "output_tps": 190,
        "it_kw": 132,
        "utilization": 0.70,
        "extra_opex_m": 0.10,
    },
}

SCALES = {
    "datacenter_100mw": {
        "label": "100 MW data center",
        "description": "100 MW utility-connected data center",
    },
    "onprem_250kw": {
        "label": "On-prem solution (250 kW)",
        "description": "On-premises deployment with a 250 kW connection envelope",
    },
}

# code, country, group, delivered electricity EUR/kWh, PUE, landed-hardware
# multiplier, nominal hurdle rate, facility capex EURm, grid/connection capex EURm,
# operations labour multiplier, technical availability, evidence grade.
C = [
    ("AT", "Austria", "EU", .1986, 1.13, 1.03, .070, .72, .30, 1.05, .994, "B"),
    ("BE", "Belgium", "EU", .1866, 1.12, 1.03, .072, .76, .34, 1.05, .994, "B"),
    ("BG", "Bulgaria", "EU", .1413, 1.17, 1.08, .095, .55, .24, .62, .988, "B"),
    ("HR", "Croatia", "EU", .1548, 1.18, 1.08, .090, .58, .25, .72, .989, "B"),
    ("CY", "Cyprus", "EU", .2429, 1.25, 1.12, .090, .70, .38, .80, .986, "B"),
    ("CZ", "Czechia", "EU", .1825, 1.13, 1.05, .080, .60, .28, .78, .992, "B"),
    ("DK", "Denmark", "EU", .1216, 1.10, 1.01, .065, .72, .35, 1.08, .996, "B"),
    ("EE", "Estonia", "EU", .1410, 1.10, 1.06, .080, .57, .27, .75, .991, "B"),
    ("FI", "Finland", "EU", .0748, 1.08, 1.01, .062, .61, .27, .95, .996, "B"),
    ("FR", "France", "EU", .1534, 1.13, 1.02, .068, .73, .36, 1.00, .995, "B"),
    ("DE", "Germany", "EU", .2264, 1.12, 1.02, .068, .78, .48, 1.08, .995, "B"),
    ("GR", "Greece", "EU", .1740, 1.21, 1.08, .090, .62, .27, .72, .988, "B"),
    ("HU", "Hungary", "EU", .2132, 1.15, 1.07, .090, .57, .25, .68, .990, "B"),
    ("IE", "Ireland", "EU", .2552, 1.11, 1.03, .070, .79, .60, 1.08, .994, "B"),
    ("IT", "Italy", "EU", .2203, 1.18, 1.06, .080, .70, .36, .88, .992, "B"),
    ("LV", "Latvia", "EU", .1361, 1.11, 1.07, .085, .55, .25, .68, .989, "B"),
    ("LT", "Lithuania", "EU", .1591, 1.11, 1.06, .082, .55, .26, .68, .990, "B"),
    ("LU", "Luxembourg", "EU", .1717, 1.12, 1.03, .065, .78, .33, 1.12, .995, "B"),
    ("MT", "Malta", "EU", .1352, 1.24, 1.12, .085, .72, .36, .82, .987, "B"),
    ("NL", "Netherlands", "EU", .1991, 1.11, 1.02, .067, .78, .54, 1.06, .995, "B"),
    ("PL", "Poland", "EU", .1915, 1.13, 1.06, .085, .60, .31, .70, .992, "B"),
    ("PT", "Portugal", "EU", .1329, 1.18, 1.06, .078, .62, .28, .76, .991, "B"),
    ("RO", "Romania", "EU", .1887, 1.16, 1.08, .098, .56, .26, .62, .988, "B"),
    ("SK", "Slovakia", "EU", .2090, 1.14, 1.06, .085, .58, .27, .72, .991, "B"),
    ("SI", "Slovenia", "EU", .1503, 1.14, 1.05, .078, .60, .27, .80, .992, "B"),
    ("ES", "Spain", "EU", .1324, 1.18, 1.04, .075, .66, .32, .78, .993, "B"),
    ("SE", "Sweden", "EU", .0970, 1.09, 1.01, .062, .63, .29, 1.02, .996, "B"),
    ("IS", "Iceland", "Europe", .0750, 1.08, 1.09, .075, .67, .30, 1.00, .992, "C"),
    ("NO", "Norway", "Europe", .0827, 1.09, 1.02, .063, .66, .30, 1.10, .996, "B"),
    ("CH", "Switzerland", "Europe", .2280, 1.12, 1.04, .058, .82, .38, 1.22, .996, "C"),
    ("GB", "United Kingdom", "Europe", .3052, 1.12, 1.03, .073, .78, .49, 1.08, .994, "B"),
    ("US", "United States", "North America", .0752, 1.15, 1.00, .080, .70, .36, 1.12, .994, "B"),
    ("CA", "Canada", "North America", .0740, 1.10, 1.03, .075, .68, .32, 1.02, .995, "C"),
    ("MX", "Mexico", "North America", .1130, 1.20, 1.10, .120, .56, .29, .62, .985, "C"),
    # China uses a Huawei sovereign-system equivalent, normalized to the same
    # output as the EUR 4m reference service unit. Public price and full
    # workload benchmarks are not yet available, hence Grade D uncertainty.
    ("CN", "China", "Asia", .0800, 1.16, 1.00, .075, .58, .28, .68, .995, "D"),
    ("IN", "India", "Asia", .0880, 1.24, 1.17, .115, .52, .25, .48, .982, "C"),
    ("JP", "Japan", "Asia", .1390, 1.17, 1.06, .052, .78, .42, .95, .997, "B"),
    ("KR", "South Korea", "Asia", .0960, 1.17, 1.06, .070, .69, .36, .86, .995, "B"),
    ("TW", "Taiwan", "Asia", .1170, 1.18, 1.05, .065, .70, .39, .82, .994, "B"),
    ("SG", "Singapore", "Asia", .1690, 1.25, 1.05, .060, .84, .55, 1.05, .997, "B"),
    ("AU", "Australia", "Other", .1210, 1.17, 1.08, .080, .72, .34, 1.02, .994, "C"),
    ("BR", "Brazil", "Other", .1270, 1.22, 1.25, .140, .60, .32, .58, .982, "C"),
    ("CL", "Chile", "Other", .1050, 1.17, 1.12, .105, .59, .30, .66, .988, "C"),
    ("SA", "Saudi Arabia", "Other", .0560, 1.24, 1.08, .078, .62, .30, .72, .992, "B"),
    ("AE", "United Arab Emirates", "Other", .1050, 1.25, 1.07, .073, .72, .37, .85, .995, "B"),
]


def crf(rate: float, years: int) -> float:
    return rate * (1 + rate) ** years / ((1 + rate) ** years - 1)


def build() -> list[dict]:
    rows = []
    for (code, country, group, power, pue, hw_mult, wacc, facility,
         connection, labour, availability, grade) in C:
        base = {
            "code": code, "country": country, "group": group,
            "electricity_eur_kwh": power, "pue": pue,
            "landed_hardware_multiplier": hw_mult, "hurdle_rate": wacc,
            "facility_capex_eur_m": facility,
            "connection_capex_eur_m": connection,
            "operations_labour_multiplier": labour,
            "availability": availability, "evidence_grade": grade,
            "scales": {}, "scale_inputs": {},
        }

        # On-premises case: one service unit in an existing company facility.
        base["scale_inputs"]["onprem_250kw"] = {
            "service_units": 1.0, "connection_mw": .25,
            "pue": pue, "electricity_eur_kwh": power,
            "electricity_usd_kwh": power * EUR_USD,
            "utilization": .70, "availability": availability,
            "facility_capex_eur_m": facility,
            "connection_capex_eur_m": connection,
        }

        # Data-center case: 100 MW at the utility meter. Commissioned IT capacity is
        # connection/PUE and is filled with 142 kW design-power service units.
        # The EUR 9m/MW base is a 2026 greenfield shell/electrical/mechanical proxy;
        # country facility inputs provide a damped local construction multiplier.
        data_center_pue = max(1.07, pue - .03)
        data_center_it_mw = 100 / data_center_pue
        data_center_units = data_center_it_mw * 1000 / 142
        data_center_wacc = max(.05, wacc - .005)
        data_center_availability = 1 - (1 - availability) * .40
        data_center_construction_m_per_it_mw = 9.0 * (facility / .65) ** .60
        data_center_facility_m = data_center_it_mw * data_center_construction_m_per_it_mw
        data_center_connection_m = connection * (100 / .25) ** .90
        base["scale_inputs"]["datacenter_100mw"] = {
            "service_units": round(data_center_units, 1), "connection_mw": 100,
            "commissioned_it_mw": round(data_center_it_mw, 1),
            "pue": data_center_pue, "electricity_eur_kwh": power * .92,
            "electricity_usd_kwh": power * .92 * EUR_USD,
            "utilization": .80, "availability": data_center_availability,
            "facility_capex_eur_m": data_center_facility_m,
            "connection_capex_eur_m": data_center_connection_m,
            "construction_eur_m_per_it_mw": data_center_construction_m_per_it_mw,
        }

        for scale in SCALES:
            base["scales"][scale] = {}
            s = base["scale_inputs"][scale]
            for key, w in WORKLOADS.items():
                local_unc = {"B": .12, "C": .18, "D": .25}[grade]
                if scale == "onprem_250kw":
                    output_tokens_m = (
                        w["output_tps"] * SECONDS_YEAR * s["utilization"]
                        * s["availability"] / 1_000_000
                    )
                    hardware_m = 4.0 * hw_mult * crf(wacc, 3)
                    site_m = (facility + connection) * crf(wacc, 12)
                    electricity_m = w["it_kw"] * pue * 8760 * power / 1_000_000
                    operations_m = .22 * hw_mult + .20 * labour + w["extra_opex_m"]
                else:
                    output_tokens_m = (
                        s["service_units"] * w["output_tps"] * SECONDS_YEAR
                        * s["utilization"] * s["availability"] / 1_000_000
                    )
                    hardware_m = (
                        s["service_units"] * 4.0 * hw_mult * .92
                        * crf(data_center_wacc, 3)
                    )
                    site_m = (
                        s["facility_capex_eur_m"] * crf(data_center_wacc, 15)
                        + s["connection_capex_eur_m"] * crf(data_center_wacc, 20)
                    )
                    electricity_m = (
                        s["service_units"] * w["it_kw"] * s["pue"] * 8760
                        * s["electricity_eur_kwh"] / 1_000_000
                    )
                    operations_m = (
                        s["service_units"] * (
                            .22 * hw_mult * .90 + w["extra_opex_m"] * .55
                        )
                        + 25 * labour + 8
                    )
                    # Data-center construction and tariff proxies add uncertainty even
                    # where the underlying country electricity observation is B.
                    local_unc = max(local_unc, .18)

                components_eur = {
                    "hardware": hardware_m * 1_000_000 / output_tokens_m,
                    "electricity": electricity_m * 1_000_000 / output_tokens_m,
                    "site_grid": site_m * 1_000_000 / output_tokens_m,
                    "operations": operations_m * 1_000_000 / output_tokens_m,
                }
                components = {k: v * EUR_USD for k, v in components_eur.items()}
                base["scales"][scale][key] = {
                    **{k: round(v, 2) for k, v in components.items()},
                    "total": round(sum(components.values()), 2),
                    "local_uncertainty_pct": round(local_unc * 100),
                }
        rows.append(base)
    return rows


def _apply_kimi_k3_calibration() -> dict:
    """Scale the provisional throughput priors to Kimi K3 API economics.

    The China 100 MW general-inference case is the calibration anchor because
    it applies the Kimi-class service economics in the provider's home market.
    Since all modeled annual costs are divided by throughput, the required
    throughput multiplier is baseline cost divided by the price-implied maximum
    production cost. Other workloads retain their original relative ratios.
    """
    blended_input_price = (
        KIMI_K3_API["cache_hit_rate"]
        * KIMI_K3_API["cache_hit_input_usd_per_m"]
        + (1 - KIMI_K3_API["cache_hit_rate"])
        * KIMI_K3_API["cache_miss_input_usd_per_m"]
    )
    api_revenue_per_m_output = (
        KIMI_K3_API["output_usd_per_m"]
        + KIMI_K3_API["general_input_output_ratio"] * blended_input_price
    )
    maximum_production_cost = (
        api_revenue_per_m_output
        / (1 - KIMI_K3_API["maximum_subsidy_share"])
    )
    baseline_rows = build()
    china = next(row for row in baseline_rows if row["code"] == "CN")
    baseline_cost = china["scales"]["datacenter_100mw"]["general"]["total"]
    multiplier = baseline_cost / maximum_production_cost
    baseline_throughputs = {
        key: assumptions["output_tps"] for key, assumptions in WORKLOADS.items()
    }
    for key, assumptions in WORKLOADS.items():
        assumptions["output_tps"] = round(
            baseline_throughputs[key] * multiplier
        )
    calibrated_china = next(row for row in build() if row["code"] == "CN")
    return {
        **KIMI_K3_API,
        "blended_input_usd_per_m": round(blended_input_price, 2),
        "api_revenue_usd_per_m_output": round(api_revenue_per_m_output, 2),
        "maximum_production_cost_usd_per_m_output": round(
            maximum_production_cost, 2
        ),
        "uncalibrated_china_cost_usd_per_m_output": baseline_cost,
        "throughput_multiplier": round(multiplier, 4),
        "baseline_output_tps": baseline_throughputs,
        "calibrated_output_tps": {
            key: assumptions["output_tps"]
            for key, assumptions in WORKLOADS.items()
        },
        "calibrated_china_cost_usd_per_m_output": calibrated_china[
            "scales"
        ]["datacenter_100mw"]["general"]["total"],
        "interpretation": (
            "General inference is price-implied; long-context and agentic "
            "throughputs apply the same K3 efficiency multiplier."
        ),
    }


KIMI_K3_CALIBRATION = _apply_kimi_k3_calibration()


def main() -> None:
    rows = build()
    (OUT / "sfii-country-estimates.json").write_text(
        json.dumps({
            "version": DATA_VERSION, "date_modified": RELEASE_DATE,
            "output_currency": "USD", "eur_usd": EUR_USD,
            "workloads": WORKLOADS, "scales": SCALES,
            "calibration": KIMI_K3_CALIBRATION, "countries": rows,
        }, indent=2) + "\n",
        encoding="utf-8",
    )
    with (OUT / "sfii-country-estimates.csv").open("w", newline="", encoding="utf-8") as f:
        fields = [
            "code", "country", "group", "output_currency", "eur_usd",
            "electricity_eur_kwh", "pue",
            "landed_hardware_multiplier", "hurdle_rate",
            "facility_capex_eur_m", "connection_capex_eur_m",
            "operations_labour_multiplier", "availability", "evidence_grade",
        ]
        for scale in SCALES:
            fields += [
                f"{scale}_{x}" for x in (
                    "service_units", "connection_mw", "commissioned_it_mw", "pue",
                    "electricity_eur_kwh", "electricity_usd_kwh", "utilization", "availability",
                    "facility_capex_eur_m", "connection_capex_eur_m",
                    "construction_eur_m_per_it_mw",
                )
            ]
            for key in WORKLOADS:
                fields += [f"{scale}_{key}_{x}" for x in ("hardware", "electricity", "site_grid", "operations", "total", "local_uncertainty_pct")]
        writer = csv.DictWriter(f, fieldnames=fields, lineterminator="\n")
        writer.writeheader()
        for row in rows:
            flat = {k: row[k] for k in fields if k in row}
            flat["output_currency"] = "USD"
            flat["eur_usd"] = EUR_USD
            for scale in SCALES:
                for metric, value in row["scale_inputs"][scale].items():
                    field = f"{scale}_{metric}"
                    if field in fields:
                        flat[field] = value
                for key in WORKLOADS:
                    for metric, value in row["scales"][scale][key].items():
                        flat[f"{scale}_{key}_{metric}"] = value
            writer.writerow(flat)

    # Pandoc includes this fragment at the end of the standalone HTML. Keeping
    # the dataset here makes the compiled paper self-contained and viewable
    # without a web server or third-party chart library.
    chart_rows = []
    for row in rows:
        chart_rows.append({
            "country": row["country"], "group": row["group"],
            "power": row["electricity_eur_kwh"], "pue": row["pue"],
            "grade": row["evidence_grade"],
            "basis": "Huawei sovereign equivalent" if row["code"] == "CN" else "GB300-class reference",
            "scales": row["scales"], "scale_inputs": row["scale_inputs"],
        })
    payload = json.dumps(chart_rows, separators=(",", ":"), ensure_ascii=False)
    chart_script = r'''<script>
(() => {
  "use strict";
  const rows = __PAYLOAD__;
  const workloads = {
    general: {title: "General inference", subtitle: "8k input · 2k visible output · ≤5 s TTFT", tps: 3683},
    long_context: {title: "Long-context analysis", subtitle: "128k input · 4k visible output · ≤20 s TTFT", tps: 1719},
    agentic: {title: "Agentic engineering", subtitle: "Multi-turn tool and code tasks · production throughput", tps: 777}
  };
  const components = [
    ["hardware", "Annualized hardware"], ["electricity", "Electricity"],
    ["site_grid", "Site & grid"], ["operations", "Operations"]
  ];
  const money = new Intl.NumberFormat("en", {style:"currency", currency:"USD", maximumFractionDigits:0});
  const one = new Intl.NumberFormat("en", {minimumFractionDigits:2, maximumFractionDigits:2});
  const scaleLabels = {
    datacenter_100mw: "100 MW data center · 80% utilization",
    onprem_250kw: "On-prem solution · 250 kW · 70% utilization"
  };
  const query = new URLSearchParams(window.location.search);
  const legacyScales = {campus_100mw: "datacenter_100mw", node_250kw: "onprem_250kw"};
  const requestedScale = legacyScales[query.get("scale")] || query.get("scale");
  let activeScale = Object.hasOwn(scaleLabels, requestedScale) ? requestedScale : "datacenter_100mw";
  const requestedWorkload = query.get("workload") === "interactive" ? "general" : query.get("workload");
  let activeWorkload = Object.hasOwn(workloads, requestedWorkload) ? requestedWorkload : "general";

  function render(key) {
    activeWorkload = key;
    const target = document.getElementById("sfii-chart");
    if (!target) return;
    const sorted = [...rows].sort((a,b) => a.scales[activeScale][key].total - b.scales[activeScale][key].total);
    const max = Math.ceil(Math.max(...sorted.map(d => d.scales[activeScale][key].total)) / 50) * 50;
    document.getElementById("sfii-chart-title").textContent = workloads[key].title;
    document.getElementById("sfii-chart-subtitle").textContent = workloads[key].subtitle;
    document.getElementById("sfii-scale-note").textContent = scaleLabels[activeScale];
    target.innerHTML = "";
    const axis = document.createElement("div");
    axis.className = "chart-axis";
    axis.innerHTML = `<span>0</span><span>${money.format(max*.25)}</span><span>${money.format(max*.5)}</span><span>${money.format(max*.75)}</span><span>${money.format(max)}</span>`;
    target.append(axis);
    for (const d of sorted) {
      const v = d.scales[activeScale][key];
      const inputs = d.scale_inputs[activeScale];
      const row = document.createElement("div");
      row.className = "chart-row";
      const tip = `${d.country}\nTotal: ${money.format(v.total)} per million output tokens\n` +
        components.map(([id,label]) => `${label}: ${money.format(v[id])}`).join("\n") +
        `\nScale: ${scaleLabels[activeScale]}\nHardware basis: ${d.basis}\nElectricity input: $${one.format(inputs.electricity_usd_kwh)}/kWh\nPUE proxy: ${inputs.pue.toFixed(2)}\nConnection: ${inputs.connection_mw} MW\nService units: ${one.format(inputs.service_units)}\nCountry-input grade: ${d.grade}\nCountry-specific uncertainty: ±${v.local_uncertainty_pct}%`;
      row.title = tip;
      const name = document.createElement("span");
      name.className = "chart-country";
      name.textContent = d.country;
      const track = document.createElement("span");
      track.className = "chart-track";
      track.setAttribute("aria-label", tip.replaceAll("\n", "; "));
      for (const [id] of components) {
        const seg = document.createElement("span");
        seg.className = `chart-segment ${id}`;
        seg.style.width = `${v[id] / max * 100}%`;
        track.append(seg);
      }
      const total = document.createElement("span");
      total.className = "chart-total";
      total.textContent = money.format(v.total);
      row.append(name, track, total);
      target.append(row);
    }
    document.querySelectorAll(".chart-tab").forEach(btn => {
      const active = btn.dataset.workload === key;
      btn.classList.toggle("active", active);
      btn.setAttribute("aria-selected", String(active));
      btn.tabIndex = active ? 0 : -1;
    });
    const url = new URL(window.location.href);
    url.searchParams.set("scale", activeScale);
    url.searchParams.set("workload", activeWorkload);
    history.replaceState(null, "", url);
  }

  document.addEventListener("DOMContentLoaded", () => {
    const selectedScale = document.querySelector(`input[name="deployment-scale"][value="${activeScale}"]`);
    if (selectedScale) selectedScale.checked = true;
    document.querySelectorAll(".chart-tab").forEach(btn => btn.addEventListener("click", () => render(btn.dataset.workload)));
    document.querySelectorAll('input[name="deployment-scale"]').forEach(input => input.addEventListener("change", () => {
      activeScale = input.value;
      render(activeWorkload);
    }));
    render(activeWorkload);
  });
})();
</script>
'''.replace("__PAYLOAD__", payload)
    (OUT / "sfii-chart-embed.html").write_text(chart_script, encoding="utf-8")


if __name__ == "__main__":
    main()