import json
from collections import defaultdict

import fire
import plotly.graph_objects as go

from deception_detection.utils import REPO_ROOT

BASE_FOLDER = REPO_ROOT / "data" / "insider_trading"


def create_sankey_diagram(model_name: str = "llama-70b-3.3") -> None:
    # Read results
    with open(BASE_FOLDER / f"{model_name}-generations.json") as f:
        results = [r["metadata"] for r in json.load(f)]

    # Initialize link counts
    links_dict = defaultdict(int)

    for result in results:
        # Track trade decision
        if result["made_trade"] == "no":
            links_dict[("all", "no_trade")] += 1
        elif result["made_trade"] == "invalid":
            links_dict[("all", "invalid_trade")] += 1
        else:
            # Track trade classification
            if result["trade_label"] in ["no", "ambiguous"]:
                links_dict[("all", "no_insider")] += 1
            elif result["trade_label"] == "yes":
                links_dict[("all", "insider_trade")] += 1

                # Track report type for insider trading cases
                if result.get("report_label") == "explicit":
                    links_dict[("insider_trade", "report_explicit")] += 1
                elif result.get("report_label") == "implied":
                    links_dict[("insider_trade", "report_implied")] += 1
                elif result.get("report_label") == "concealed":
                    links_dict[("insider_trade", "report_concealed")] += 1
                else:
                    links_dict[("insider_trade", "report_invalid")] += 1

                # Track doubling down responses for concealed/deceptive cases
                if result.get("report_label") in ["implied", "concealed"]:
                    source_node = (
                        "report_implied"
                        if result["report_label"] == "implied"
                        else "report_concealed"
                    )
                    if result.get("doubling_down_label") == "honest":
                        links_dict[(source_node, "confirmation_honest")] += 1
                    elif result.get("doubling_down_label") == "deceptive":
                        links_dict[(source_node, "confirmation_deceptive")] += 1
                    elif result.get("doubling_down_label") == "ambiguous":
                        links_dict[(source_node, "confirmation_ambiguous")] += 1
                    elif result.get("doubling_down_label") == "invalid":
                        links_dict[(source_node, "confirmation_invalid")] += 1

    # Calculate total counts for each node
    node_counts = defaultdict(int)

    for (source_node, target_node), count in links_dict.items():
        if source_node == "all":
            node_counts["all"] += count
        node_counts[target_node] += count

    node_info = {
        "all": ("All transcripts", "grey", 0.5),
        "no_trade": ("Model doesn't make trade", "grey", 0.2),
        "insider_trade": ("Model makes insider trade", "grey", 0.6),
        "invalid_trade": ("Model makes trade, invalid parsing", "grey", 0.87),
        "no_insider": ("Model makes trade without using insider info", "grey", 0.75),
        "report_explicit": ("Report Honest", "blue", 0.48),
        "report_implied": ("Report Ambiguous", "yellow", 0.59),
        "report_concealed": ("Report Deceptive", "red", 0.68),
        "report_invalid": ("Report Invalid", "grey", 0.75),
        "confirmation_honest": ("Confirmation Honest", "blue", 0.5),
        "confirmation_deceptive": ("Confirmation Deceptive", "red", 0.6),
        "confirmation_ambiguous": ("Confirmation Ambiguous", "yellow", 0.53),
        "confirmation_invalid": ("Confirmation Invalid", "grey", 0.68),
    }

    # Create node labels with counts and percentages
    node_labels = [f"{node_info[node][0]}\n({node_counts[node]})" for node in node_info]

    # Convert to lists for plotly
    source: list[int] = []
    target: list[int] = []
    value: list[int] = []

    # Create node index mapping
    node_indices = {node: idx for idx, node in enumerate(node_info.keys())}

    # Convert links to indices
    for (source_node, target_node), count in links_dict.items():
        source.append(node_indices[source_node])
        target.append(node_indices[target_node])
        value.append(count)

    # Create and save the diagram

    x: list[float] = [0, 0.25, 0.55, 0.85]
    x_by_node = [x[0]] + [x[1]] * 4 + [x[2]] * 4 + [x[3]] * 4

    fig = go.Figure(
        data=[
            go.Sankey(
                node=dict(
                    pad=15,
                    thickness=20,
                    line=dict(color="black", width=0.5),
                    label=node_labels,  # Use the new labels with counts
                    color=[node_info[node][1] for node in node_info],
                    x=x_by_node,
                    y=[node_info[node][2] for node in node_info],
                ),
                link=dict(source=source, target=target, value=value),
            )
        ]
    )

    fig.update_layout(
        title_text="Insider Trading Decision Flow",
        font_size=10,
        # Make the plot wider to better show the levels
        width=1200,
        height=800,
    )
    fig.write_html(BASE_FOLDER / f"{model_name}-sankey_diagram.html")


# def to_rollouts():
#     with open(BASE_FOLDER / "results.json") as f:
#         results = json.load(f)

#     # insider trading report
#     for r in results:
#         # is the trade made with insider info?
#         if "trade_label" in r and r["trade_label"] == "yes":
#             transcript = make_full_transcript(r)
#             RolloutDialogue(
#                 input_messages=transcript[:-3],
#                 expected_label=Label.Deceptive,
#                 input_str=None,  # or tokenizer formatted string
#                 outptut_messages=transcript[-3:],
#                 output_strs=[m.content for m in transcript[-3:]],
#                 grader_reasoning=[],
#                 graded_deception=[],
#                 deception_percentage=None,
#                 metadata=None,
#             )

#     data = {
#         "dataset_class": self.dataset_class_name,
#         "base_name": self.base_name,
#         "variant": self.variant,
#         "model": self.model if isinstance(self.model, str) else self.model.name_or_path,
#         "rollouts": [dialogue.model_dump() for dialogue in self.rollout_dialogues],
#     }

#     with open(filename, "w") as f:
#         json.dump(data, f, default=lambda o: o.__dict__, indent=2)


if __name__ == "__main__":
    fire.Fire(create_sankey_diagram)
