from gurobipy import *
import numpy as np
import random


def probs(Tree):
    """Funzione che, a partire da un albero, restituisce il vettore,
    di lunghezza pari al numero dei nodi, delle probabilità di ogni nodo"""
    p=[]
    nodes=Tree.nodes()
    for i in range(Tree.number_of_nodes()):
        p.append(nodes[i]['prob'])
    return p

def time_period(Tree):
    """Funzione che, a partire da un albero, restituisce il vettore,
    di lunghezza pari al numero dei nodi, dell'istante di tempo di ogni nodo"""
    T=[]
    nodes=Tree.nodes()
    for i in range(Tree.number_of_nodes()):
        T.append(nodes[i]['t'])
    return T

def nodes_leaves(Tree):
    """Funzione che, a partire da un albero, restituisce il vettore degli
    indici dei nodi e quello degli indici delle foglie"""
    N=Tree.nodes()
    L=Tree.get_leaves()
    return N,L

def predecessor(Tree):
    """Funzione che, a partire da un albero, restituisce il vettore,
    di lunghezza pari al numero dei nodi, del predecessore di ogni nodo"""
    a=[]
    a.append(None)
    for i in range(1,Tree.number_of_nodes()):
        for predecessor in Tree.predecessors(i):
            a.append(predecessor)
    return a

def ancestor(Tree):
    """Funzione che, a partire da un albero, restituisce la matrice,
    di dimensioni pari al numero dei nodi ed ai periodi di tempo, che 
    ha in posizione (i,j) l'indice dell'antenato del nodo i al tempo j,
    quando questoha senso (ovvero se l'istante di tempo j è precedente
    all'istante di tempo del nodo i)"""
    T=time_period(Tree)
    time_periods=max(T)+1
    omega=np.ndarray(shape=(Tree.number_of_nodes(),time_periods),dtype=int)
    omega[0,]=0
    omega[:,0]=0
    for n in range(1,Tree.number_of_nodes()):
        for t in range(1,time_periods):
            if t<T[n]:
                ancestor=n
                for i in range(t,T[n]):
                    for predecessor in Tree.predecessors(ancestor):
                        ancestor=predecessor
                omega[n,t]=ancestor
            else:
                omega[n,t]=0
    return omega

def successors_t(Tree):
    """Funzione che, a partire da un albero, restituisce il dizionario 
    che per ogni coppia (n,t) contiene l'elenco dei successori del nodo n
    all'istante di tempo t, quando t è maggiore del periodo di tempo 
    del nodo n"""
    T=time_period(Tree)
    time_periods=max(T)+1
    sigma={}
    for n in range(0,Tree.number_of_nodes()):
        for t in range(1,time_periods):
            if t>T[n]:
                sigma[(n,t)]=[]
                succ=[]
                succ.append(n)
                for i in range(T[n],t):
                    for j in range(len(succ)):
                        for successor in Tree.successors(succ[j]):
                            if T[successor]==t:
                                sigma[(n,t)].append(successor)
                            succ.append(successor)
    return sigma