src/main.py - master - d.mit/zadanie

src/main.py
1 346 строк42 KB
d.mit
create: .cm-token, docker-compose.yml, Dockerfile, instruction.md, README.md, requirements.txt, main.py, README.md
20 апр 2026, 23:46
20 апр 2026, 23:46c103833
from __future__ import annotations

import bisect
import heapq
import math
from collections import defaultdict
from dataclasses import dataclass
from datetime import datetime
from typing import Any, Dict, List, Optional, Tuple

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse

EPS = 1e-9
VIS_EPS = 1e-6
TAU = 2.0 * math.pi

app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None)


class InputError(Exception):
    pass


def bad_request() -> JSONResponse:
    return JSONResponse(status_code=400, content={"status": "incorrect_input"})


def ensure_keys(obj: Any, keys: set[str]) -> None:
    if not isinstance(obj, dict) or set(obj.keys()) != keys:
        raise InputError()


def is_number(x: Any) -> bool:
    return isinstance(x, (int, float)) and not isinstance(x, bool) and math.isfinite(float(x))


def as_float(x: Any, *, low: Optional[float] = None, high: Optional[float] = None) -> float:
    if not is_number(x):
        raise InputError()
    v = float(x)
    if low is not None and v < low:
        raise InputError()
    if high is not None and v > high:
        raise InputError()
    return v


def as_int(x: Any, *, low: Optional[int] = None, high: Optional[int] = None) -> int:
    if not isinstance(x, int) or isinstance(x, bool):
        raise InputError()
    v = x
    if low is not None and v < low:
        raise InputError()
    if high is not None and v > high:
        raise InputError()
    return v


def as_str(x: Any) -> str:
    if not isinstance(x, str) or x == "":
        raise InputError()
    return x


def as_bool(x: Any) -> bool:
    if not isinstance(x, bool):
        raise InputError()
    return x


# =========================================================
# TASK 1
# =========================================================

@dataclass(frozen=True)
class GravityAssist:
    velocity_gain: float
    fuel_consumption: int
    time_to_execute: int


@dataclass
class CruiseState:
    pos: float
    time: float
    speed: float
    fuel: int


class FlightEngine:
    def __init__(self, mass_shuttle: float, mass_fuel_unit: float, power_per_unit: float):
        self.mass_shuttle = mass_shuttle
        self.mass_fuel_unit = mass_fuel_unit
        self.power_per_unit = power_per_unit
        self._accel_cache: Dict[int, Tuple[List[float], List[float]]] = {}
        self._brake_cache: Dict[int, List[float]] = {}

    def _ensure_accel(self, fuel: int) -> None:
        if fuel in self._accel_cache:
            return
        max_burns = fuel // 2
        dist = [0.0]
        vel = [0.0]
        v = 0.0
        s = 0.0
        for i in range(1, max_burns + 1):
            # после мгновенного расхода топлива масса уменьшена на 1 единицу топлива
            mass_now = self.mass_shuttle + (fuel - i) * self.mass_fuel_unit
            a = self.power_per_unit / mass_now
            s += v + 0.5 * a
            v += a
            dist.append(s)
            vel.append(v)
        self._accel_cache[fuel] = (dist, vel)

    def accel_profile(self, fuel: int, total_distance: float) -> Tuple[int, float, float, int]:
        self._ensure_accel(fuel)
        dist, vel = self._accel_cache[fuel]
        limit = total_distance / 2.0
        burns = bisect.bisect_right(dist, limit + EPS) - 1
        if burns < 0:
            burns = 0
        return burns, dist[burns], vel[burns], fuel - burns

    def _ensure_brake(self, fuel_available: int) -> None:
        if fuel_available in self._brake_cache:
            return
        arr = []
        for i in range(1, fuel_available + 1):
            mass_now = self.mass_shuttle + (fuel_available - i) * self.mass_fuel_unit
            a = self.power_per_unit / mass_now
            arr.append(a)
        self._brake_cache[fuel_available] = arr

    def brake_profile(self, speed: float, fuel_available: int) -> Optional[Tuple[float, int]]:
        """
        Возвращает:
        (distance_to_stop, braking_time_seconds)
        или None, если остановиться нельзя.
        """
        if speed <= EPS:
            return 0.0, 0
        if fuel_available <= 0:
            return None

        self._ensure_brake(fuel_available)
        decels = self._brake_cache[fuel_available]

        v = speed
        dist = 0.0
        for i, a in enumerate(decels, start=1):
            if v <= a + EPS:
                dist += v * v / (2.0 * a)
                return dist, i
            dist += v - 0.5 * a
            v -= a
        return None


def parse_task1(payload: Any) -> Dict[str, Any]:
    ensure_keys(payload, {
        "mass_shuttle",
        "mass_fuel_unit",
        "power_per_unit",
        "oxygen_time",
        "total_fuel",
        "fuel_consumption",
        "bodies",
        "edges",
    })

    mass_shuttle = as_float(payload["mass_shuttle"], low=0.0)
    if mass_shuttle <= 0:
        raise InputError()
    mass_fuel_unit = as_float(payload["mass_fuel_unit"], low=0.0)
    if mass_fuel_unit <= 0:
        raise InputError()
    power_per_unit = as_float(payload["power_per_unit"], low=0.0)
    if power_per_unit <= 0:
        raise InputError()

    oxygen_time = as_int(payload["oxygen_time"], low=1)
    total_fuel = as_int(payload["total_fuel"], low=0)
    fuel_consumption = as_float(payload["fuel_consumption"], low=0.0)
    if fuel_consumption < 0:
        raise InputError()

    bodies = payload["bodies"]
    if not isinstance(bodies, list) or len(bodies) > 100:
        raise InputError()

    body_assists: Dict[str, List[GravityAssist]] = {}
    body_ids = set()
    for body in bodies:
        ensure_keys(body, {"id", "gravity_assists"})
        body_id = as_str(body["id"])
        if body_id in body_ids or body_id in {"start_point", "rescue_point"}:
            raise InputError()
        body_ids.add(body_id)

        ga_list = body["gravity_assists"]
        if not isinstance(ga_list, list):
            raise InputError()

        assists: List[GravityAssist] = []
        for ga in ga_list:
            ensure_keys(ga, {"velocity_gain", "fuel_consumption", "time_to_execute"})
            velocity_gain = as_float(ga["velocity_gain"], low=0.0)
            fuel_cost = as_int(ga["fuel_consumption"], low=0)
            time_to_execute = as_int(ga["time_to_execute"], low=0)
            assists.append(GravityAssist(
                velocity_gain=velocity_gain,
                fuel_consumption=fuel_cost,
                time_to_execute=time_to_execute,
            ))
        body_assists[body_id] = assists

    edges = payload["edges"]
    if not isinstance(edges, list) or len(edges) == 0:
        raise InputError()

    valid_nodes = set(body_ids) | {"start_point", "rescue_point"}
    adj: Dict[str, List[Tuple[str, float]]] = defaultdict(list)
    edge_map: Dict[Tuple[str, str], float] = {}

    start_seen = False
    rescue_seen = False

    for edge in edges:
        ensure_keys(edge, {"from", "to", "distance"})
        frm = as_str(edge["from"])
        to = as_str(edge["to"])
        distance = as_float(edge["distance"], low=0.0)
        if distance <= 0:
            raise InputError()
        if frm not in valid_nodes or to not in valid_nodes:
            raise InputError()
        if (frm, to) in edge_map:
            raise InputError()

        edge_map[(frm, to)] = distance
        adj[frm].append((to, distance))

        if frm == "start_point" or to == "start_point":
            start_seen = True
        if frm == "rescue_point" or to == "rescue_point":
            rescue_seen = True

    if not start_seen or not rescue_seen:
        raise InputError()

    for node in adj:
        adj[node].sort(key=lambda x: x[1])

    return {
        "mass_shuttle": mass_shuttle,
        "mass_fuel_unit": mass_fuel_unit,
        "power_per_unit": power_per_unit,
        "oxygen_time": oxygen_time,
        "total_fuel": total_fuel,
        "fuel_consumption": fuel_consumption,
        "body_assists": body_assists,
        "adj": adj,
        "edge_map": edge_map,
    }


def max_orbit_fuel(
    total_fuel: int,
    mass_shuttle: float,
    mass_fuel_unit: float,
    fuel_consumption: float,
) -> int:
    # F + fuel_consumption * (mass_shuttle + F * mass_fuel_unit) <= total_fuel
    denom = 1.0 + fuel_consumption * mass_fuel_unit
    rhs = total_fuel - fuel_consumption * mass_shuttle
    if rhs < -EPS:
        return 0
    return max(0, int(math.floor((rhs + EPS) / denom)))


def enumerate_k_shortest_simple_paths(
    adj: Dict[str, List[Tuple[str, float]]],
    start: str,
    target: str,
    k: int = 300,
    max_expansions: int = 200000,
) -> List[Tuple[float, List[str]]]:
    heap: List[Tuple[float, int, List[str], frozenset[str]]] = []
    counter = 0
    heapq.heappush(heap, (0.0, counter, [start], frozenset({start})))
    result: List[Tuple[float, List[str]]] = []
    expansions = 0

    while heap and len(result) < k and expansions < max_expansions:
        dist_so_far, _, path, used = heapq.heappop(heap)
        node = path[-1]
        if node == target:
            result.append((dist_so_far, path))
            continue

        expansions += 1
        for nxt, w in adj.get(node, []):
            if nxt in used:
                continue
            counter += 1
            heapq.heappush(
                heap,
                (dist_so_far + w, counter, path + [nxt], used | {nxt}),
            )

    return result


def prune_cruise_states(states: List[CruiseState]) -> List[CruiseState]:
    by_fuel: Dict[int, List[CruiseState]] = defaultdict(list)
    for st in states:
        arr = by_fuel[st.fuel]
        dominated = False
        kept: List[CruiseState] = []
        for old in arr:
            if old.time <= st.time + EPS and old.speed >= st.speed - EPS:
                dominated = True
                kept.append(old)
            elif st.time <= old.time + EPS and st.speed >= old.speed - EPS:
                continue
            else:
                kept.append(old)
        if not dominated:
            kept.append(st)
        by_fuel[st.fuel] = kept

    out: List[CruiseState] = []
    for fuel_val, arr in by_fuel.items():
        arr.sort(key=lambda s: (s.time, -s.speed))
        out.extend(arr)
    return out


def finish_time_from_state(
    st: CruiseState,
    total_distance: float,
    engine: FlightEngine,
) -> Optional[float]:
    rem = total_distance - st.pos
    if rem < -EPS:
        return None
    if rem <= EPS:
        if st.speed <= EPS:
            return st.time
        return None
    if st.speed <= EPS:
        return None

    br = engine.brake_profile(st.speed, st.fuel)
    if br is None:
        return None
    brake_dist, brake_time = br
    if brake_dist > rem + EPS:
        return None

    coast_dist = max(0.0, rem - brake_dist)
    return st.time + coast_dist / st.speed + brake_time


def evaluate_path(
    path: List[str],
    total_distance: float,
    edge_map: Dict[Tuple[str, str], float],
    body_assists: Dict[str, List[GravityAssist]],
    engine: FlightEngine,
    fmax: int,
    oxygen_time: int,
) -> Optional[float]:
    # позиции внутренних тел на маршруте
    internal_nodes: List[Tuple[str, float]] = []
    pref = 0.0
    for i in range(1, len(path)):
        pref += edge_map[(path[i - 1], path[i])]
        if i < len(path) - 1:
            internal_nodes.append((path[i], pref))

    best_time: Optional[float] = None

    # Полный перебор количества топлива на орбите.
    # Это "честно", но может быть тяжеловато, если fmax очень большой.
    for initial_fuel in range(1, fmax + 1):
        accel_burns, dist_acc, speed_acc, fuel_after_acc = engine.accel_profile(initial_fuel, total_distance)

        if speed_acc <= EPS:
            continue

        states: List[CruiseState] = [
            CruiseState(pos=dist_acc, time=float(accel_burns), speed=speed_acc, fuel=fuel_after_acc)
        ]

        for st in states:
            ft = finish_time_from_state(st, total_distance, engine)
            if ft is not None and ft <= oxygen_time + EPS:
                if best_time is None or ft < best_time:
                    best_time = ft

        for node_id, node_pos in internal_nodes:
            if node_pos + EPS < dist_acc:
                continue

            next_states: List[CruiseState] = []
            options = body_assists.get(node_id, [])

            for st in states:
                if node_pos + EPS < st.pos:
                    continue
                if st.speed <= EPS:
                    continue

                # можно ли вообще долететь до этого тела, не начиная тормозить раньше?
                br = engine.brake_profile(st.speed, st.fuel)
                if br is None:
                    continue
                brake_dist_now, _ = br
                if brake_dist_now > (total_distance - node_pos) + EPS:
                    continue

                travel_dist = max(0.0, node_pos - st.pos)
                arrive_time = st.time + travel_dist / st.speed

                # без маневра
                next_states.append(
                    CruiseState(
                        pos=node_pos,
                        time=arrive_time,
                        speed=st.speed,
                        fuel=st.fuel,
                    )
                )

                # с одним маневром
                for ga in options:
                    if ga.fuel_consumption <= st.fuel:
                        next_states.append(
                            CruiseState(
                                pos=node_pos,
                                time=arrive_time + ga.time_to_execute,
                                speed=st.speed + ga.velocity_gain,
                                fuel=st.fuel - ga.fuel_consumption,
                            )
                        )

            states = prune_cruise_states(next_states)
            if not states:
                break

            for st in states:
                ft = finish_time_from_state(st, total_distance, engine)
                if ft is not None and ft <= oxygen_time + EPS:
                    if best_time is None or ft < best_time:
                        best_time = ft

    return best_time


def solve_task1(data: Dict[str, Any]) -> Dict[str, Any]:
    fmax = max_orbit_fuel(
        data["total_fuel"],
        data["mass_shuttle"],
        data["mass_fuel_unit"],
        data["fuel_consumption"],
    )
    if fmax <= 0:
        return {"can_reach": False}

    adj = data["adj"]
    edge_map = data["edge_map"]
    body_assists = data["body_assists"]

    paths = enumerate_k_shortest_simple_paths(adj, "start_point", "rescue_point", k=300)
    if not paths:
        return {"can_reach": False}

    engine = FlightEngine(
        mass_shuttle=data["mass_shuttle"],
        mass_fuel_unit=data["mass_fuel_unit"],
        power_per_unit=data["power_per_unit"],
    )

    best_route: Optional[List[str]] = None
    best_time: Optional[float] = None

    for total_distance, path in paths:
        t = evaluate_path(
            path=path,
            total_distance=total_distance,
            edge_map=edge_map,
            body_assists=body_assists,
            engine=engine,
            fmax=fmax,
            oxygen_time=data["oxygen_time"],
        )
        if t is None:
            continue
        if best_time is None or t < best_time:
            best_time = t
            best_route = path

    if best_time is None or best_route is None:
        return {"can_reach": False}

    return {
        "can_reach": True,
        "min_flight_time": float(f"{best_time:.1f}"),
        "route": best_route,
    }


# =========================================================
# TASK 2
# =========================================================

@dataclass(frozen=True)
class HarmonicComp:
    radius: float
    omega: float
    phi: float
    a_cos: float
    a_sin: float
    b_cos: float
    b_sin: float


@dataclass
class CompiledBody:
    body_id: str
    radius: float
    center_along: float
    center_cross: float
    comps: List[HarmonicComp]
    speed_bound: float
    reach_bound: float

    def eval_ab(self, t: float) -> Tuple[float, float]:
        a = self.center_along
        b = self.center_cross
        for c in self.comps:
            ang = math.fmod(c.phi + c.omega * t, TAU)
            cs = math.cos(ang)
            sn = math.sin(ang)
            a += c.a_cos * cs + c.a_sin * sn
            b += c.b_cos * cs + c.b_sin * sn
        return a, b

    def blocked(self, t: float) -> bool:
        a, b = self.eval_ab(t)
        return a > 0.0 and abs(b) <= self.radius + 1e-12

    def may_ever_block_by_bound(self) -> bool:
        if self.center_along + self.reach_bound <= 0:
            return False
        if abs(self.center_cross) > self.reach_bound + self.radius:
            return False
        return True


def parse_iso_z(value: Any) -> datetime:
    s = as_str(value)
    try:
        return datetime.strptime(s, "%Y-%m-%dT%H:%M:%SZ")
    except Exception as e:
        raise InputError() from e


def parse_task2(payload: Any) -> Dict[str, Any]:
    ensure_keys(payload, {"target_star_vector", "celestial_bodies", "observation_params"})

    ensure_keys(payload["target_star_vector"], {"x", "y"})
    vx = as_float(payload["target_star_vector"]["x"])
    vy = as_float(payload["target_star_vector"]["y"])
    norm = math.hypot(vx, vy)
    if norm <= 0:
        raise InputError()
    ux, uy = vx / norm, vy / norm

    bodies = payload["celestial_bodies"]
    if not isinstance(bodies, list) or len(bodies) > 100:
        raise InputError()

    raw: Dict[str, Dict[str, Any]] = {}
    for body in bodies:
        if not isinstance(body, dict):
            raise InputError()
        btype = body.get("type")
        if btype == "star":
            ensure_keys(body, {"type", "id", "position", "radius"})
            body_id = as_str(body["id"])
            if body_id in raw or body_id == "Atlas":
                raise InputError()
            ensure_keys(body["position"], {"x", "y"})
            x = as_float(body["position"]["x"])
            y = as_float(body["position"]["y"])
            radius = as_float(body["radius"], low=0.0)
            if radius <= 0:
                raise InputError()
            raw[body_id] = {
                "type": "star",
                "id": body_id,
                "x": x,
                "y": y,
                "radius": radius,
            }

        elif btype in {"planet", "moon"}:
            ensure_keys(body, {
                "type",
                "id",
                "parent_id",
                "orbit_radius",
                "angular_velocity",
                "initial_angle",
                "radius",
                "rotation_clockwise",
            })
            body_id = as_str(body["id"])
            if body_id in raw or body_id == "Atlas":
                raise InputError()
            parent_id = as_str(body["parent_id"])
            orbit_radius = as_float(body["orbit_radius"], low=0.0)
            angular_velocity = as_float(body["angular_velocity"], low=0.0)
            if orbit_radius <= 0 or angular_velocity <= 0:
                raise InputError()
            initial_angle = as_float(body["initial_angle"], low=0.0, high=360.0)
            if not (0.0 <= initial_angle < 360.0):
                raise InputError()
            radius = as_float(body["radius"], low=0.0)
            if radius <= 0:
                raise InputError()
            clockwise = as_bool(body["rotation_clockwise"])
            raw[body_id] = {
                "type": btype,
                "id": body_id,
                "parent_id": parent_id,
                "orbit_radius": orbit_radius,
                "angular_velocity": angular_velocity,
                "initial_angle": initial_angle,
                "radius": radius,
                "rotation_clockwise": clockwise,
            }
        else:
            raise InputError()

    ensure_keys(payload["observation_params"], {"start_time", "required_transmission_time"})
    _ = parse_iso_z(payload["observation_params"]["start_time"])
    required_transmission_time = as_int(payload["observation_params"]["required_transmission_time"], low=1, high=10**6)

    compiled: Dict[str, CompiledBody] = {}
    visiting: set[str] = set()

    def make_comp(orbit_radius: float, omega_deg: float, phi_deg: float, clockwise: bool) -> HarmonicComp:
        omega = math.radians(omega_deg)
        if clockwise:
            omega = -omega
        phi = math.radians(phi_deg)
        return HarmonicComp(
            radius=orbit_radius,
            omega=omega,
            phi=phi,
            a_cos=orbit_radius * ux,
            a_sin=orbit_radius * uy,
            b_cos=-orbit_radius * uy,
            b_sin=orbit_radius * ux,
        )

    def compile_one(body_id: str) -> CompiledBody:
        if body_id in compiled:
            return compiled[body_id]
        if body_id in visiting:
            raise InputError()
        visiting.add(body_id)

        body = raw.get(body_id)
        if body is None:
            raise InputError()

        if body["type"] == "star":
            cx, cy = body["x"], body["y"]
            center_along = cx * ux + cy * uy
            center_cross = cx * (-uy) + cy * ux
            result = CompiledBody(
                body_id=body_id,
                radius=body["radius"],
                center_along=center_along,
                center_cross=center_cross,
                comps=[],
                speed_bound=0.0,
                reach_bound=0.0,
            )
        elif body["type"] == "planet":
            parent_id = body["parent_id"]
            parent = raw.get(parent_id)
            if parent is None or parent["type"] != "star":
                raise InputError()
            p = compile_one(parent_id)
            comp = make_comp(
                body["orbit_radius"],
                body["angular_velocity"],
                body["initial_angle"],
                body["rotation_clockwise"],
            )
            result = CompiledBody(
                body_id=body_id,
                radius=body["radius"],
                center_along=p.center_along,
                center_cross=p.center_cross,
                comps=[comp],
                speed_bound=abs(comp.omega) * comp.radius,
                reach_bound=comp.radius,
            )
        else:
            # moon
            parent_id = body["parent_id"]
            comp = make_comp(
                body["orbit_radius"],
                body["angular_velocity"],
                body["initial_angle"],
                body["rotation_clockwise"],
            )
            if parent_id == "Atlas":
                result = CompiledBody(
                    body_id=body_id,
                    radius=body["radius"],
                    center_along=0.0,
                    center_cross=0.0,
                    comps=[comp],
                    speed_bound=abs(comp.omega) * comp.radius,
                    reach_bound=comp.radius,
                )
            else:
                parent = raw.get(parent_id)
                if parent is None or parent["type"] != "planet":
                    raise InputError()
                p = compile_one(parent_id)
                result = CompiledBody(
                    body_id=body_id,
                    radius=body["radius"],
                    center_along=p.center_along,
                    center_cross=p.center_cross,
                    comps=p.comps + [comp],
                    speed_bound=p.speed_bound + abs(comp.omega) * comp.radius,
                    reach_bound=p.reach_bound + comp.radius,
                )

        visiting.remove(body_id)
        compiled[body_id] = result
        return result

    compiled_bodies = [compile_one(body_id) for body_id in raw.keys()]

    return {
        "compiled_bodies": compiled_bodies,
        "required_transmission_time": required_transmission_time,
    }


def collect_intervals_for_body(body: CompiledBody, horizon: float) -> List[Tuple[float, float]]:
    if horizon <= 0:
        return []
    if not body.may_ever_block_by_bound():
        return []

    out: List[Tuple[float, float]] = []

    a0, b0 = body.eval_ab(0.0)
    a1, b1 = body.eval_ab(horizon)
    stack: List[Tuple[float, float, float, float, float, float]] = [(0.0, horizon, a0, b0, a1, b1)]

    while stack:
        t0, t1, a_left, b_left, a_right, b_right = stack.pop()
        h = t1 - t0
        L = body.speed_bound

        # гарантированно всегда сзади
        if max(a_left, a_right) + L * h <= 0.0:
            continue

        m_left = abs(b_left) - body.radius
        m_right = abs(b_right) - body.radius

        # гарантированно вне полосы
        if min(m_left, m_right) - L * h > 0.0:
            continue

        # гарантированно блокирует весь интервал
        if min(a_left, a_right) - L * h > 0.0 and max(m_left, m_right) + L * h <= 0.0:
            out.append((t0, t1))
            continue

        if h <= VIS_EPS:
            # переоценка маленьким кусочком в сторону безопасности
            out.append((t0, t1))
            continue

        tm = 0.5 * (t0 + t1)
        am, bm = body.eval_ab(tm)
        stack.append((tm, t1, am, bm, a_right, b_right))
        stack.append((t0, tm, a_left, b_left, am, bm))

    return out


def merge_intervals(intervals: List[Tuple[float, float]]) -> List[Tuple[float, float]]:
    if not intervals:
        return []
    intervals.sort()
    merged: List[List[float]] = [[intervals[0][0], intervals[0][1]]]
    for l, r in intervals[1:]:
        if l <= merged[-1][1] + VIS_EPS:
            merged[-1][1] = max(merged[-1][1], r)
        else:
            merged.append([l, r])
    return [(l, r) for l, r in merged]


def inspect_visibility(
    merged: List[Tuple[float, float]],
    horizon: float,
    required: int,
) -> Dict[str, Any]:
    if not merged:
        # весь [0, horizon] пока виден
        tail_end = int(math.floor(horizon))
        if tail_end >= required:
            return {"kind": "tail", "start": 0}
        return {"kind": "none"}

    prev_end: Optional[float] = None

    for l, r in merged:
        gap_start = 0 if prev_end is None else int(math.floor(prev_end)) + 1
        gap_end = int(math.floor(l))
        duration = gap_end - gap_start
        if duration >= required:
            return {
                "kind": "closed",
                "start": gap_start,
                "duration": duration,
            }
        prev_end = r

    gap_start = 0 if prev_end is None else int(math.floor(prev_end)) + 1
    gap_end = int(math.floor(horizon))
    duration = gap_end - gap_start
    if duration >= required:
        return {"kind": "tail", "start": gap_start}
    return {"kind": "none"}


def solve_task2(data: Dict[str, Any]) -> Dict[str, Any]:
    bodies: List[CompiledBody] = data["compiled_bodies"]
    required: int = data["required_transmission_time"]

    # Если по грубой геометрической оценке никто вообще не может закрыть луч
    if all(not b.may_ever_block_by_bound() for b in bodies):
        return {
            "found": True,
            "next_fitting_interval_in": 0,
            "interval_duration": "inf",
        }

    wait_cap = 1_000_000_000 + required
    horizon = float(min(wait_cap, max(10000, required * 2)))

    while True:
        intervals: List[Tuple[float, float]] = []
        for body in bodies:
            intervals.extend(collect_intervals_for_body(body, horizon))

        merged = merge_intervals(intervals)
        info = inspect_visibility(merged, horizon, required)

        if info["kind"] == "closed":
            return {
                "found": True,
                "next_fitting_interval_in": info["start"],
                "interval_duration": info["duration"],
            }

        if info["kind"] == "tail":
            # Нужно расширять горизонт, чтобы узнать реальную длину интервала,
            # если только не удалось строго понять, что затмений вообще не бывает.
            if horizon >= 1e12:
                # аварийный предохранитель
                return {
                    "found": True,
                    "next_fitting_interval_in": info["start"],
                    "interval_duration": "inf",
                }
            horizon *= 2.0
            continue

        # fitting gap пока нет
        if horizon >= wait_cap:
            return {"found": False}

        horizon = min(float(wait_cap), horizon * 2.0)


# =========================================================
# TASK 3
# =========================================================

class DSU:
    def __init__(self, n: int):
        self.p = list(range(n))
        self.r = [0] * n

    def find(self, x: int) -> int:
        while self.p[x] != x:
            self.p[x] = self.p[self.p[x]]
            x = self.p[x]
        return x

    def union(self, a: int, b: int) -> None:
        a = self.find(a)
        b = self.find(b)
        if a == b:
            return
        if self.r[a] < self.r[b]:
            a, b = b, a
        self.p[b] = a
        if self.r[a] == self.r[b]:
            self.r[a] += 1


def parse_task3(payload: Any) -> Dict[str, Any]:
    ensure_keys(payload, {"stars", "cluster_params", "target_constellation"})

    stars = payload["stars"]
    if not isinstance(stars, list) or len(stars) > 1000:
        raise InputError()

    star_names = []
    coords = []
    seen_names = set()

    for star in stars:
        ensure_keys(star, {"name", "x", "y", "z"})
        name = as_str(star["name"])
        if name in seen_names:
            raise InputError()
        seen_names.add(name)
        x = as_float(star["x"])
        y = as_float(star["y"])
        z = as_float(star["z"])
        star_names.append(name)
        coords.append((x, y, z))

    ensure_keys(payload["cluster_params"], {"min_size", "max_size", "max_neighbor_distance"})
    min_size = as_int(payload["cluster_params"]["min_size"], low=1)
    max_size = as_int(payload["cluster_params"]["max_size"], low=1)
    if min_size > max_size:
        raise InputError()
    max_neighbor_distance = as_float(payload["cluster_params"]["max_neighbor_distance"], low=0.0)
    if max_neighbor_distance < 0:
        raise InputError()

    ensure_keys(payload["target_constellation"], {"edges"})
    target_edges_raw = payload["target_constellation"]["edges"]
    if not isinstance(target_edges_raw, list):
        raise InputError()

    pairs = set()
    target_edges: List[Tuple[int, int, float]] = []
    max_v = -1
    vertices_used = set()

    for e in target_edges_raw:
        ensure_keys(e, {"from", "to", "distance"})
        u = as_int(e["from"], low=0)
        v = as_int(e["to"], low=0)
        if u == v:
            raise InputError()
        d = as_float(e["distance"], low=0.0)
        if d <= 0:
            raise InputError()
        pair = (min(u, v), max(u, v))
        if pair in pairs:
            raise InputError()
        pairs.add(pair)
        target_edges.append((u, v, d))
        vertices_used.add(u)
        vertices_used.add(v)
        max_v = max(max_v, u, v)

    target_n = max_v + 1
    if target_n < 2 or target_n > 50:
        raise InputError()
    if vertices_used != set(range(target_n)):
        raise InputError()
    if len(target_edges) != target_n - 1:
        raise InputError()

    # проверка, что это дерево
    dsu = DSU(target_n)
    for u, v, _ in target_edges:
        if dsu.find(u) == dsu.find(v):
            raise InputError()
        dsu.union(u, v)
    root = dsu.find(0)
    for i in range(target_n):
        if dsu.find(i) != root:
            raise InputError()

    return {
        "star_names": star_names,
        "coords": coords,
        "min_size": min_size,
        "max_size": max_size,
        "max_neighbor_distance": max_neighbor_distance,
        "target_edges": target_edges,
        "target_n": target_n,
    }


def dist3(a: Tuple[float, float, float], b: Tuple[float, float, float]) -> float:
    return math.dist(a, b)


def build_components(coords: List[Tuple[float, float, float]], dmax: float) -> List[List[int]]:
    n = len(coords)
    dsu = DSU(n)
    limit = dmax + EPS
    for i in range(n):
        for j in range(i + 1, n):
            if dist3(coords[i], coords[j]) <= limit:
                dsu.union(i, j)

    groups: Dict[int, List[int]] = defaultdict(list)
    for i in range(n):
        groups[dsu.find(i)].append(i)
    return list(groups.values())


def prim_mst(component: List[int], coords: List[Tuple[float, float, float]]) -> List[Tuple[int, int, float]]:
    k = len(component)
    used = [False] * k
    min_d = [float("inf")] * k
    parent = [-1] * k
    min_d[0] = 0.0

    for _ in range(k):
        v = -1
        best = float("inf")
        for i in range(k):
            if not used[i] and min_d[i] < best:
                best = min_d[i]
                v = i
        if v == -1:
            raise InputError()

        used[v] = True
        for u in range(k):
            if used[u]:
                continue
            d = dist3(coords[component[v]], coords[component[u]])
            if d < min_d[u]:
                min_d[u] = d
                parent[u] = v

    edges: List[Tuple[int, int, float]] = []
    for u in range(1, k):
        edges.append((parent[u], u, min_d[u]))
    return edges


def distance_groups(lengths: List[float]) -> Tuple[List[int], List[int]]:
    order = sorted(range(len(lengths)), key=lambda i: lengths[i])
    labels = [-1] * len(lengths)
    group_sizes: List[int] = []
    g = 0
    i = 0
    while i < len(order):
        j = i + 1
        while j < len(order) and math.isclose(lengths[order[i]], lengths[order[j]], rel_tol=1e-9, abs_tol=1e-9):
            j += 1
        group_sizes.append(j - i)
        for t in range(i, j):
            labels[order[t]] = g
        g += 1
        i = j
    return group_sizes, labels


def candidate_labels_from_target_groups(candidate_lengths: List[float], target_group_sizes: List[int]) -> List[int]:
    order = sorted(range(len(candidate_lengths)), key=lambda i: candidate_lengths[i])
    labels = [-1] * len(candidate_lengths)
    pos = 0
    for group_id, size in enumerate(target_group_sizes):
        for _ in range(size):
            if pos >= len(order):
                raise InputError()
            labels[order[pos]] = group_id
            pos += 1
    if pos != len(order):
        raise InputError()
    return labels


def build_adj(n: int, edges: List[Tuple[int, int, int]]) -> List[List[Tuple[int, int]]]:
    adj: List[List[Tuple[int, int]]] = [[] for _ in range(n)]
    for u, v, color in edges:
        adj[u].append((v, color))
        adj[v].append((u, color))
    return adj


def tree_centers(adj: List[List[Tuple[int, int]]]) -> List[int]:
    n = len(adj)
    deg = [len(x) for x in adj]
    leaves = [i for i in range(n) if deg[i] <= 1]
    removed = len(leaves)

    while removed < n:
        new_leaves = []
        for leaf in leaves:
            for nei, _ in adj[leaf]:
                deg[nei] -= 1
                if deg[nei] == 1:
                    new_leaves.append(nei)
        if removed + len(new_leaves) >= n:
            return new_leaves
        removed += len(new_leaves)
        leaves = new_leaves

    return leaves


def rooted_code(
    adj: List[List[Tuple[int, int]]],
    node: int,
    parent: int,
    memo: Dict[Tuple[int, int], Tuple],
) -> Tuple:
    key = (node, parent)
    if key in memo:
        return memo[key]
    parts = []
    for nei, color in adj[node]:
        if nei == parent:
            continue
        parts.append((color, rooted_code(adj, nei, node, memo)))
    parts.sort(key=lambda x: (x[0], x[1]))
    memo[key] = tuple(parts)
    return memo[key]


def edge_color(adj: List[List[Tuple[int, int]]], a: int, b: int) -> int:
    for nei, color in adj[a]:
        if nei == b:
            return color
    raise InputError()


def map_rooted(
    adj_t: List[List[Tuple[int, int]]],
    adj_c: List[List[Tuple[int, int]]],
    u: int,
    pu: int,
    v: int,
    pv: int,
    memo_t: Dict[Tuple[int, int], Tuple],
    memo_c: Dict[Tuple[int, int], Tuple],
    mapping: Dict[int, int],
) -> bool:
    if u in mapping and mapping[u] != v:
        return False
    mapping[u] = v

    ch_t = []
    for nei, color in adj_t[u]:
        if nei == pu:
            continue
        ch_t.append((color, rooted_code(adj_t, nei, u, memo_t), nei))
    ch_c = []
    for nei, color in adj_c[v]:
        if nei == pv:
            continue
        ch_c.append((color, rooted_code(adj_c, nei, v, memo_c), nei))

    ch_t.sort(key=lambda x: (x[0], x[1]))
    ch_c.sort(key=lambda x: (x[0], x[1]))

    if len(ch_t) != len(ch_c):
        return False
    if [(a, b) for a, b, _ in ch_t] != [(a, b) for a, b, _ in ch_c]:
        return False

    for (_, _, nt), (_, _, nc) in zip(ch_t, ch_c):
        if not map_rooted(adj_t, adj_c, nt, u, nc, v, memo_t, memo_c, mapping):
            return False
    return True


def match_colored_trees(
    adj_t: List[List[Tuple[int, int]]],
    adj_c: List[List[Tuple[int, int]]],
) -> Optional[Dict[int, int]]:
    ct = tree_centers(adj_t)
    cc = tree_centers(adj_c)
    if len(ct) != len(cc):
        return None

    memo_t: Dict[Tuple[int, int], Tuple] = {}
    memo_c: Dict[Tuple[int, int], Tuple] = {}

    if len(ct) == 1:
        rt = ct[0]
        rc = cc[0]
        if rooted_code(adj_t, rt, -1, memo_t) != rooted_code(adj_c, rc, -1, memo_c):
            return None
        mapping: Dict[int, int] = {}
        if map_rooted(adj_t, adj_c, rt, -1, rc, -1, memo_t, memo_c, mapping):
            return mapping
        return None

    # два центра
    at, bt = ct
    x, y = cc
    if edge_color(adj_t, at, bt) != edge_color(adj_c, x, y):
        return None

    for cx, cy in [(x, y), (y, x)]:
        if rooted_code(adj_t, at, bt, memo_t) != rooted_code(adj_c, cx, cy, memo_c):
            continue
        if rooted_code(adj_t, bt, at, memo_t) != rooted_code(adj_c, cy, cx, memo_c):
            continue

        mapping: Dict[int, int] = {at: cx, bt: cy}
        ok1 = map_rooted(adj_t, adj_c, at, bt, cx, cy, memo_t, memo_c, mapping)
        ok2 = map_rooted(adj_t, adj_c, bt, at, cy, cx, memo_t, memo_c, mapping)
        if ok1 and ok2:
            return mapping

    return None


def solve_task3(data: Dict[str, Any]) -> Dict[str, Any]:
    star_names = data["star_names"]
    coords = data["coords"]
    min_size = data["min_size"]
    max_size = data["max_size"]
    dmax = data["max_neighbor_distance"]
    target_edges = data["target_edges"]
    target_n = data["target_n"]

    # target tree colors
    target_lengths = [d for _, _, d in target_edges]
    target_group_sizes, target_labels = distance_groups(target_lengths)
    target_colored_edges = [
        (u, v, target_labels[i])
        for i, (u, v, _) in enumerate(target_edges)
    ]
    target_adj = build_adj(target_n, target_colored_edges)

    components = build_components(coords, dmax)

    matches: List[List[str]] = []

    for comp in components:
        if len(comp) < min_size or len(comp) > max_size:
            continue
        if len(comp) != target_n:
            continue

        mst = prim_mst(comp, coords)
        cand_lengths = [d for _, _, d in mst]
        cand_labels = candidate_labels_from_target_groups(cand_lengths, target_group_sizes)
        cand_colored_edges = [
            (u, v, cand_labels[i])
            for i, (u, v, _) in enumerate(mst)
        ]
        cand_adj = build_adj(target_n, cand_colored_edges)

        mapping = match_colored_trees(target_adj, cand_adj)
        if mapping is None:
            continue

        matched_stars = [None] * target_n
        for tv in range(target_n):
            cv = mapping[tv]
            matched_stars[tv] = star_names[comp[cv]]
        matches.append(matched_stars)

    if len(matches) != 1:
        return {"found": False}

    return {
        "found": True,
        "matched_stars": matches[0],
    }


# =========================================================
# HTTP
# =========================================================

@app.post("/api/v1/robinson_cruise")
async def robinson_cruise(request: Request):
    try:
        payload = await request.json()
    except Exception:
        return bad_request()
    try:
        data = parse_task1(payload)
    except InputError:
        return bad_request()
    return JSONResponse(status_code=200, content=solve_task1(data))


@app.post("/api/v1/star_visibility")
async def star_visibility(request: Request):
    try:
        payload = await request.json()
    except Exception:
        return bad_request()
    try:
        data = parse_task2(payload)
    except InputError:
        return bad_request()
    return JSONResponse(status_code=200, content=solve_task2(data))


@app.post("/api/v1/constellation_finder")
async def constellation_finder(request: Request):
    try:
        payload = await request.json()
    except Exception:
        return bad_request()
    try:
        data = parse_task3(payload)
    except InputError:
        return bad_request()
    return JSONResponse(status_code=200, content=solve_task3(data))


if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="0.0.0.0", port=8000, app_dir="src")