import random
import networkx as nx
from IPython.core.display import SVG
import pyomo.environ as pyo
from pysat.solvers import Solver
from pysat.formula import CNF 
import py_svg_combinatorics as psc
from ipywidgets import widgets, HBox
from collections import Counter
from pprint import pprint
from random import randint
import numpy as np
from IPython.display import IFrame
import IPython
from copy import copy
import os
from pathlib import Path

def print_solution(m):
    for v in m.component_data_objects(pyo.Var):
        if v.value and v.value > 0:
            print(str(v), v.value)

nbname = ''
try:
    nbname = __vsc_ipynb_file__
except:
    if 'COCALC_JUPYTER_FILENAME' in os.environ:
        nbname = os.environ['COCALC_JUPYTER_FILENAME']    
title_ = Path(nbname).stem.replace('-', '_').title()    
IFrame(f'https://discopal.ispras.ru/index.php?title=Hardprob/{title_}&useskin=cleanmonobook', width=1280, height=300)

def visme(G, m=None):
    def norm_size(sz):
        return int(sz*100)

    pos = nx.get_node_attributes(G, "pos")
    if not pos:
        pos = nx.spring_layout(G)
        nx.set_node_attributes(G, pos, "pos")
    cut_edges = None 
    v_colors = ['blue']*len(G.nodes())
    if m:
        p_colors_ = psc.get_hex_palette_for_n(m.k)
        cut_edges = [e for e in m.E if pyo.value(m.y[e])>0]
        for v in G.nodes():
            v_colors[v] = p_colors_[int(pyo.value(m.x[v]))]

    nx.draw_networkx(
        G,
        pos,
        with_labels=True,
        edge_color="green",
        node_color=v_colors,
        node_size=300,
        style=':',
        width=0.4,
        font_size=6
    )

    if cut_edges:
        nx.draw_networkx_edges(
            G,
            pos,
            edgelist=cut_edges,
            edge_color="blue",
            node_size=200,
            width=0.7,
        )

# G = nx.fast_gnp_random_graph(10, 0.5)
G = nx.random_lobster(10, 0.5, 0.5)
for (u, v) in G.edges():
    G.edges[u, v]['weight'] = np.random.rand()

visme(G)

def get_model(G, k):
    m = pyo.ConcreteModel()
    m.k = k
    m.E = list(G.edges()) 
    m.V = list(G.nodes())
    m.W = dict([(e[:2], e[2]['weight']) for e in G.edges(data=True)])
    m.m = len(m.E)
    m.n = len(m.V)
    bigM = 10000
    eps = 0.1

    # Номер разбиения
    m.x = pyo.Var(m.V, domain=pyo.NonNegativeIntegers)

    @m.Constraint(m.V)
    def каждая_вершина_в_каком_нибудь_разбиении_не_большем_k(m, v):
        return m.x[v] <= m.k - 1


    # 0-1 — на разрезе или нет.
    m.y = pyo.Var(m.E, domain=pyo.Binary)

    # техническая переменная «v<u»
    m.u_gt_v = pyo.Var(m.E, domain=pyo.Binary)
    # техническая переменная «v>u»
    m.v_gt_u = pyo.Var(m.E, domain=pyo.Binary)


    @m.Constraint(m.E)
    def выставь_u_gt_v_u(m, u, v):
        return m.x[u] - m.x[v] <= k * m.u_gt_v[(u, v)]

    @m.Constraint(m.E)
    def выставь_u_gt_v_d(m, u, v):
        return m.x[u] - m.x[v] >= eps -k * (1 - m.u_gt_v[(u, v)])


    @m.Constraint(m.E, ['u', 'd'])
    def выставь_v_gt_u(m, u, v, what):
        if what == 'u':
            return m.x[v] - m.x[u] <= k * m.v_gt_u[(u, v)]
        if what == 'd':
            return m.x[v] - m.x[u] >= eps -k * (1 - m.v_gt_u[(u, v)])
        return pyo.Constraint.Skip

    @m.Constraint(m.E)
    def ребро_на_разрезе_только_когда_вершины_в_разных_разбиениях(m, u, v):
        return m.y[(u,v)] == m.v_gt_u[(u, v)] + m.u_gt_v[(u, v)]

    m.cut_weight = pyo.Objective(expr = sum(m.W[e] * m.y[e] for e in m.E), sense=pyo.maximize)
    return m

m = get_model(G, 5)
#m.E, m.V

solver = pyo.SolverFactory('cbc')
solver.solve(m).write()
#print_solution(m)

visme(G, m)