from pydynamo.world3 import var_color
# AFFICHAGE D'UNE BOUCLE DE RETROACTION
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
import re

def afficher_boucle_retroaction_industrie():
    plt.rcParams['mathtext.fontset'] = 'custom'
    plt.rcParams['mathtext.it'] = 'STIXGeneral:italic'
    plt.rcParams['mathtext.bf'] = 'STIXGeneral:bold'

    # Coordonnées des noeuds de la figure
    x = 1/np.sqrt(3)
    d = (0, 0.5)
    pos = {
        'ic': np.array((0, 1)),
        'io': np.array((-x, 0)),
        'icir': np.array((x, 0))
    }
    vv = set(pos)
    # On récupère le graphique d'influence des noeuds
    G = w.get_influence_graph()

    # On détermine les positions fixes et non fixes des noeuds
    fixed = set(pos.keys())
    non_fixed = {a for a, b in G.in_edges(pos)}.difference(fixed)
    ioloop = nx.DiGraph(nx.subgraph(G, fixed.union(non_fixed)))
    for f in fixed:
        aa = {a for a, b in ioloop.in_edges(f) if a not in fixed}
        for a, posa in nx.spring_layout(nx.subgraph(ioloop, aa), k=0.2).items():
            pos[a] = pos[f] + posa
    pos = nx.spring_layout(ioloop, k=0.2, pos=pos, fixed=fixed)

    # On dessine le graphe
    def change_eq(eq):
        return re.sub('(io|ic|icir)(?!\w)', r'$\\mathbf{\1}$', eq.replace('/', '\n/').replace('.k', '').replace('.j', '').replace('=', '=\n'))

    nx.draw(ioloop, pos = pos, with_labels=True, node_color= ['tab:blue' if i not in vv else 'tab:red' for i in pos], node_size=1000, width=3, edge_color=['gray' if not (a in vv and b in vv) else 'tab:orange' for a, b in ioloop.edges])
    nx.draw_networkx_edge_labels(ioloop, pos, edge_labels={(a, b): change_eq(w.eqs['update'][b]['raw_line']) for a,b in nx.subgraph(ioloop, fixed).edges() if a!=b}, rotate=False, font_size=14)
    # plt.title('Boucle de rétroaction positive pour la production industrielle');

# ECHELLES ET COULEURS DE MON MODELE
def get_ant_mod_scales_and_color(antmod):
    return {
        'A': max(antmod.A),
        'S': antmod.si,
        'R': antmod.R[0],
        'H': antmod.HB,
        'X': antmod.X[0]*400,
        'V': 1,

    #     'NT': max(s.NT),
    #     'CT': s.si/100,
    #     'ET': s.ri/100,
    #     'YT': s.vi*s.HB*s.MI,
    #     'FT': s.HB*s.V/10,
    #     'XT': s.X[0]*400/30,

    #     'FP': s.HB/100,
    #     'LP': s.HB/100,

        'I': antmod.MI,
    #     'VV': 1,
    }, {
        'CT': 'darkolivegreen',
         'A': 'r',
         'S': 'g',
         'V': 'lime',
         'NT': 'black',
         'H': 'brown',
         'FP': 'peru',
         'LP': 'sandybrown',
         'ET': 'magenta',
         'YT': 'yellow',
         'R': 'gray',
         'I': 'orange',
         'XT': 'slateblue',
         'FT': 'darkorange',
         'VV': 'aquamarine',
         'X': 'purple'
    }