# import plotly.graph_objects as go
# import textwrap
# import re
# from collections import defaultdict

# def generate_subplot(paraphrased_sentence, scheme_sentences, sampled_sentence, highlight_info):
#     # Combine nodes into one list with appropriate labels
#     nodes = [paraphrased_sentence] + scheme_sentences + sampled_sentence
#     nodes[0] += ' L0'  # Paraphrased sentence is level 0
#     para_len = len(scheme_sentences)
#     for i in range(1, para_len + 1):
#         nodes[i] += ' L1'  # Scheme sentences are level 1
#     for i in range(para_len + 1, len(nodes)):
#         nodes[i] += ' L2'  # Sampled sentences are level 2

#     # Define the highlight_words function
#     def highlight_words(sentence, color_map):
#         for word, color in color_map.items():
#             sentence = re.sub(f"\\b{word}\\b", f"{{{{{word}}}}}", sentence, flags=re.IGNORECASE)
#         return sentence

#     # Clean and wrap nodes, and highlight specified words globally
#     cleaned_nodes = [re.sub(r'\sL[0-9]$', '', node) for node in nodes]
#     global_color_map = dict(highlight_info)
#     highlighted_nodes = [highlight_words(node, global_color_map) for node in cleaned_nodes]
#     wrapped_nodes = ['<br>'.join(textwrap.wrap(node, width=30)) for node in highlighted_nodes]

#     # Function to determine tree levels and create edges dynamically
#     def get_levels_and_edges(nodes):
#         levels = {}
#         edges = []
#         for i, node in enumerate(nodes):
#             level = int(node.split()[-1][1])
#             levels[i] = level

#         # Add edges from L0 to all L1 nodes
#         root_node = next(i for i, level in levels.items() if level == 0)
#         for i, level in levels.items():
#             if level == 1:
#                 edges.append((root_node, i))

#         # Add edges from each L1 node to their corresponding L2 nodes
#         l1_indices = [i for i, level in levels.items() if level == 1]
#         l2_indices = [i for i, level in levels.items() if level == 2]

#         for i, l1_node in enumerate(l1_indices):
#             l2_start = i * 4
#             for j in range(4):
#                 l2_index = l2_start + j
#                 if l2_index < len(l2_indices):
#                     edges.append((l1_node, l2_indices[l2_index]))

#         # Add edges from each L2 node to their corresponding L3 nodes
#         l2_indices = [i for i, level in levels.items() if level == 2]
#         l3_indices = [i for i, level in levels.items() if level == 3]

#         l2_to_l3_map = {l2_node: [] for l2_node in l2_indices}

#         # Map L3 nodes to L2 nodes
#         for l3_node in l3_indices:
#             l2_node = l3_node % len(l2_indices)
#             l2_to_l3_map[l2_indices[l2_node]].append(l3_node)

#         for l2_node, l3_nodes in l2_to_l3_map.items():
#             for l3_node in l3_nodes:
#                 edges.append((l2_node, l3_node))

#         return levels, edges

#     # Get levels and dynamic edges
#     levels, edges = get_levels_and_edges(nodes)
#     max_level = max(levels.values(), default=0)

#     # Calculate positions
#     positions = {}
#     level_heights = defaultdict(int)
#     for node, level in levels.items():
#         level_heights[level] += 1

#     y_offsets = {level: - (height - 1) / 2 for level, height in level_heights.items()}
#     x_gap = 2
#     l1_y_gap = 10
#     l2_y_gap = 6

#     for node, level in levels.items():
#         if level == 1:
#             positions[node] = (-level * x_gap, y_offsets[level] * l1_y_gap)
#         elif level == 2:
#             positions[node] = (-level * x_gap, y_offsets[level] * l2_y_gap)
#         else:
#             positions[node] = (-level * x_gap, y_offsets[level] * l2_y_gap)
#         y_offsets[level] += 1

#     # Function to highlight words in a wrapped node string
#     def color_highlighted_words(node, color_map):
#         parts = re.split(r'(\{\{.*?\}\})', node)
#         colored_parts = []
#         for part in parts:
#             match = re.match(r'\{\{(.*?)\}\}', part)
#             if match:
#                 word = match.group(1)
#                 color = color_map.get(word, 'black')
#                 colored_parts.append(f"<span style='color: {color};'>{word}</span>")
#             else:
#                 colored_parts.append(part)
#         return ''.join(colored_parts)

#     # Create figure
#     fig = go.Figure()

#     # Add nodes to the figure
#     for i, node in enumerate(wrapped_nodes):
#         colored_node = color_highlighted_words(node, global_color_map)
#         x, y = positions[i]
#         fig.add_trace(go.Scatter(
#             x=[-x],  # Reflect the x coordinate
#             y=[y],
#             mode='markers',
#             marker=dict(size=10, color='blue'),
#             hoverinfo='none'
#         ))
#         fig.add_annotation(
#             x=-x,  # Reflect the x coordinate
#             y=y,
#             text=colored_node,
#             showarrow=False,
#             xshift=15,
#             align="center",
#             font=dict(size=8),
#             bordercolor='black',
#             borderwidth=1,
#             borderpad=2,
#             bgcolor='white',
#             width=150
#         )

#     # Add edges to the figure
#     for edge in edges:
#         x0, y0 = positions[edge[0]]
#         x1, y1 = positions[edge[1]]
#         fig.add_trace(go.Scatter(
#             x=[-x0, -x1],  # Reflect the x coordinates
#             y=[y0, y1],
#             mode='lines',
#             line=dict(color='black', width=1)
#         ))

#     fig.update_layout(
#         showlegend=False,
#         margin=dict(t=20, b=20, l=20, r=20),
#         xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
#         yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
#         width=1200,  # Adjusted width to accommodate more levels
#         height=1000   # Adjusted height to accommodate more levels
#     )

#     return fig

import plotly.graph_objects as go
import textwrap
import re
from collections import defaultdict

def generate_subplot(paraphrased_sentence, scheme_sentences, sampled_sentence, highlight_info):
    # Combine nodes into one list with appropriate labels
    nodes = [paraphrased_sentence] + scheme_sentences + sampled_sentence
    nodes[0] += ' L0'  # Paraphrased sentence is level 0
    para_len = len(scheme_sentences)
    for i in range(1, para_len + 1):
        nodes[i] += ' L1'  # Scheme sentences are level 1
    for i in range(para_len + 1, len(nodes)):
        nodes[i] += ' L2'  # Sampled sentences are level 2

    # Define the highlight_words function
    def highlight_words(sentence, color_map):
        for word, color in color_map.items():
            sentence = re.sub(f"\\b{word}\\b", f"{{{{{word}}}}}", sentence, flags=re.IGNORECASE)
        return sentence

    # Clean and wrap nodes, and highlight specified words globally
    cleaned_nodes = [re.sub(r'\sL[0-9]$', '', node) for node in nodes]
    global_color_map = dict(highlight_info)
    highlighted_nodes = [highlight_words(node, global_color_map) for node in cleaned_nodes]
    wrapped_nodes = ['<br>'.join(textwrap.wrap(node, width=30)) for node in highlighted_nodes]

    # Function to determine tree levels and create edges dynamically
    def get_levels_and_edges(nodes):
        levels = {}
        edges = []
        for i, node in enumerate(nodes):
            level = int(node.split()[-1][1])
            levels[i] = level

        # Add edges from L0 to all L1 nodes
        root_node = next(i for i, level in levels.items() if level == 0)
        for i, level in levels.items():
            if level == 1:
                edges.append((root_node, i))

        # Add edges from each L1 node to their corresponding L2 nodes
        l1_indices = [i for i, level in levels.items() if level == 1]
        l2_indices = [i for i, level in levels.items() if level == 2]

        for i, l1_node in enumerate(l1_indices):
            l2_start = i * 4
            for j in range(4):
                l2_index = l2_start + j
                if l2_index < len(l2_indices):
                    edges.append((l1_node, l2_indices[l2_index]))

        # Add edges from each L2 node to their corresponding L3 nodes
        l2_indices = [i for i, level in levels.items() if level == 2]
        l3_indices = [i for i, level in levels.items() if level == 3]

        l2_to_l3_map = {l2_node: [] for l2_node in l2_indices}

        # Map L3 nodes to L2 nodes
        for l3_node in l3_indices:
            l2_node = l3_node % len(l2_indices)
            l2_to_l3_map[l2_indices[l2_node]].append(l3_node)

        for l2_node, l3_nodes in l2_to_l3_map.items():
            for l3_node in l3_nodes:
                edges.append((l2_node, l3_node))

        return levels, edges

    # Get levels and dynamic edges
    levels, edges = get_levels_and_edges(nodes)
    max_level = max(levels.values(), default=0)

    # Calculate positions
    positions = {}
    level_heights = defaultdict(int)
    for node, level in levels.items():
        level_heights[level] += 1

    y_offsets = {level: - (height - 1) / 2 for level, height in level_heights.items()}
    x_gap = 2
    l1_y_gap = 10
    l2_y_gap = 6

    for node, level in levels.items():
        if level == 1:
            positions[node] = (-level * x_gap, y_offsets[level] * l1_y_gap)
        elif level == 2:
            positions[node] = (-level * x_gap, y_offsets[level] * l2_y_gap)
        else:
            positions[node] = (-level * x_gap, y_offsets[level] * l2_y_gap)
        y_offsets[level] += 1

    # Function to highlight words in a wrapped node string
    def color_highlighted_words(node, color_map):
        parts = re.split(r'(\{\{.*?\}\})', node)
        colored_parts = []
        for part in parts:
            match = re.match(r'\{\{(.*?)\}\}', part)
            if match:
                word = match.group(1)
                color = color_map.get(word, 'black')
                colored_parts.append(f"<span style='color: {color};'>{word}</span>")
            else:
                colored_parts.append(part)
        return ''.join(colored_parts)

    # Define the text for each edge
    edge_texts = [
        "Highest Entropy Masking",
        "Pseudo-random Masking",
        "Random Masking",
        "Greedy Sampling",
        "Temperature Sampling",
        "Exponential Minimum Sampling",
        "Inverse Transform Sampling",
        "Greedy Sampling",
        "Temperature Sampling",
        "Exponential Minimum Sampling",
        "Inverse Transform Sampling",
        "Greedy Sampling",
        "Temperature Sampling",
        "Exponential Minimum Sampling",
        "Inverse Transform Sampling"
    ]

    # Create figure
    fig = go.Figure()

    # Add nodes to the figure
    for i, node in enumerate(wrapped_nodes):
        colored_node = color_highlighted_words(node, global_color_map)
        x, y = positions[i]
        fig.add_trace(go.Scatter(
            x=[-x],  # Reflect the x coordinate
            y=[y],
            mode='markers',
            marker=dict(size=10, color='blue'),
            hoverinfo='none'
        ))
        fig.add_annotation(
            x=-x,  # Reflect the x coordinate
            y=y,
            text=colored_node,
            showarrow=False,
            xshift=15,
            align="center",
            font=dict(size=8),
            bordercolor='black',
            borderwidth=1,
            borderpad=2,
            bgcolor='white',
            width=150
        )

    # Add edges and text above each edge
    for i, edge in enumerate(edges):
        x0, y0 = positions[edge[0]]
        x1, y1 = positions[edge[1]]
        fig.add_trace(go.Scatter(
            x=[-x0, -x1],  # Reflect the x coordinates
            y=[y0, y1],
            mode='lines',
            line=dict(color='black', width=1)
        ))

        # Calculate the midpoint of the edge
        mid_x = (-x0 + -x1) / 2
        mid_y = (y0 + y1) / 2

        # Adjust y position to shift text upwards
        text_y_position = mid_y + 0.8  # Increase this value to shift the text further upwards

        # Add text annotation above the edge
        fig.add_annotation(
            x=mid_x,
            y=text_y_position,
            text=edge_texts[i],  # Use the text specific to this edge
            showarrow=False,
            font=dict(size=10),
            align="center"
        )

    fig.update_layout(
        showlegend=False,
        margin=dict(t=20, b=20, l=20, r=20),
        xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
        yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
        width=1200,  # Adjusted width to accommodate more levels
        height=1000   # Adjusted height to accommodate more levels
    )

    return fig