"""
Example 2: Pathfinding in a 2D Grid
Solve grid-based pathfinding using the DPO (Debt-Payment Optimization) algorithm.

DPO key idea: intentionally accepts worse moves and records the degradation as
'debt'. When an improvement is found it repays that debt with interest, forcing
the search to overshoot local optima and converge aggressively to the global best.

  X_new = X + β·debt (repayment) + γ·debt (overshoot) + δ·(gBest - X) (global pull)
"""

import numpy as np
from dpo.core.problem import Problem
from dpo.core.solution import CombinatoricSolution
from dpo.core.universal import DPO_Presets, DPO_Universal


# ---------------------------------------------------------------------------
# Grid helpers
# ---------------------------------------------------------------------------

def create_grid_with_obstacles(size: int, num_obstacles: int = 5) -> np.ndarray:
    """Return a (size × size) binary grid; 1 = obstacle, 0 = free."""
    grid = np.zeros((size, size), dtype=int)
    for _ in range(num_obstacles):
        x = np.random.randint(2, size - 2)
        y = np.random.randint(2, size - 2)
        w = np.random.randint(2, 5)
        h = np.random.randint(2, 5)
        grid[x : min(x + w, size), y : min(y + h, size)] = 1
    # Keep start and goal free
    grid[0, 0] = 0
    grid[-1, -1] = 0
    return grid


# ---------------------------------------------------------------------------
# Solution representation
# ---------------------------------------------------------------------------

class PathSolution(CombinatoricSolution):
    """
    A candidate path encoded as an ordered list of (row, col) waypoints.
    Start and goal are pinned; intermediate waypoints are evolved by DPO.
    """

    def __init__(self, waypoints: list, grid_size: int, start: tuple, goal: tuple):
        super().__init__(waypoints, grid_size)
        self._grid_size = grid_size
        self._start = start
        self._goal = goal

    # -- DPO mutation operator -----------------------------------------------
    def mutate(self, mutation_strength: float = 0.15) -> "PathSolution":
        mutant_wps = list(self.sequence)  # shallow copy
        num_mutations = max(1, int(len(mutant_wps) * mutation_strength))
        for _ in range(num_mutations):
            if len(mutant_wps) <= 2:
                break
            idx = np.random.randint(1, len(mutant_wps) - 1)
            mutant_wps[idx] = (
                np.random.randint(self._grid_size),
                np.random.randint(self._grid_size),
            )
        return PathSolution(mutant_wps, self._grid_size, self._start, self._goal)

    # -- DPO crossover operator ----------------------------------------------
    def crossover(self, other: "PathSolution") -> "PathSolution":
        n = min(len(self.sequence), len(other.sequence))
        split = np.random.randint(1, max(2, n - 1))
        child_wps = list(self.sequence[:split]) + list(other.sequence[split:])
        # Always keep the correct start / goal
        child_wps[0] = self._start
        child_wps[-1] = self._goal
        return PathSolution(child_wps, self._grid_size, self._start, self._goal)

    def copy(self) -> "PathSolution":
        return PathSolution(
            list(self.sequence), self._grid_size, self._start, self._goal
        )

    def to_dict(self) -> dict:
        return {"waypoints": list(self.sequence), "num_waypoints": len(self.sequence)}


# ---------------------------------------------------------------------------
# Problem definition (DPO Problem ABC)
# ---------------------------------------------------------------------------

class PathfindingProblem(Problem):
    """
    2-D grid pathfinding formulated as a DPO Problem.

    Fitness = path_length + 100 × collisions + 10 × dist_to_goal
    (lower is better; DPO minimises fitness)
    """

    # Required metrics keys expected by DPO's fitness weights
    _MAX_DIAGONAL = None  # set in __init__

    def __init__(self, grid: np.ndarray, start: tuple, goal: tuple):
        self.grid = grid
        self.start = start
        self.goal = goal
        self._grid_size = grid.shape[0]
        # Maximum possible diagonal — used to normalise latency_ms / flops_m
        self._MAX_DIAGONAL = float(np.sqrt(2) * self._grid_size)

    # -- Core evaluation (called by DPO every iteration) --------------------
    def evaluate(self, solution, **kwargs):
        """
        Returns (fitness: float, metrics: dict).
        DPO requires metrics to contain: accuracy, latency_ms, memory_mb, flops_m.
        """
        waypoints = list(solution.sequence)

        total_length = 0.0
        collisions = 0

        for i in range(len(waypoints) - 1):
            r1, c1 = waypoints[i]
            r2, c2 = waypoints[i + 1]
            seg_len = np.sqrt((r2 - r1) ** 2 + (c2 - c1) ** 2)
            total_length += seg_len
            # Penalise any segment ending on an obstacle cell
            if 0 <= r2 < self._grid_size and 0 <= c2 < self._grid_size:
                if self.grid[r2, c2] == 1:
                    collisions += 1

        dist_to_goal = np.sqrt(
            (waypoints[-1][0] - self.goal[0]) ** 2
            + (waypoints[-1][1] - self.goal[1]) ** 2
        )

        fitness = total_length + 100.0 * collisions + 10.0 * dist_to_goal

        # DPO-required metrics -----------------------------------------------
        # accuracy  : higher = better (DPO maximises this internally)
        # latency_ms: mapped to path length (normalised)
        # memory_mb : mapped to waypoint count (light proxy)
        # flops_m   : mapped to evaluations performed (proxy)
        accuracy = float(1.0 / (1.0 + fitness / (self._MAX_DIAGONAL * 10.0)))
        metrics = {
            "accuracy":    accuracy,
            "latency_ms":  total_length,
            "memory_mb":   float(len(waypoints)) * 0.1,
            "flops_m":     float(len(waypoints)) * 1.0,
            # extra diagnostics (ignored by DPO weights, useful for reporting)
            "path_length":     total_length,
            "collisions":      collisions,
            "distance_to_goal": dist_to_goal,
        }
        return fitness, metrics

    # -- Solution factory (DPO calls this to seed the initial population) ---
    def create_solution(self, **kwargs) -> PathSolution:
        num_intermediate = np.random.randint(3, 8)
        waypoints = (
            [self.start]
            + [
                (
                    np.random.randint(self._grid_size),
                    np.random.randint(self._grid_size),
                )
                for _ in range(num_intermediate)
            ]
            + [self.goal]
        )
        return PathSolution(waypoints, self._grid_size, self.start, self.goal)

    # -- Problem metadata (DPO uses 'type' for auto-preset selection) -------
    def get_problem_info(self) -> dict:
        return {
            "name": "GridPathfinding",
            "type": "pathfinding",         # triggers DPO pathfinding preset
            "grid_size": self._grid_size,
        }

    # -- Hard constraints (none here; obstacle penalty handled in fitness) ---
    def get_constraints(self) -> dict:
        return {}


# ---------------------------------------------------------------------------
# Main example
# ---------------------------------------------------------------------------

def example_pathfinding_grid():
    """Run DPO on the 2-D grid pathfinding problem and report results."""
    print("=" * 60)
    print("DPO Pathfinding Example: 2D Grid with Obstacles")
    print("=" * 60)

    # ── Problem setup ────────────────────────────────────────────
    grid_size = 20
    start = (0, 0)
    goal  = (grid_size - 1, grid_size - 1)
    np.random.seed(42)
    grid  = create_grid_with_obstacles(grid_size, num_obstacles=5)

    problem = PathfindingProblem(grid, start, goal)

    # ── DPO configuration (Pathfinding preset) ───────────────────
    # alpha=0.25  — widest exploration of all presets
    # beta=0.90   — moderate repayment force
    # gamma=1.20  — strong overshoot (helps escape local optima)
    # delta=0.15  — light global-best pull (keeps diversity)
    # 5 islands, migration_freq=8  — fast inter-island mixing
    # force_late_debt=True from 60 %  — keeps escaping late stagnation
    config = DPO_Presets.Pathfinding_Config(
        population_size=60,
        max_iterations=100,
    )
    config.verbose = False  # silence per-iteration logs

    # ── Run DPO ─────────────────────────────────────────────────
    optimizer = DPO_Universal(problem=problem, config=config)
    result    = optimizer.optimize()

    # ── Results ─────────────────────────────────────────────────
    best_sol = optimizer.get_best_solution()
    history  = optimizer.get_history()

    print("\n" + "─" * 60)
    print("DPO Pathfinding Results")
    print("─" * 60)
    print(f"  Best fitness       : {result['best_fitness']:.4f}")
    print(f"  Best accuracy      : {result['best_accuracy']:.4f}")

    # Retrieve per-metric values from the history acceptance stats if available
    stats = result.get("acceptance_stats", {})
    print(f"  Total candidates   : {stats.get('total_candidates', 'N/A')}")
    print(f"  Accepted (better)  : {stats.get('accepted_better', 'N/A')}")
    print(f"  Accepted worse     : {stats.get('accepted_worse', 'N/A')}  ← DPO debt events")
    print(f"  Rejected           : {stats.get('rejected', 'N/A')}")

    print("\n  Convergence milestones:")
    print(f"    Iter to 95 % best : {history.get('time_to_95', 'N/A')}")
    print(f"    Iter to 99 % best : {history.get('time_to_99', 'N/A')}")
    print(f"    AUC @ 50 %        : {history.get('auc_50', 0.0):.4f}")

    if best_sol:
        wps = best_sol.get("waypoints", [])
        print(f"\n  Waypoints in best path : {len(wps)}")
        if wps:
            print(f"  Path start → goal      : {wps[0]} → {wps[-1]}")

    return result


if __name__ == "__main__":
    result = example_pathfinding_grid()

    print("\n" + "=" * 60)
    print("Pathfinding optimisation complete!")
    print("=" * 60)
