"""
transpiler/structural_analysis.py — Advanced CFG Analysis for Restructuring.
"""

from __future__ import annotations
from typing import Dict, Set, List, Optional
from transpiler.middleend import cfg, ir
from transpiler.optimization import loop_detector
from transpiler.infrastructure import visitor, logger

class PostDominatorAnalysis:
    """Computes immediate post-dominators using the Cooper et al. RPO-intersect
    algorithm on the *reverse CFG*.

    The reverse CFG has the same nodes but all edges flipped — dominators on
    the reverse CFG are exactly post-dominators on the original CFG.  The
    intersect algorithm converges in O(N) passes for reducible graphs and
    remains correct for irreducible ones.
    """

    def __init__(self, target_cfg: cfg.CFG):
        self.cfg = target_cfg
        self.idoms: Dict[int, Optional[int]] = {}  # block_id -> ipostdom_id
        # post_dominators retained for any callers that use it
        self.post_dominators: Dict[int, Set[int]] = {}

    def run(self):
        self._compute_idoms_intersect()
        # Build post_dominators set from the idom tree (walk up from each node)
        self._build_post_dominators_from_idoms()

    # ------------------------------------------------------------------
    # Core algorithm: Cooper et al. intersect on reverse CFG
    # ------------------------------------------------------------------

    def _compute_idoms_intersect(self):
        blocks = self.cfg.blocks

        # Identify exit nodes (sources in the reverse CFG / roots of the post-dom tree)
        exit_ids: Set[int] = {b.id for b in self.cfg.exit_blocks}
        if not exit_ids:
            exit_ids = {bid for bid, b in blocks.items() if not blocks[bid].successors}

        # --- RPO on the reverse CFG (DFS from exit nodes, following predecessors) ---
        rpo: List[int] = []
        visited: Set[int] = set()

        def dfs(node: int):
            visited.add(node)
            # In the reverse CFG, edges go FROM successor TO predecessor of the original.
            # So "successors in reverse" = predecessors in original.
            for pred in blocks[node].predecessors:
                if pred.id not in visited:
                    dfs(pred.id)
            rpo.append(node)

        for eid in sorted(exit_ids):  # stable ordering
            if eid not in visited:
                dfs(eid)
        rpo.reverse()  # RPO = reverse of post-order
        rpo_index: Dict[int, int] = {bid: i for i, bid in enumerate(rpo)}

        # --- Initialise idoms ---
        UNDEFINED = -1
        idoms: Dict[int, int] = {bid: UNDEFINED for bid in blocks}
        for eid in exit_ids:
            idoms[eid] = eid  # roots are their own ipostdom

        def intersect(b1: int, b2: int) -> int:
            """Walk both fingers up the dominator tree until they meet."""
            while b1 != b2:
                while rpo_index.get(b1, 0) > rpo_index.get(b2, 0):
                    b1 = idoms[b1]
                    if b1 == UNDEFINED:
                        return b2
                while rpo_index.get(b2, 0) > rpo_index.get(b1, 0):
                    b2 = idoms[b2]
                    if b2 == UNDEFINED:
                        return b1
            return b1

        # --- Fixed-point iteration in RPO order ---
        changed = True
        while changed:
            changed = False
            for bid in rpo:
                if bid in exit_ids:
                    continue
                # In the reverse CFG, predecessors of bid = *successors* of bid in original
                rev_pred_ids = [s.id for s in blocks[bid].successors]

                # Seed with the first already-processed predecessor
                new_idom = UNDEFINED
                for pred_id in rev_pred_ids:
                    if idoms.get(pred_id, UNDEFINED) != UNDEFINED:
                        new_idom = pred_id
                        break

                if new_idom == UNDEFINED:
                    continue

                for pred_id in rev_pred_ids:
                    if pred_id == new_idom:
                        continue
                    if idoms.get(pred_id, UNDEFINED) != UNDEFINED:
                        new_idom = intersect(new_idom, pred_id)

                if idoms.get(bid, UNDEFINED) != new_idom:
                    idoms[bid] = new_idom
                    changed = True

        # Store: root nodes (idoms[n]==n) have no proper ipostdom
        for bid in blocks:
            val = idoms.get(bid, UNDEFINED)
            self.idoms[bid] = None if (val == UNDEFINED or val == bid) else val

    # ------------------------------------------------------------------
    # Derived: post_dominators set (for callers that need the full set)
    # ------------------------------------------------------------------

    def _build_post_dominators_from_idoms(self):
        """Walk the idom tree upwards from each node to assemble the full
        post_dominators set (block -> set of its post-dominators)."""
        blocks = self.cfg.blocks
        for bid in blocks:
            pdoms: Set[int] = {bid}
            curr: Optional[int] = self.idoms.get(bid)
            while curr is not None:
                pdoms.add(curr)
                curr = self.idoms.get(curr)
            self.post_dominators[bid] = pdoms


class LoopAnalyzer:
    """Analyzes properties of a loop for restructuring."""
    def __init__(self, target_cfg: cfg.CFG, loop_header: int, loop_blocks: Set[int]):
        self.cfg = target_cfg
        self.header_id = loop_header
        self.blocks = loop_blocks
        self.exits: Set[int] = set()   # blocks outside the loop reached from inside
        self.latches: Set[int] = set() # blocks inside the loop that jump back to header

    def analyze(self):
        for bid in self.blocks:
            block = self.cfg.blocks[bid]
            for succ in block.successors:
                if succ.id not in self.blocks:
                    self.exits.add(succ.id)
                if succ.id == self.header_id:
                    self.latches.add(bid)

    def is_single_exit(self) -> bool:
        return len(self.exits) <= 1
