"""Hill-climbing solver for slab placement.

We don't try to enumerate the full N! search space — for a 34-slot grid the
number of placements is enormous. Instead we use multi-start random + first-
improvement hill climbing with swap/rotate/move moves. This converges quickly
to good local optima and is plenty fast for typical inputs (≤ 60 slabs).
"""

from __future__ import annotations

import random
import time
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple

from .slabs import (
    GRID_COLS,
    GridConfig,
    Placement,
    SLABS_BY_VALUE,
    all_slot_ids,
    generate_grid_config,
    score,
)


@dataclass
class Solution:
    placements: List[Placement] = field(default_factory=list)
    grid: GridConfig = field(default_factory=list)
    score: int = 0
    slot_num: int = 34


def _initial(
    slab_values: List[str], grid: GridConfig, rng: random.Random
) -> List[Placement]:
    slots = all_slot_ids(grid)
    rng.shuffle(slots)
    placements: List[Placement] = []
    for v, sid in zip(slab_values, slots):
        slab = SLABS_BY_VALUE.get(v)
        rot = rng.randint(0, 3) if slab and slab.rotate else 0
        placements.append(Placement(value=v, slot_id=sid, rotation=rot))
    return placements


def _hill_climb(
    placements: List[Placement],
    grid: GridConfig,
    rng: random.Random,
    max_iters: int,
    time_limit: float,
) -> Tuple[List[Placement], int]:
    cur = [Placement(p.value, p.slot_id, p.rotation) for p in placements]
    cur_score = score(cur, grid)
    used = {p.slot_id for p in cur}
    free = [sid for sid in all_slot_ids(grid) if sid not in used]
    start = time.monotonic()

    for _ in range(max_iters):
        if time.monotonic() - start > time_limit:
            break
        move = rng.random()
        if move < 0.45 and len(cur) > 1:
            # swap two placements
            i, j = rng.sample(range(len(cur)), 2)
            cur[i].slot_id, cur[j].slot_id = cur[j].slot_id, cur[i].slot_id
            ns = score(cur, grid)
            if ns > cur_score:
                cur_score = ns
            else:
                cur[i].slot_id, cur[j].slot_id = cur[j].slot_id, cur[i].slot_id
        elif move < 0.75 and free:
            # move one placement to a free slot
            i = rng.randint(0, len(cur) - 1)
            old = cur[i].slot_id
            new_idx = rng.randint(0, len(free) - 1)
            new_slot = free[new_idx]
            cur[i].slot_id = new_slot
            ns = score(cur, grid)
            if ns > cur_score:
                cur_score = ns
                free[new_idx] = old
            else:
                cur[i].slot_id = old
        else:
            # rotate one placement
            i = rng.randint(0, len(cur) - 1)
            slab = SLABS_BY_VALUE.get(cur[i].value)
            if not slab or not slab.rotate:
                continue
            old_rot = cur[i].rotation
            cur[i].rotation = (old_rot + rng.choice([1, 2, 3])) % 4
            ns = score(cur, grid)
            if ns > cur_score:
                cur_score = ns
            else:
                cur[i].rotation = old_rot

    return cur, cur_score


def solve(
    slab_values: List[str],
    slot_num: int = 34,
    *,
    restarts: int = 6,
    iters_per_restart: int = 8000,
    time_limit: float = 4.0,
    seed: Optional[int] = None,
) -> Solution:
    """Find a good placement for the given slabs.

    `slab_values` is a list of slab `value` strings (duplicates allowed).
    Slabs that don't fit (more slabs than slots) are dropped from the tail.
    """
    grid = generate_grid_config(slot_num)
    capacity = sum(r["cols"] for r in grid)
    if len(slab_values) > capacity:
        slab_values = slab_values[:capacity]
    if not slab_values:
        return Solution([], grid, 0, slot_num)

    rng = random.Random(seed)
    best: List[Placement] = []
    best_score = -10**9

    per_restart_budget = time_limit / max(restarts, 1)
    for _ in range(restarts):
        init = _initial(slab_values, grid, rng)
        sol, sc = _hill_climb(init, grid, rng, iters_per_restart, per_restart_budget)
        if sc > best_score:
            best_score = sc
            best = [Placement(p.value, p.slot_id, p.rotation) for p in sol]

    return Solution(best, grid, best_score, slot_num)