import sys

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import plotly.io as pio

import roadmap_python.analysis_utils as utils

print("Python version:", sys.version.split()[0])
print("pandas version:", pd.__version__)
print("numpy version:", np.__version__)
print("matplotlib version:", plt.matplotlib.__version__)

%load_ext autoreload
%autoreload 2

paths = utils.get_paths()
smaple_proj_f = "3PointPlan-example.xml"
tasks_ms_df, links_df = utils.read_mspdi_tasks_and_links(
    paths.raw / "open-msp-viewer-sample" / smaple_proj_f
)

# magnify the dates so span multiple FY, for demo purposes later
factor = 18  # 6x longer overall (tweak until it spans multiple FYs)
anchor = tasks_ms_df["start"].min()
for col in ["start", "finish"]:
    tasks_ms_df[col] = anchor + (tasks_ms_df[col] - anchor) * factor

display(tasks_ms_df)

pio.renderers.default = "notebook_connected"  # good default for classic notebooks

df = tasks_ms_df.copy()

# Keep only tasks that can be drawn as bars
df = df[df["start"].notna() & df["finish"].notna()].copy()

# Ensure one row per task on the y-axis (avoid duplicate-name collisions)
# Prefer WBS if present; fall back to UID.
df["task_label"] = df.apply(
    lambda r: f"{r['wbs']} — {r['name']}"
    if pd.notna(r.get("wbs")) and str(r.get("wbs")).strip()
    else f"UID {r['uid']} — {r['name']}",
    axis=1,
)

# Milestones often have start == finish; give them a tiny width so they render
is_milestone = (df["milestone"]) | (df["start"] == df["finish"])
df["finish_plot"] = df["finish"]
df.loc[is_milestone, "finish_plot"] = df.loc[is_milestone, "finish"] + pd.Timedelta(weeks=2)

# Pick a stable order (WBS-ish if available, otherwise by start)
sort_cols = ["start", "outline_level", "uid"]
if "wbs" in df.columns and df["wbs"].notna().any():
    sort_cols = ["wbs", "start", "outline_level", "uid"]

df = df.sort_values(sort_cols, kind="stable")

fig = px.timeline(
    df,
    x_start="start",
    x_end="finish_plot",
    y="task_label",
    hover_data={
        "uid": True,
        "name": True,
        "wbs": True,
        "outline_level": True,
        "summary": True,
        "milestone": True,
        "start": "|%Y-%m-%d",
        "finish": "|%Y-%m-%d",
        "finish_plot": False,
    },
)

fig.update_yaxes(autorange="reversed")  # top-to-bottom like a plan
fig.update_layout(
    height=max(750, 18 * len(df)),
    xaxis_rangeslider_visible=True,  # handy for zooming
    margin=dict(l=20, r=20, t=40, b=20),
)

display(fig)
fig.write_html(paths.figures / "fig0-simple-timeline.html")

df = tasks_ms_df.copy()
# Keep only tasks that can be drawn as bars
df = df[df["start"].notna() & df["finish"].notna()].copy()
dd_path = paths.raw / "roadmap_data_dictionary_tasks.csv"
dd = pd.read_csv(dd_path)
dd = dd[
    [
        "uid",
        "y_id",
        "y_label",
        "item_type",
        "group",
        "shape",
        "color",
        "include",
        "fill_alpha",
        "bar_pattern",
        "bar_line_color",
        "bar_line_width",
        "marker_size",
        "marker_line_color",
        "marker_line_width",
    ]
].copy()

df_p = df.merge(dd, on="uid", how="left")
# display(df_p)
# Fallbacks so you can leave blanks in the CSV and still render
df_p["y_label"] = df_p["y_label"].fillna(df_p["name"])
df_p["item_type"] = df_p["item_type"].fillna(
    df_p["milestone"].fillna(False).map(lambda m: "milestone" if m else "task")
)
df_p["shape"] = df_p["shape"].fillna(
    df_p["item_type"].map(lambda t: "star" if t == "milestone" else "rectangle")
)
# make sure it's clean 0..1 floats
df_p["fill_alpha"] = pd.to_numeric(df_p["fill_alpha"], errors="coerce").fillna(1.0).clip(0, 1)

df_p = df_p[(df_p["include"])]

df_p.head()

#################################
# Tasks
#################################
df_tasks = df_p[
    (df_p["item_type"] == "task") & df_p["start"].notna() & df_p["finish"].notna()
].copy()
# df_tasks = df_tasks[(df_tasks["outline_level"] >= 2)]

y_col = "y_id"  # whatever you’re using on your y-axis
fig = px.timeline(
    df_tasks,
    x_start="start",
    x_end="finish",
    y=y_col,
    hover_data=["uid", "wbs", "outline_level", "summary"],
    text="name",  # labels from your CSV
)

# Put text inside bars (timeline() doesn't take textposition directly)
fig.update_traces(
    textposition="inside", insidetextanchor="middle", selector=dict(type="bar")
)  # :contentReference[oaicite:3]{index=3}

fig.update_traces(
    marker_color=df_tasks["color"].tolist(),
    marker_pattern_shape=df_tasks["bar_pattern"].fillna("").tolist(),
    marker_opacity=df_tasks["fill_alpha"].tolist(),
    selector=dict(type="bar"),
)

# y tick labels
fig.update_yaxes(
    tickmode="array",
    tickvals=df_tasks["y_id"].tolist(),
    ticktext=df_tasks["y_label"].tolist(),
)

# # Apply per-bar colors from the dictionary (no legend, but obeys CSV)
# fig.update_traces(
#     marker=dict(color=df_tasks["color"].tolist()),
#     selector=dict(type="bar")
# )  # :contentReference[oaicite:4]{index=4}


#################################
# Milestones
#################################
df_ms = df_p[(df_p["item_type"] == "milestone") & df_p["start"].notna()].copy()

fig.add_trace(
    go.Scatter(
        x=df_ms["start"],
        y=df_ms[y_col],
        mode="markers",
        marker=dict(
            symbol=df_ms["shape"].tolist(),  # e.g. "star"
            color=df_ms["color"].tolist(),  # e.g. "gold" or "#FFD700"
            size=14,
            line=dict(width=1, color="black"),
        ),
        text=df_ms["name"],
        hovertemplate="%{text}<extra></extra>",
        name="Milestones",
    )
)  # :contentReference[oaicite:5]{index=5}


#################################
# Overall Figure Styling
#################################

# Update timeline to show FY quarters
fy_start_month = 10  # change if your FY starts in a different month

xmin = df_p["start"].min()
xmax = (df_p["finish_plot"] if "finish_plot" in df_p.columns else df_p["finish"]).max()

fig.update_yaxes(autorange="reversed")  # top-to-bottom like a plan
fig.update_layout(
    height=max(1000, 18 * len(df)),
    xaxis_rangeslider_visible=True,  # handy for zooming
    margin=dict(l=20, r=20, t=40, b=20),
)


display(fig)
fig.write_html(paths.figures / "fig1-timeline-add-datadict.html")

df = tasks_ms_df.copy()
ddv2_path = paths.raw / "roadmap_data_dictionary_tasks_v2.csv"
# display(df)
fig, df_merged = utils.build_roadmap_timeline_from_dd(df, ddv2_path, y_col="y_id")

fig.update_layout(
    height=max(1000, 18 * len(df)),
    xaxis_rangeslider_visible=True,  # handy for zooming
    margin=dict(l=20, r=20, t=40, b=20),
)

display(fig)
fig.write_html(paths.figures / "fig2-timeline-altview.html")

paths = utils.get_paths()
smaple_proj_f = "dod_acquisition_lifecycle_demo.xml"
tasks_ms_df, links_df = utils.read_mspdi_tasks_and_links(paths.raw / smaple_proj_f)
df = tasks_ms_df.copy()
# Keep only tasks that can be drawn as bars
df = df[df["start"].notna() & df["finish"].notna()].copy()
display(df.head())

ddv_path = paths.raw / "roadmap_data_dictionary_dod_demo_v1.csv"
fig, df_merged = utils.build_roadmap_timeline_from_dd(df, ddv_path, y_col="y_id", y_label="y_label")

fig.update_layout(
    height=max(2000, 18 * len(df)),
    xaxis_rangeslider_visible=True,  # handy for zooming
    margin=dict(l=20, r=20, t=40, b=20),
)
display(fig)
fig.write_html(paths.figures / "fig3-timeline-dodacq-view1.html")

ddv_path = paths.raw / "roadmap_data_dictionary_dod_demo_v4.csv"
fig, df_merged = utils.build_roadmap_timeline_from_dd(
    df, ddv_path, bar_text_wrap_width=32, xaxis_fy_quarters=True
)

fig.update_layout(
    height=max(1000, 18 * len(df)),
    margin=dict(l=20, r=20, t=40, b=20),
)

display(fig)
fig.write_html(paths.figures / "fig4-timeline-dodacq-view2.html")