import networkx as nx
import matplotlib.pyplot as plt
import matplotlib.patches as patches

# Create a directed graph
G = nx.DiGraph()

# Add nodes for personality traits
G.add_node("Extraversion")
G.add_node("Agreeableness")
G.add_node("Conscientiousness")
G.add_node("Neuroticism")
G.add_node("Openness")

# Add nodes for motivational factors
G.add_node("Information Seeking & Sharing")
G.add_node("Self-Expression & Exhibition")
G.add_node("Social Interaction")
G.add_node("Escapism & Relaxation")
G.add_node("Norm & Trend Following")

# Add the outcome variable
G.add_node("Time Spent on Instagram")

# Add edges (relationships)
G.add_edges_from([
    ("Extraversion", "Information Seeking & Sharing"),
    ("Extraversion", "Self-Expression & Exhibition"),
    ("Extraversion", "Social Interaction"),
    ("Extraversion", "Escapism & Relaxation"),
    ("Extraversion", "Norm & Trend Following"),
    ("Extraversion", "Time Spent on Instagram"),
    ("Agreeableness", "Information Seeking & Sharing"),
    ("Agreeableness", "Self-Expression & Exhibition"),
    ("Agreeableness", "Social Interaction"),
    ("Agreeableness", "Escapism & Relaxation"),
    ("Agreeableness", "Norm & Trend Following"),
    ("Agreeableness", "Time Spent on Instagram"),
    ("Conscientiousness", "Information Seeking & Sharing"),
    ("Conscientiousness", "Self-Expression & Exhibition"),
    ("Conscientiousness", "Social Interaction"),
    ("Conscientiousness", "Escapism & Relaxation"),
    ("Conscientiousness", "Norm & Trend Following"),
    ("Conscientiousness", "Time Spent on Instagram"),
    ("Neuroticism", "Information Seeking & Sharing"),
    ("Neuroticism", "Self-Expression & Exhibition"),
    ("Neuroticism", "Social Interaction"),
    ("Neuroticism", "Escapism & Relaxation"),
    ("Neuroticism", "Norm & Trend Following"),
    ("Neuroticism", "Time Spent on Instagram"),
    ("Openness", "Information Seeking & Sharing"),
    ("Openness", "Self-Expression & Exhibition"),
    ("Openness", "Social Interaction"),
    ("Openness", "Escapism & Relaxation"),
    ("Openness", "Norm & Trend Following"),
    ("Openness", "Time Spent on Instagram"),
    ("Information Seeking & Sharing", "Time Spent on Instagram"),
    ("Self-Expression & Exhibition", "Time Spent on Instagram"),
    ("Social Interaction", "Time Spent on Instagram"),
    ("Escapism & Relaxation", "Time Spent on Instagram"),
    ("Norm & Trend Following", "Time Spent on Instagram")
])

# Define positions for the nodes with offsets for edge placement
pos_ordered = {
    "Extraversion": (-4, 3),
    "Agreeableness": (-4, 1.5),
    "Conscientiousness": (-4, 0),
    "Neuroticism": (-4, -1.5),
    "Openness": (-4, -3),
    "Information Seeking & Sharing": (1, 2),
    "Self-Expression & Exhibition": (1, 1),
    "Social Interaction": (1, 0),
    "Escapism & Relaxation": (1, -1),
    "Norm & Trend Following": (1, -2),
    "Time Spent on Instagram": (4, 0)  # Moved further to the right
}

# Define node width and height
node_width = 2.7
node_height = 0.35

# Increase the canvas size and draw the graph
plt.figure(figsize=(16, 12))

# Draw edges with offsets
for start, end in G.edges():
    start_pos = pos_ordered[start]
    end_pos = pos_ordered[end]
    if start_pos[0] < end_pos[0]:  # From left to right
        offset = (0.5 * node_width, 0)
    else:  # From right to left
        offset = (-0.5 * node_width, 0)
    adjusted_start_pos = (start_pos[0] + offset[0], start_pos[1] + offset[1])
    adjusted_end_pos = (end_pos[0] - offset[0], end_pos[1] - offset[1])
    plt.plot([adjusted_start_pos[0], adjusted_end_pos[0]], 
             [adjusted_start_pos[1], adjusted_end_pos[1]], 
             color='gray', linewidth=1, zorder=1)
    plt.arrow(adjusted_start_pos[0], adjusted_start_pos[1], 
              adjusted_end_pos[0] - adjusted_start_pos[0], 
              adjusted_end_pos[1] - adjusted_start_pos[1], 
              head_width=0.1, head_length=0.1, fc='gray', ec='gray', zorder=2)

# Draw the nodes as rectangles
ax = plt.gca()
for node, (x, y) in pos_ordered.items():
    rect = patches.FancyBboxPatch((x - 1.35, y - 0.175), node_width, node_height, 
                                  boxstyle="round,pad=0.05", 
                                  edgecolor='black', facecolor='lightgreen', 
                                  linewidth=1, zorder=3)
    ax.add_patch(rect)
    ax.text(x, y, node, ha='center', va='center', fontsize=15, fontweight='bold', zorder=4)

# Set limits and title
plt.xlim(-6, 6)
plt.ylim(-5, 5)
plt.title("Theoretical Framework: Personality Traits, Motivational Factors, and Time Spent on Instagram")

# Display the plot
plt.show()