#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# IMPORTS
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

from enum import Enum;

from src.local.typing import *;

from src.stacks.stack import *;
from src.graphs.graph import *

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# EXPORTS
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

__all__ = [
    'tarjan_algorithm',
];

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# CONSTANTS
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

class State(Enum):
    UNTOUCHED = 0;
    PENDING = 1;
    FINISHED = 2;

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# Tarjan Algorithm
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

class NodeInformation:
    root: int;
    index: int;
    state: State;

    def __init__(self):
        self.root = self.index = 0;
        self.state = State.UNTOUCHED;

def tarjan_algorithm(G: Graph) -> List[Any]:
    '''
    runs the Tarjan-Algorithm to compute the strongly connected components
    '''

    # sonst Stack erstellen und Bearbeitungszustand der Knoten im Speicher notieren:
    S = Stack();
    index = 0;
    infos = { u: NodeInformation() for u in G.nodes };
    components = [];

    def tarjan_visit(v: Any):
        '''
        recursive depth-first search algorithm to compute components
        '''
        nonlocal G, S, index, components;

        if not(infos[v].state == State.UNTOUCHED):
            return;

        index += 1;
        infos[v].index = infos[v].root = index;
        S.push(v);

        infos[v].state = State.PENDING;
        # depth first search:
        for u in G.successor(v):
            tarjan_visit(u);
            # remains relevant for v, provided u still in Stack:
            if u in S:
                infos[v].root = min(infos[v].root, infos[u].root);
        infos[v].state = State.FINISHED;

        # if at root of component pop everything from stack up to root and add component:
        if infos[v].index == infos[v].root:
            component = [];
            condition_reached_root = False;
            while not condition_reached_root:
                u = S.pop();
                component.append(u);
                condition_reached_root = (u == v);
            components.append(component);
        return;

    for v in G.nodes:
        tarjan_visit(v);
    return components;