"""
Build a KMC (Kangaroo Mother Care) QI Run Chart:
Baseline average = 4 hrs/day  ->  Target = 8 hrs/day
Output: Excel workbook with data table + native run chart (12 weekly data points).
"""

from openpyxl import Workbook
from openpyxl.styles import Font, PatternFill, Alignment, Border, Side
from openpyxl.chart import LineChart, Reference, Series
from openpyxl.chart.marker import Marker
from openpyxl.chart.data_source import NumDataSource, NumRef
from openpyxl.utils import get_column_letter
from openpyxl.drawing.line import LineProperties
from openpyxl.chart.shapes import GraphicalProperties
from openpyxl.comments import Comment

FONT_NAME = "Arial"

wb = Workbook()
ws = wb.active
ws.title = "KMC Run Chart"

# ---------------------------------------------------------------
# Data: 12 weekly data points
# Weeks 1-4 = baseline phase (avg 4 hrs/day)
# Weeks 5-12 = improvement phase (PDSA cycles) trending to target 8 hrs/day
# ---------------------------------------------------------------
weeks = list(range(1, 13))
kmc_hours = [3.8, 4.2, 3.9, 4.1,      # baseline (avg = 4.0)
             4.5, 5.0, 5.3, 5.8,      # PDSA 1 & 2
             6.4, 7.0, 7.5, 8.1]      # PDSA 3 -> target reached

baseline_median = 4.0
target = 8.0

annotations = {
    1: "Baseline data collection start",
    4: "Baseline avg = 4.0 hrs/day",
    5: "PDSA 1: Staff training on KMC benefits",
    8: "PDSA 2: Dedicated KMC corner + reclining chairs",
    11: "PDSA 3: Family counseling + flexible visiting hours",
    12: "Target of 8 hrs/day achieved",
}

# ---------------------------------------------------------------
# Title block
# ---------------------------------------------------------------
ws.merge_cells("A1:F1")
ws["A1"] = "KMC Quality Improvement Run Chart"
ws["A1"].font = Font(name=FONT_NAME, size=16, bold=True, color="1F4E78")
ws["A1"].alignment = Alignment(horizontal="center")

ws.merge_cells("A2:F2")
ws["A2"] = "Measure: Average daily duration of Kangaroo Mother Care (hrs/day) | Baseline = 4 hrs/day -> Target = 8 hrs/day"
ws["A2"].font = Font(name=FONT_NAME, size=11, italic=True, color="595959")
ws["A2"].alignment = Alignment(horizontal="center")

start_row = 4
headers = ["Week", "KMC Duration (hrs/day)", "Baseline Median", "Target", "PDSA / Annotation"]
header_fill = PatternFill(start_color="1F4E78", end_color="1F4E78", fill_type="solid")

for i, h in enumerate(headers, start=1):
    c = ws.cell(row=start_row, column=i, value=h)
    c.font = Font(name=FONT_NAME, size=11, bold=True, color="FFFFFF")
    c.fill = header_fill
    c.alignment = Alignment(horizontal="center", vertical="center", wrap_text=True)

thin = Side(style="thin", color="BFBFBF")
border = Border(left=thin, right=thin, top=thin, bottom=thin)

for idx, wk in enumerate(weeks):
    r = start_row + 1 + idx
    ws.cell(row=r, column=1, value=wk).font = Font(name=FONT_NAME, size=10)
    ws.cell(row=r, column=2, value=kmc_hours[idx]).font = Font(name=FONT_NAME, size=10)
    ws.cell(row=r, column=3, value=baseline_median).font = Font(name=FONT_NAME, size=10, color="808080")
    ws.cell(row=r, column=4, value=target).font = Font(name=FONT_NAME, size=10, color="C00000")
    note = annotations.get(wk, "")
    ws.cell(row=r, column=5, value=note).font = Font(name=FONT_NAME, size=10, italic=True)
    for col in range(1, 6):
        ws.cell(row=r, column=col).border = border
        ws.cell(row=r, column=col).alignment = Alignment(horizontal="center", vertical="center", wrap_text=True)

last_row = start_row + len(weeks)

# Column widths
widths = [8, 22, 16, 12, 42]
for i, w in enumerate(widths, start=1):
    ws.column_dimensions[get_column_letter(i)].width = w

# ---------------------------------------------------------------
# Run chart rule notes
# ---------------------------------------------------------------
notes_row = last_row + 2
ws.cell(row=notes_row, column=1, value="Run Chart Interpretation Notes:").font = Font(name=FONT_NAME, size=11, bold=True)
rules = [
    "- Baseline median (dashed grey line) = 4.0 hrs/day, calculated from Weeks 1-4.",
    "- Target (dashed red line) = 8.0 hrs/day.",
    "- A trend (5+ consecutive points moving in one direction) is visible from Week 4 to Week 12, indicating a statistically meaningful improvement (not random variation).",
    "- Weeks 5, 8 and 11 mark PDSA cycle interventions associated with the sustained upward trend.",
    "- Target of 8 hrs/day achieved at Week 12.",
]
for i, txt in enumerate(rules, start=1):
    ws.cell(row=notes_row + i, column=1, value=txt).font = Font(name=FONT_NAME, size=9, italic=True, color="595959")
    ws.merge_cells(start_row=notes_row + i, start_column=1, end_row=notes_row + i, end_column=6)

# ---------------------------------------------------------------
# Chart
# ---------------------------------------------------------------
chart = LineChart()
chart.title = "KMC Duration Run Chart (Baseline 4 hrs/day -> Target 8 hrs/day)"
chart.style = 2
chart.y_axis.title = "KMC Duration (hrs/day)"
chart.x_axis.title = "Week"
chart.y_axis.scaling.min = 0
chart.y_axis.scaling.max = 9
chart.height = 10
chart.width = 24

cats = Reference(ws, min_col=1, min_row=start_row + 1, max_row=last_row)

# Series 1: KMC duration (actual data) with markers
data_kmc = Reference(ws, min_col=2, min_row=start_row, max_row=last_row)
chart.add_data(data_kmc, titles_from_data=True)

# Series 2: Baseline median line
data_base = Reference(ws, min_col=3, min_row=start_row, max_row=last_row)
chart.add_data(data_base, titles_from_data=True)

# Series 3: Target line
data_target = Reference(ws, min_col=4, min_row=start_row, max_row=last_row)
chart.add_data(data_target, titles_from_data=True)

chart.set_categories(cats)

# Style series
s_kmc = chart.series[0]
s_kmc.marker = Marker(symbol="circle", size=7)
s_kmc.marker.graphicalProperties.solidFill = "1F4E78"
s_kmc.marker.graphicalProperties.line.solidFill = "1F4E78"
s_kmc.graphicalProperties.line.solidFill = "1F4E78"
s_kmc.graphicalProperties.line.width = 22000
s_kmc.smooth = False

s_base = chart.series[1]
s_base.graphicalProperties.line.solidFill = "808080"
s_base.graphicalProperties.line.width = 15000
s_base.graphicalProperties.line.dashStyle = "dash"
s_base.marker = Marker(symbol="none")
s_base.smooth = False

s_target = chart.series[2]
s_target.graphicalProperties.line.solidFill = "C00000"
s_target.graphicalProperties.line.width = 15000
s_target.graphicalProperties.line.dashStyle = "dash"
s_target.marker = Marker(symbol="none")
s_target.smooth = False

chart.series[0].tx.strRef = None  # keep default titles from header row

ws.add_chart(chart, "H4")

out_path = "/tmp/workspace/kmc-qi-run-chart/KMC_QI_Run_Chart.xlsx"
wb.save(out_path)
print("Saved:", out_path)