"""Auto-layout mesh generation for CUBE NOC router mesh.

Generates cube_mesh.yaml describing the internal router grid, PE/UCIe/XBAR
attachments, and HBM exclusion zone.  The file is cached with a source_hash
so it is only regenerated when relevant topology parameters change.

Algorithm (final, per Phase 1 design iteration):
  cols          = physical_cols  (PE x-positions + relay cols for max_spacing)
  rows_per_half = ceil(n_connections / 2)
  total_rows    = rows_per_half * 2 + 2   (+ 2 HBM rows)
  PEs: 1 PE per row when rows available, corners at fixed positions
  Hot path: min_connections = max(n_connections, 2)
"""
from __future__ import annotations

import hashlib
import math
from pathlib import Path
from typing import Any

import yaml


# ── Public API ────────────────────────────────────────────────────────


def ensure_mesh_file(cube_spec: dict, mesh_path: Path) -> dict:
    """Generate cube_mesh.yaml if needed, return parsed mesh dict."""
    source_hash = _compute_source_hash(cube_spec)

    if mesh_path.exists():
        existing = yaml.safe_load(mesh_path.read_text(encoding="utf-8"))
        if existing and existing.get("source_hash") == source_hash:
            return existing

    mesh = _generate_mesh(cube_spec, source_hash)
    mesh_path.write_text(
        yaml.dump(mesh, default_flow_style=False, sort_keys=False),
        encoding="utf-8",
    )
    return mesh


# ── Hash ──────────────────────────────────────────────────────────────


def _compute_source_hash(cube_spec: dict) -> str:
    """Hash relevant topology params that determine mesh layout."""
    relevant = {
        "geometry": cube_spec["geometry"],
        "pe_layout": cube_spec["pe_layout"],
        "ucie_n_connections": cube_spec["ucie"]["n_connections"],
        "hbm_mapping_mode": cube_spec.get("memory_map", {}).get(
            "hbm_mapping_mode", "n_to_one"
        ),
        "placement": cube_spec.get("placement", {}),
    }
    raw = yaml.dump(relevant, sort_keys=True)
    return hashlib.sha256(raw.encode()).hexdigest()[:16]


# ── Layout helpers ────────────────────────────────────────────────────


def _corner_pe_positions(
    cube_w: float, cube_h: float
) -> dict[str, list[tuple[float, float]]]:
    """PE center positions per corner, relative to cube origin."""
    return {
        "NW": [(1.5, 1.5), (4.5, 1.5)],
        "NE": [(cube_w - 4.5, 1.5), (cube_w - 1.5, 1.5)],
        "SW": [(1.5, cube_h - 1.5), (4.5, cube_h - 1.5)],
        "SE": [(cube_w - 4.5, cube_h - 1.5), (cube_w - 1.5, cube_h - 1.5)],
    }


def _compute_col_positions(cube_w: float, pe_positions: dict) -> list[float]:
    """Compute X positions for grid columns based on PE positions + relay spacing."""
    xs: set[float] = set()
    for positions in pe_positions.values():
        for x, _y in positions:
            xs.add(x)

    sorted_xs = sorted(xs)
    # Insert relay columns for gaps > max_spacing (3mm)
    max_spacing = 3.0
    result: list[float] = []
    for i, x in enumerate(sorted_xs):
        if i > 0:
            gap = x - result[-1]
            while gap > max_spacing + 0.01:
                mid = result[-1] + max_spacing
                if mid < x - 0.5:
                    result.append(round(mid, 1))
                    gap = x - result[-1]
                else:
                    break
        result.append(x)
    return result


def _compute_row_positions(
    cube_h: float, n_connections: int, pe_positions: dict
) -> tuple[list[float], int]:
    """Compute Y positions for grid rows.

    Returns (y_positions, rows_per_half).
    Layout: [top PE rows] [HBM row top] [HBM row bot] [bottom PE rows]
    """
    n_conn = max(n_connections, 2)  # hot path minimum
    rows_per_half = math.ceil(n_conn / 2)

    # Top half: evenly spaced from top PE y to just above HBM zone
    top_pe_y = 1.5
    hbm_gap = 1.5  # minimum gap between PE rows and HBM rows
    hbm_top_y = cube_h / 2 - 1.5  # ~5.5 for h=14
    hbm_bot_y = cube_h / 2 + 1.5  # ~8.5 for h=14
    bot_pe_y = cube_h - 1.5

    top_rows: list[float] = []
    if rows_per_half == 1:
        top_rows = [top_pe_y]
    else:
        # End before HBM zone with gap
        top_end = hbm_top_y - hbm_gap
        step = (top_end - top_pe_y) / (rows_per_half - 1) if rows_per_half > 1 else 0
        for i in range(rows_per_half):
            top_rows.append(round(top_pe_y + i * step, 1))

    # HBM rows
    hbm_rows = [round(hbm_top_y, 1), round(hbm_bot_y, 1)]

    # Bottom half: mirror of top, start after HBM zone with gap
    bot_rows: list[float] = []
    if rows_per_half == 1:
        bot_rows = [bot_pe_y]
    else:
        bot_start = hbm_bot_y + hbm_gap
        step = (bot_pe_y - bot_start) / (rows_per_half - 1) if rows_per_half > 1 else 0
        for i in range(rows_per_half):
            bot_rows.append(round(bot_start + i * step, 1))

    return top_rows + hbm_rows + bot_rows, rows_per_half


# ── Mesh generation ──────────────────────────────────────────────────


def _generate_mesh(cube_spec: dict, source_hash: str) -> dict:
    geom = cube_spec["geometry"]
    cube_w = geom["cube_mm"]["w"]
    cube_h = geom["cube_mm"]["h"]
    pe_layout = cube_spec["pe_layout"]
    corners = pe_layout["corners"]
    pe_per_corner = pe_layout["pe_per_corner"]
    n_connections = cube_spec["ucie"]["n_connections"]

    pe_positions = _corner_pe_positions(cube_w, cube_h)
    col_xs = _compute_col_positions(cube_w, pe_positions)
    row_ys, rows_per_half = _compute_row_positions(
        cube_h, n_connections, pe_positions
    )
    n_rows = len(row_ys)
    n_cols = len(col_xs)

    # HBM exclusion zone: center rows, center cols
    hbm_row_start = rows_per_half      # first HBM row index
    hbm_row_end = rows_per_half + 1    # last HBM row index (inclusive)
    hbm_col_start = n_cols // 2 - 1    # center-left col
    hbm_col_end = n_cols // 2          # center-right col

    # Build routers dict
    routers: dict[str, Any] = {}
    for r in range(n_rows):
        for c in range(n_cols):
            key = f"r{r}c{c}"
            if (hbm_row_start <= r <= hbm_row_end
                    and hbm_col_start <= c <= hbm_col_end):
                routers[key] = None  # HBM excluded
            else:
                routers[key] = {
                    "pos_mm": [col_xs[c], row_ys[r]],
                    "attach": [],
                }

    # PE assignment: map each PE to a router based on corner and position.
    # All PEs in the same corner share one row. Corner order determines row:
    #   Top half:  NW → row 0, NE → row 1
    #   Bottom half: SW → row 4, SE → row 5  (for rows_per_half=2)
    pe_idx = 0
    top_pe_routers: list[str] = []
    bot_pe_routers: list[str] = []

    top_corners = [c for c in corners if c in ("NW", "NE")]
    bot_corners = [c for c in corners if c in ("SW", "SE")]

    for corner in corners:
        is_top = corner in ("NW", "NE")

        if is_top:
            corner_idx = top_corners.index(corner)
            row = corner_idx if corner_idx < rows_per_half else rows_per_half - 1
        else:
            corner_idx = bot_corners.index(corner)
            bot_start = hbm_row_end + 1
            row = bot_start + corner_idx if (bot_start + corner_idx) < n_rows else n_rows - 1

        for ci in range(pe_per_corner):
            pe_x, _pe_y = pe_positions[corner][ci]
            col = min(range(n_cols), key=lambda c: abs(col_xs[c] - pe_x))

            key = f"r{row}c{col}"
            router = routers[key]
            if router is not None:
                router["attach"].append(f"pe{pe_idx}.dma")
                router["attach"].append(f"pe{pe_idx}.cpu")
                router["attach"].append(f"pe{pe_idx}.hbm")
                if is_top:
                    top_pe_routers.append(key)
                else:
                    bot_pe_routers.append(key)

            pe_idx += 1

    # M_CPU and SRAM attachments: find nearest router to configured position
    placement = cube_spec.get("placement", {})

    def _nearest_router(target_mm: list[float]) -> str | None:
        best_key, best_dist = None, float("inf")
        for rk, rv in routers.items():
            if rv is None:
                continue
            rx, ry = rv["pos_mm"]
            dist = math.sqrt((rx - target_mm[0]) ** 2 + (ry - target_mm[1]) ** 2)
            if dist < best_dist:
                best_dist = dist
                best_key = rk
        return best_key

    mcpu_pos = placement.get("m_cpu", {}).get("pos_mm", [1.5, 5.5])
    mcpu_key = _nearest_router(mcpu_pos)
    if mcpu_key and routers.get(mcpu_key) is not None:
        routers[mcpu_key]["attach"].append("m_cpu")

    sram_pos = placement.get("sram", {}).get("pos_mm", [1.5, 8.5])
    sram_key = _nearest_router(sram_pos)
    if sram_key and routers.get(sram_key) is not None:
        routers[sram_key]["attach"].append("sram")

    # UCIe PE rows: top-half rows + bottom-half rows (1 per PE row)
    ucie_pe_rows = []
    for r in range(rows_per_half):
        ucie_pe_rows.append(r)
    for r in range(rows_per_half):
        ucie_pe_rows.append(hbm_row_end + 1 + r)

    # UCIe-E distribution: 1 per PE row, rightmost column
    for i, row in enumerate(ucie_pe_rows):
        key = f"r{row}c{n_cols - 1}"
        router = routers.get(key)
        if router is not None:
            router["attach"].append(f"ucie_e.c{i}")

    # UCIe-W distribution: 1 per PE row, leftmost column (mirror of E)
    for i, row in enumerate(ucie_pe_rows):
        key = f"r{row}c0"
        router = routers.get(key)
        if router is not None:
            router["attach"].append(f"ucie_w.c{i}")

    # UCIe PE columns: left-half + right-half PE columns (for N/S distribution)
    pe_xs = set()
    for positions in pe_positions.values():
        for x, _y in positions:
            pe_xs.add(x)
    left_pe_cols = sorted(c for c in range(n_cols)
                          if col_xs[c] in pe_xs and c < hbm_col_start)
    right_pe_cols = sorted(c for c in range(n_cols)
                           if col_xs[c] in pe_xs and c > hbm_col_end)
    n_ucie = len(ucie_pe_rows)
    half_n = n_ucie // 2
    ucie_pe_cols = left_pe_cols[:half_n] + right_pe_cols[:n_ucie - half_n]

    # UCIe-N distribution: PE columns on top row (row 0)
    for i, col in enumerate(ucie_pe_cols):
        key = f"r0c{col}"
        router = routers.get(key)
        if router is not None:
            router["attach"].append(f"ucie_n.c{i}")

    # UCIe-S distribution: PE columns on bottom row (row n_rows-1)
    for i, col in enumerate(ucie_pe_cols):
        key = f"r{n_rows - 1}c{col}"
        router = routers.get(key)
        if router is not None:
            router["attach"].append(f"ucie_s.c{i}")

    return {
        "source_hash": source_hash,
        "mesh": {
            "rows": n_rows,
            "cols": n_cols,
        },
        "routers": routers,
    }
