#!/usr/bin/env python3
"""Deterministic single-server FIFO timing model.

FIFO admission occurs at request start. At most one request is in service;
all pre-start waiting is modeled on the client side. There is no server queue,
network, batching, token generation, timeout, drop, random noise, or GPU work.
"""

from __future__ import annotations

import argparse
import csv
from dataclasses import asdict, dataclass
from fractions import Fraction
import json
from math import ceil
from pathlib import Path
from statistics import mean
from typing import Iterable


@dataclass(frozen=True)
class Request:
    request_id: int
    intended_schedule_ms: int
    credit_issued_ms: int
    request_start_ms: int
    response_end_ms: int

    @property
    def request_latency_ms(self) -> int:
        return self.response_end_ms - self.request_start_ms

    @property
    def effective_latency_ms(self) -> int:
        return self.response_end_ms - self.credit_issued_ms

    @property
    def scheduled_latency_ms(self) -> int:
        return self.response_end_ms - self.intended_schedule_ms

    @property
    def credit_to_start_wait_ms(self) -> int:
        return self.request_start_ms - self.credit_issued_ms

    @property
    def issuance_delay_ms(self) -> int:
        return self.credit_issued_ms - self.intended_schedule_ms

    def record(self) -> dict:
        return asdict(self) | {
            "request_latency_ms": self.request_latency_ms,
            "effective_latency_ms": self.effective_latency_ms,
            "scheduled_latency_ms": self.scheduled_latency_ms,
            "credit_to_start_wait_ms": self.credit_to_start_wait_ms,
            "issuance_delay_ms": self.issuance_delay_ms,
        }


def simulate(
    n: int,
    interval_ms: int,
    service_ms: int = 400,
    issuance_delays_ms: Iterable[int] | None = None,
) -> list[Request]:
    """Run a drained finite batch, with request IDs in FIFO credit order.

    s_i = i * interval; c_i = s_i + delay_i
    start_i = max(c_i, end_(i-1)); end_i = start_i + service

    The empty system is available at t=0. Credits must be nondecreasing.
    Equal-time credits use increasing request ID order.
    """
    if n < 1 or interval_ms < 1 or service_ms < 1:
        raise ValueError("n, interval_ms, and service_ms must be positive")
    delays = [0] * n if issuance_delays_ms is None else list(issuance_delays_ms)
    if len(delays) != n or any(d < 0 for d in delays):
        raise ValueError("Provide one nonnegative issuance delay per request")
    rows: list[Request] = []
    previous_end = 0
    previous_credit = -1
    for i, delay in enumerate(delays):
        scheduled = i * interval_ms
        credit = scheduled + delay
        if credit < previous_credit:
            raise ValueError("This model requires credits in nondecreasing order")
        start = max(credit, previous_end)
        end = start + service_ms
        rows.append(Request(i, scheduled, credit, start, end))
        previous_credit, previous_end = credit, end
    return rows


def nearest_rank(values: Iterable[int], percentile: int | float) -> int:
    """Return sorted[ceil(N*p/100)-1], with exact decimal rank arithmetic.

    Domain is 0 < p <= 100. No interpolation, extrapolation, or p=0 convention.
    """
    ordered = sorted(values)
    p = Fraction(str(percentile))
    if not ordered or not 0 < p <= 100:
        raise ValueError("Nonempty values and 0 < percentile <= 100 required")
    return ordered[ceil(len(ordered) * p / 100) - 1]


def timeline(rows: list[Request]) -> list[dict]:
    """Post-event state, after all events at each distinct timestamp.

    At ties, completion is applied before credit issuance before dispatch.
    Counts below are equivalent to that ordering. Waiting excludes service.
    """
    times = sorted({0} | {
        t for row in rows
        for t in (row.credit_issued_ms, row.request_start_ms, row.response_end_ms)
    })
    result = []
    for t in times:
        issued = sum(r.credit_issued_ms <= t for r in rows)
        started = sum(r.request_start_ms <= t for r in rows)
        completed = sum(r.response_end_ms <= t for r in rows)
        result.append({
            "time_ms": t,
            "issued_cumulative": issued,
            "started_cumulative": started,
            "completed_cumulative": completed,
            "waiting": issued - started,
            "in_service": started - completed,
            "in_system": issued - completed,
        })
    return result


def state_area_ms(states: list[dict], key: str) -> int:
    return sum(
        current[key] * (following["time_ms"] - current["time_ms"])
        for current, following in zip(states, states[1:])
    )


def summarize(rows: list[Request], interval_ms: int, slo_ms: int = 1000) -> dict:
    states = timeline(rows)
    out = {
        "n": len(rows),
        "target_interval_ms": interval_ms,
        "target_arrival_rate_per_s": 1000 / interval_ms,
        "service_ms": rows[0].request_latency_ms,
        "service_capacity_per_s": 1000 / rows[0].request_latency_ms,
        "target_load_ratio": rows[0].request_latency_ms / interval_ms,
        "slo_threshold_ms_inclusive": slo_ms,
        "first_credit_ms": rows[0].credit_issued_ms,
        "last_credit_ms": rows[-1].credit_issued_ms,
        "first_start_ms": rows[0].request_start_ms,
        "last_start_ms": rows[-1].request_start_ms,
        "last_end_ms": rows[-1].response_end_ms,
        "mean_dispatch_gap_ms": mean([
            b.request_start_ms - a.request_start_ms
            for a, b in zip(rows, rows[1:])
        ]) if len(rows) > 1 else None,
        "mean_credit_to_start_wait_ms": mean(r.credit_to_start_wait_ms for r in rows),
        "max_waiting_requests": max(s["waiting"] for s in states),
        "queue_area_request_ms": state_area_ms(states, "waiting"),
        "in_service_area_request_ms": state_area_ms(states, "in_service"),
        "in_system_area_request_ms": state_area_ms(states, "in_system"),
        "completed": states[-1]["completed_cumulative"],
        "failed": 0,
        "dropped": 0,
    }
    for metric in ("request", "effective", "scheduled"):
        values = [getattr(r, f"{metric}_latency_ms") for r in rows]
        passing = sum(v <= slo_ms for v in values)
        out.update({
            f"{metric}_mean_ms": mean(values),
            f"{metric}_p95_ms": nearest_rank(values, 95),
            f"{metric}_p99_ms": nearest_rank(values, 99),
            f"{metric}_max_ms": max(values),
            f"{metric}_slo_pass_count": passing,
            f"{metric}_slo_pass_fraction": passing / len(rows),
        })
    return out


def write_csv(path: Path, rows: list[dict]) -> None:
    with path.open("w", newline="", encoding="utf-8") as stream:
        writer = csv.DictWriter(stream, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)


def generate(output_dir: Path) -> dict:
    output_dir.mkdir(parents=True, exist_ok=True)
    request_rows, state_rows, summaries = [], [], []
    for interval in (100, 200, 400, 800):
        requests = simulate(n=20, interval_ms=interval)
        prefix = {"scenario": f"interval_{interval}ms", "target_interval_ms": interval}
        request_rows.extend(prefix | row.record() for row in requests)
        state_rows.extend(prefix | state for state in timeline(requests))
        summaries.append(summarize(requests, interval))
    four = simulate(4, 100)
    delayed = simulate(1, 100, issuance_delays_ms=[300])
    payload = {
        "experiment": "single_server_fifo_timing_boundaries",
        "schema_version": 1,
        "is_real_inference_benchmark": False,
        "units": "milliseconds unless named otherwise",
        "assumptions": [
            "Deterministic FIFO single service slot; empty system at time zero.",
            "Constant 400 ms service, with request start equal to service admission.",
            "Intended schedule and issuance coincide in the main sweep.",
            "Credit-to-start waiting is modeled at the client-side dispatch boundary.",
            "Infinite waiting capacity; all requests succeed; full post-arrival drain.",
            "No network, server queue, batching, tokens, timeouts, drops or GPU work.",
            "Finite N=20 runs are illustrative, not steady-state estimates.",
            "Percentiles use nearest rank ceil(N*p/100), without interpolation.",
            "The one-second SLO is inclusive: latency <= 1000 ms.",
            "This is not an AIPerf scheduler implementation or performance result.",
        ],
        "sweep": summaries,
        "four_request_example": [r.record() for r in four],
        "four_request_summary": summarize(four, 100),
        "delayed_issuance_example": delayed[0].record(),
    }
    write_csv(output_dir / "requests.csv", request_rows)
    write_csv(output_dir / "timeline.csv", state_rows)
    write_csv(output_dir / "sweep.csv", summaries)
    write_csv(output_dir / "four_requests.csv", [r.record() for r in four])
    write_csv(output_dir / "delayed_issuance.csv", [delayed[0].record()])
    (output_dir / "results.json").write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
    return payload


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output-dir", type=Path, default=Path(__file__).parent / "data")
    args = parser.parse_args()
    payload = generate(args.output_dir)
    print("Model configuration: deterministic FIFO service with client-side credit-to-start waiting.")
    print("interval_ms | request_mean_ms | effective_mean_ms | effective_p95_ms | <=1000ms request/effective")
    for s in payload["sweep"]:
        print(f"{s['target_interval_ms']:11} | {s['request_mean_ms']:15} | "
              f"{s['effective_mean_ms']:17} | {s['effective_p95_ms']:16} | "
              f"{s['request_slo_pass_count']}/20 vs {s['effective_slo_pass_count']}/20")
    d = payload["delayed_issuance_example"]
    print("Delayed issuance (schedule, credit, start, end) = (0, 300, 300, 700) ms")
    print(f"Latencies: request={d['request_latency_ms']}, effective={d['effective_latency_ms']}, "
          f"scheduled={d['scheduled_latency_ms']} ms")
    print(f"Wrote six data files to {args.output_dir}")


if __name__ == "__main__":
    main()