import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import networkx as nx
import numpy as np
import json
import sys
import random

def generate_tree(current_x, current_y, depth, max_depth, max_nodes, x_range, G, parent=None, node_count_per_depth=None):
    """Generates a tree of nodes with positions adjusted on the x-axis, y-axis, and number of nodes on the z-axis."""
    if node_count_per_depth is None:
        node_count_per_depth = {}

    if depth > max_depth:
        return node_count_per_depth

    if depth not in node_count_per_depth:
        node_count_per_depth[depth] = 0

    num_children = random.randint(1, max_nodes)
    x_positions = [current_x + i * x_range / (num_children + 1) for i in range(num_children)]

    for x in x_positions:
        node_id = len(G.nodes)
        node_count_per_depth[depth] += 1
        prob = random.uniform(0, 1)
        G.add_node(node_id, pos=(x, prob, depth))
        if parent is not None:
            G.add_edge(parent, node_id)
        generate_tree(x, current_y + 1, depth + 1, max_depth, max_nodes, x_range, G, parent=node_id, node_count_per_depth=node_count_per_depth)

    return node_count_per_depth


def build_graph_from_json(json_data, G):
    """Builds a graph from JSON data."""
    data = json.loads(json_data)

    def add_event(parent_id, event_data, depth):
        node_id = len(G.nodes)
        prob = event_data['probability'] / 100.0
        pos = (depth, prob, event_data['event_number'])
        label = event_data['name']
        G.add_node(node_id, pos=pos, label=label)
        if parent_id is not None:
            G.add_edge(parent_id, node_id)

        subevents = event_data.get('subevents', {}).get('event', [])
        if not isinstance(subevents, list):
            subevents = [subevents]

        for subevent in subevents:
            add_event(node_id, subevent, depth + 1)

    root_event = list(data.get('events', {}).values())[0]
    root_id = len(G.nodes)
    G.add_node(root_id, pos=(0, root_event['probability'] / 100.0, root_event['event_number']), label=root_event['name'])
    add_event(None, root_event, 0)


def find_paths(G):
    """Finds paths with highest/lowest probability and longest/shortest durations."""
    best_path, worst_path = None, None
    longest_path, shortest_path = None, None
    best_mean_prob, worst_mean_prob = -1, float('inf')
    max_duration, min_duration = -1, float('inf')

    # Use nx.all_pairs_shortest_path for efficiency
    all_paths_dict = dict(nx.all_pairs_shortest_path(G))

    for source, paths_from_source in all_paths_dict.items():
        for target, path in paths_from_source.items():
            if source != target and all('pos' in G.nodes[node] for node in path):
                probabilities = [G.nodes[node]['pos'][1] for node in path]
                mean_prob = np.mean(probabilities)

                if mean_prob > best_mean_prob:
                    best_mean_prob = mean_prob
                    best_path = path
                if mean_prob < worst_mean_prob:
                    worst_mean_prob = mean_prob
                    worst_path = path

                x_positions = [G.nodes[node]['pos'][0] for node in path]
                duration = max(x_positions) - min(x_positions)

                if duration > max_duration:
                    max_duration = duration
                    longest_path = path
                if duration < min_duration and duration > 0:  # Avoid paths with 0 duration
                    min_duration = duration
                    shortest_path = path

    return best_path, best_mean_prob, worst_path, worst_mean_prob, longest_path, shortest_path

def draw_path_3d(G, path, filename='path_plot_3d.png', highlight_color='blue'):
    """Draws a specific path in 3D."""
    H = G.subgraph(path).copy()
    pos = nx.get_node_attributes(G, 'pos')
    x_vals, y_vals, z_vals = zip(*[pos[node] for node in path])

    fig = plt.figure(figsize=(16, 12))
    ax = fig.add_subplot(111, projection='3d')

    node_colors = ['red' if prob < 0.33 else 'blue' if prob < 0.67 else 'green' for _, prob, _ in [pos[node] for node in path]]
    ax.scatter(x_vals, y_vals, z_vals, c=node_colors, s=700, edgecolors='black', alpha=0.7)

    for edge in H.edges():
        x_start, y_start, z_start = pos[edge[0]]
        x_end, y_end, z_end = pos[edge[1]]
        ax.plot([x_start, x_end], [y_start, y_end], [z_start, z_end], color=highlight_color, lw=2)

    for node, (x, y, z) in pos.items():
        if node in path:
            ax.text(x, y, z, str(node), fontsize=12, color='black')

    ax.set_xlabel('Time (weeks)')
    ax.set_ylabel('Event Probability')
    ax.set_zlabel('Event Number')
    ax.set_title('3D Event Tree - Path')

    plt.savefig(filename, bbox_inches='tight')
    plt.close()


def draw_global_tree_3d(G, filename='global_tree.png'):
    """Draws the entire graph in 3D."""
    pos = nx.get_node_attributes(G, 'pos')
    labels = nx.get_node_attributes(G, 'label')

    if not pos:
        print("Graph is empty. No nodes to visualize.")
        return

    x_vals, y_vals, z_vals = zip(*pos.values())
    fig = plt.figure(figsize=(16, 12))
    ax = fig.add_subplot(111, projection='3d')

    node_colors = ['red' if prob < 0.33 else 'blue' if prob < 0.67 else 'green' for _, prob, _ in pos.values()]
    ax.scatter(x_vals, y_vals, z_vals, c=node_colors, s=700, edgecolors='black', alpha=0.7)

    for edge in G.edges():
        x_start, y_start, z_start = pos[edge[0]]
        x_end, y_end, z_end = pos[edge[1]]
        ax.plot([x_start, x_end], [y_start, y_end], [z_start, z_end], color='gray', lw=2)

    for node, (x, y, z) in pos.items():
        label = labels.get(node, f"{node}")
        ax.text(x, y, z, label, fontsize=12, color='black')

    ax.set_xlabel('Time')
    ax.set_ylabel('Probability')
    ax.set_zlabel('Event Number')
    ax.set_title('3D Event Tree')

    plt.savefig(filename, bbox_inches='tight')
    plt.close()

def main(json_data):
    G = nx.DiGraph()
    build_graph_from_json(json_data, G)

    if mode == 'random':
        generate_tree(0, 0, 0, 5, 3, 10, G)
    elif mode == 'json' and input_file:
        with open(input_file, 'r') as file:
            json_data = file.read()
        build_graph_from_json(json_data, G)
    else:
        print("Invalid mode or input file not provided.")
        return

    draw_global_tree_3d(G, filename='global_tree.png')

    best_path, best_mean_prob, worst_path, worst_mean_prob, longest_path, shortest_path = find_paths(G)

    if best_path:
        print(f"\nPath with the highest average probability: {' -> '.join(map(str, best_path))}")
        print(f"Average probability: {best_mean_prob:.2f}")
    if worst_path:
        print(f"\nPath with the lowest average probability: {' -> '.join(map(str, worst_path))}")
        print(f"Average probability: {worst_mean_prob:.2f}")
    if longest_path:
        print(f"\nPath with the longest duration: {' -> '.join(map(str, longest_path))}")
        print(f"Duration: {max(G.nodes[node]['pos'][0] for node in longest_path) - min(G.nodes[node]['pos'][0] for node in longest_path):.2f}")
    if shortest_path:
        print(f"\nPath with the shortest duration: {' -> '.join(map(str, shortest_path))}")
        print(f"Duration: {max(G.nodes[node]['pos'][0] for node in shortest_path) - min(G.nodes[node]['pos'][0] for node in shortest_path):.2f}")

    draw_global_tree_3d(G, filename='global_tree.png')

    if best_path:
        draw_path_3d(G, best_path, 'best_path.png', 'blue')
    if worst_path:
        draw_path_3d(G, worst_path, 'worst_path.png', 'red')
    if longest_path:
        draw_path_3d(G, longest_path, 'longest_duration_path.png', 'green')
    if shortest_path:
        draw_path_3d(G, shortest_path, 'shortest_duration_path.png', 'purple')

    return 'global_tree.png'


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python script.py <mode> [input_file]")
    else:
        mode = sys.argv[1]
        input_file = sys.argv[2] if len(sys.argv) > 2 else None
        main(mode, input_file)