#!/usr/bin/env python3
"""Calculate the 2010-2026 fixed-service AI Cost Index backcast."""

from __future__ import annotations

import csv
import json
import statistics
from pathlib import Path

from build_sfii_estimates import (
    C, EUR_USD, KIMI_K3_CALIBRATION, SECONDS_YEAR, WORKLOADS, crf,
)


OUT = Path(__file__).resolve().parent
SCALES = ("datacenter_100mw", "onprem_250kw")


def load_inputs() -> dict[tuple[int, str], dict]:
    inputs = {}
    with (OUT / "historical-country-inputs.csv").open(encoding="utf-8") as source:
        for row in csv.DictReader(source):
            for field in (
                "electricity_usd_2026_kwh", "wacc",
                "construction_real_factor", "operations_real_factor",
            ):
                row[field] = float(row[field])
            row["year"] = int(row["year"])
            inputs[(row["year"], row["code"])] = row
    return inputs


def calculate() -> list[dict]:
    history = load_inputs()
    results = []

    for (code, country, group, _power, pue, hw_mult, _current_wacc,
         facility, connection, labour, availability, grade) in C:
        for year in range(2010, 2027):
            annual = history[(year, code)]
            power_usd = annual["electricity_usd_2026_kwh"]
            wacc = annual["wacc"]
            construction_factor = annual["construction_real_factor"]
            operations_factor = annual["operations_real_factor"]

            data_center_pue = max(1.07, pue - .03)
            data_center_units = (100 / data_center_pue) * 1000 / 142
            data_center_wacc = max(.05 if year == 2026 else .035, wacc - .005)
            data_center_availability = 1 - (1 - availability) * .40
            construction_eur_m_per_it_mw = 9.0 * (facility / .65) ** .60
            facility_eur_m = (100 / data_center_pue) * construction_eur_m_per_it_mw
            connection_eur_m = connection * (100 / .25) ** .90

            for workload, assumptions in WORKLOADS.items():
                for scale in SCALES:
                    if scale == "onprem_250kw":
                        output_tokens_m = (
                            assumptions["output_tps"] * SECONDS_YEAR * .70
                            * availability / 1_000_000
                        )
                        hardware_m = (
                            4.0 * EUR_USD * hw_mult * crf(wacc, 3)
                        )
                        site_grid_m = (
                            (facility + connection) * EUR_USD
                            * construction_factor * crf(wacc, 12)
                        )
                        electricity_m = (
                            assumptions["it_kw"] * pue * 8760
                            * power_usd / 1_000_000
                        )
                        operations_m = (
                            .22 * hw_mult + .20 * labour
                            + assumptions["extra_opex_m"]
                        ) * EUR_USD * operations_factor
                    else:
                        output_tokens_m = (
                            data_center_units * assumptions["output_tps"]
                            * SECONDS_YEAR * .80 * data_center_availability
                            / 1_000_000
                        )
                        hardware_m = (
                            data_center_units * 4.0 * EUR_USD * hw_mult
                            * .92 * crf(data_center_wacc, 3)
                        )
                        site_grid_m = (
                            facility_eur_m * EUR_USD * construction_factor
                            * crf(data_center_wacc, 15)
                            + connection_eur_m * EUR_USD * construction_factor
                            * crf(data_center_wacc, 20)
                        )
                        electricity_m = (
                            data_center_units * assumptions["it_kw"]
                            * data_center_pue * 8760 * power_usd * .92
                            / 1_000_000
                        )
                        operations_m = (
                            data_center_units * (
                                .22 * hw_mult * .90
                                + assumptions["extra_opex_m"] * .55
                            ) + 25 * labour + 8
                        ) * EUR_USD * operations_factor

                    components = {
                        "hardware": hardware_m * 1_000_000 / output_tokens_m,
                        "electricity": electricity_m * 1_000_000 / output_tokens_m,
                        "site_grid": site_grid_m * 1_000_000 / output_tokens_m,
                        "operations": operations_m * 1_000_000 / output_tokens_m,
                    }
                    results.append({
                        "year": year, "code": code, "country": country,
                        "group": group, "scale": scale,
                        "workload": workload,
                        **{key: round(value, 2) for key, value in components.items()},
                        "total": round(sum(components.values()), 2),
                        "electricity_usd_2026_kwh": round(power_usd, 4),
                        "wacc": round(wacc, 4),
                        "construction_real_factor": construction_factor,
                        "power_source": annual["power_source"],
                        "evidence_grade": grade,
                    })
    return results


def build_summary(results: list[dict]) -> dict:
    summaries = {}
    for scale in SCALES:
        summaries[scale] = {}
        for workload in WORKLOADS:
            by_year = {}
            for year in range(2010, 2027):
                rows = [
                    row for row in results
                    if row["scale"] == scale and row["workload"] == workload
                    and row["year"] == year
                ]
                ordered = sorted(rows, key=lambda row: row["total"])
                medians = {
                    component: statistics.median(row[component] for row in rows)
                    for component in (
                        "hardware", "electricity", "site_grid", "operations", "total"
                    )
                }
                by_year[year] = {
                    "median": round(medians["total"], 2),
                    "mean": round(statistics.mean(row["total"] for row in rows), 2),
                    "minimum": round(ordered[0]["total"], 2),
                    "minimum_country": ordered[0]["country"],
                    "maximum": round(ordered[-1]["total"], 2),
                    "maximum_country": ordered[-1]["country"],
                    "median_components": {
                        key: round(value, 2) for key, value in medians.items()
                        if key != "total"
                    },
                }
            base = by_year[2026]["median"]
            for year in by_year:
                by_year[year]["index_2026_100"] = round(
                    by_year[year]["median"] / base * 100, 1
                )
            summaries[scale][workload] = by_year
    return summaries


def main() -> None:
    results = calculate()
    summary = build_summary(results)
    payload = {
        "method": "Fixed 2026 service and hardware basket repriced under annual conditions",
        "currency": "constant 2026 USD",
        "throughput_calibration": KIMI_K3_CALIBRATION,
        "years": list(range(2010, 2027)),
        "summary": summary,
        "results": results,
    }
    (OUT / "historical-ai-cost-index.json").write_text(
        json.dumps(payload, indent=2) + "\n", encoding="utf-8"
    )
    fields = list(results[0])
    with (OUT / "historical-ai-cost-index.csv").open("w", newline="", encoding="utf-8") as output:
        writer = csv.DictWriter(output, fieldnames=fields, lineterminator="\n")
        writer.writeheader()
        writer.writerows(results)
    print(f"Wrote {len(results)} historical cost rows")


if __name__ == "__main__":
    main()