import json
import random
import networkx as nx
import os
from visualizer import get_sorted_nodes

def prepare_edges_for_json(G):
    nodes_list = get_sorted_nodes(G)
    nodes_list_dict = {str(i+1): node for i, node in enumerate(nodes_list)}
    coord_to_id = {v: k for k, v in nodes_list_dict.items()}
    edges_formatted = []
    for u, v in G.edges():
        if u in coord_to_id and v in coord_to_id:
            edges_formatted.append({"room1": coord_to_id[u], "room2": coord_to_id[v]})
    return edges_formatted, list(nodes_list_dict.keys()), nodes_list_dict

def prepare_parameter_for_json(G, I, nodes_list_dict):
    n_count = len(G.nodes())
    if n_count == 0:
        return [], [], [], [], [], [], [], [], [], []
        
    weights = [n_count / (n_count * (1 + (((i + 1) * 2) / 30))) for i in range(n_count)]
    m_weights = random.choices(I, weights=weights, k=5)
    t_weights_probs = [n_count / (n_count * (1 + (((i + 1) * 2) / 5))) for i in range(10)]
    t_weights = random.choices(range(1, 11), weights=t_weights_probs, k=5)

    dismantled, conditioningDuration, assignment, help_list = [], [], [], []

    for m in range(5):
        dismantled.append({"m": str(m + 1), "i": str(m_weights[m]), "t": t_weights[m], "value": 1})
        conditioningDuration.append({"m": str(m + 1), "value": 1})
        x = random.randint(1, 3)
        if m > 2:
            if 1 not in help_list: x = 1
            if 2 not in help_list: x = 2
            if 3 not in help_list: x = 3
        help_list.append(x)
        assignment.append({"m": str(m + 1), "r": str(x), "value": 1})

    t_weights_del = random.choices(range(1, 11), weights=t_weights_probs[:10], k=3)
    delivered = [{"r": str(r+1), "i": "1", "t": t_weights_del[r], "value": 1} for r in range(3)]
    conditioningCapacity = [{"r": str(r+1), "value": 1} for r in range(3)]

    CostMT, CostMB, CostRT, CostRB, Coord = [], [], [], [], []
    for i in range(n_count):
        s_id = str(i + 1)
        CostMT.append({"i": s_id, "value": random.choice([2, 5])})
        CostMB.append({"i": s_id, "value": random.choice([5, 10, 30])})
        CostRT.append({"i": s_id, "value": random.choice([4, 10])})
        CostRB.append({"i": s_id, "value": 1000 if i==0 else random.choice([20, 30, 100])})
        if s_id in nodes_list_dict:
            Coord.append({"i": s_id, "Coordinates": nodes_list_dict[s_id]})

    return dismantled, assignment, delivered, conditioningCapacity, conditioningDuration, CostMT, CostMB, CostRT, CostRB, Coord

def generate_full_json_dict(G, loop=0):
    edges, I, nodes_list_dict = prepare_edges_for_json(G)
    dismantled, assignment, delivered, condCap, condDur, CostMT, CostMB, CostRT, CostRB, Coord = prepare_parameter_for_json(G, I, nodes_list_dict)
    sets = {
        "I": I, "E": {"bidirectional": True, "seed": 1, "edges": edges},
        "M": ["1", "2", "3", "4", "5"], "R": ["1", "2", "3"]
    }
    params = {
        "defaults": { "V": 1000, "CostMB": 100, "CostMT": 20, "CostRB": 300, "CostRT": 50 },
        "t_max": 100, "V": [{"m": "1", "i": "1", "value": 42}],
        "dismantled": dismantled, "delivered": delivered,
        "conditioningCapacity": condCap, "conditioningDuration": condDur,
        "assignment": assignment, "Coord": Coord,
        "CostMT": CostMT, "CostMB": CostMB, "CostRT": CostRT, "CostRB": CostRB, "CostZR": 9, "CostZH": 5
    }
    return {"description": "Generated by Gradio", "sets": sets, "params": params}

def load_graph_from_data(data, name):
    """Core function to parse loaded JSON data into a NetworkX graph."""
    G = nx.Graph()
    id_to_coord = {}

    if "params" in data and "Coord" in data["params"]:
        for item in data["params"]["Coord"]:
            coord = tuple(item["Coordinates"])
            id_to_coord[item["i"]] = coord
            G.add_node(coord)

    if "sets" in data and "E" in data["sets"] and "edges" in data["sets"]["E"]:
        for edge in data["sets"]["E"]["edges"]:
            r1 = edge["room1"]
            r2 = edge["room2"]
            if r1 in id_to_coord and r2 in id_to_coord:
                G.add_edge(id_to_coord[r1], id_to_coord[r2])

    width = max([n[0] for n in G.nodes()]) if len(G.nodes()) > 0 else 10
    height = max([n[1] for n in G.nodes()]) if len(G.nodes()) > 0 else 10
    width = int(width + max(2, width * 0.1))
    height = int(height + max(2, height * 0.1))

    return {"name": name, "graph": G, "width": width, "height": height}

def load_graph_from_json(filepath):
    with open(filepath, 'r') as f:
        data = json.load(f)
    return load_graph_from_data(data, os.path.basename(filepath))