import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.patches import FancyBboxPatch, FancyArrowPatch

fig, ax = plt.subplots(figsize=(28, 38))
ax.set_xlim(0, 28)
ax.set_ylim(0, 38)
ax.axis('off')
fig.patch.set_facecolor('#F8F9FA')

# ─── COLOUR PALETTE ───────────────────────────────────────────────
C_TITLE   = '#1A237E'   # deep navy
C_START   = '#1565C0'   # blue  – entry node
C_SYMPTOM = '#0277BD'   # mid-blue – symptom boxes
C_INVEST  = '#00838F'   # teal – investigation
C_CANCER  = '#BF360C'   # brick-red – cancer type boxes
C_STAGE   = '#6A1B9A'   # purple – staging
C_TREAT   = '#2E7D32'   # green – treatment
C_REFER   = '#F57F17'   # amber – referral / flags
C_WHITE   = '#FFFFFF'
C_LTGRAY  = '#ECEFF1'

def box(ax, x, y, w, h, text, fc, tc='white', fs=9, bold=False,
        radius=0.35, ec=None, lw=1.5, va='center', wrap=False):
    ec = ec or fc
    fancy = FancyBboxPatch((x - w/2, y - h/2), w, h,
                            boxstyle=f"round,pad=0.05,rounding_size={radius}",
                            facecolor=fc, edgecolor=ec, linewidth=lw, zorder=3)
    ax.add_patch(fancy)
    weight = 'bold' if bold else 'normal'
    ax.text(x, y, text, ha='center', va=va, fontsize=fs, color=tc,
            fontweight=weight, zorder=4,
            wrap=wrap,
            multialignment='center')

def diamond(ax, x, y, w, h, text, fc, tc='white', fs=8.5):
    dx, dy = w/2, h/2
    coords = [(x, y+dy), (x+dx, y), (x, y-dy), (x-dx, y)]
    poly = plt.Polygon(coords, closed=True, facecolor=fc, edgecolor=fc, linewidth=1.5, zorder=3)
    ax.add_patch(poly)
    ax.text(x, y, text, ha='center', va='center', fontsize=fs, color=tc,
            fontweight='bold', zorder=4, multialignment='center')

def arrow(ax, x1, y1, x2, y2, color='#455A64', lw=1.5, style='->', label=''):
    ax.annotate('', xy=(x2, y2), xytext=(x1, y1),
                arrowprops=dict(arrowstyle=style, color=color, lw=lw),
                zorder=2)
    if label:
        mx, my = (x1+x2)/2, (y1+y2)/2
        ax.text(mx+0.15, my, label, fontsize=7.5, color=color, fontstyle='italic')

# ══════════════════════════════════════════════════════
# TITLE
# ══════════════════════════════════════════════════════
ax.add_patch(FancyBboxPatch((0.3, 36.5), 27.4, 1.3,
             boxstyle="round,pad=0.1", facecolor=C_TITLE, edgecolor=C_TITLE, zorder=3))
ax.text(14, 37.15, 'DIAGNOSTIC ALGORITHM FOR COMMON CANCER PRESENTATIONS',
        ha='center', va='center', fontsize=15, color='white', fontweight='bold', zorder=4)
ax.text(14, 36.7, 'OPD / Primary Care Setting  |  Orris Medical Reference',
        ha='center', va='center', fontsize=9, color='#B3C5EF', zorder=4)

# ══════════════════════════════════════════════════════
# ENTRY  (top-centre)
# ══════════════════════════════════════════════════════
box(ax, 14, 35.6, 7, 0.7,
    'PATIENT PRESENTS TO OPD WITH SYMPTOMS / COMPLAINT',
    C_START, fs=10, bold=True)

arrow(ax, 14, 35.25, 14, 34.65)

# RED FLAG SCREEN
box(ax, 14, 34.3, 9, 0.65,
    'RED FLAG SYMPTOM SCREEN  (weight loss · unexplained anaemia · non-healing ulcer · painless lump · unexplained bleeding)',
    C_REFER, fs=8, bold=True, tc='white')

# YES / NO branches
arrow(ax, 9.5, 34.3,  5, 34.3, color=C_REFER, label='NO red flags')
arrow(ax, 18.5, 34.3, 23, 34.3, color=C_CANCER, label='RED FLAGS present')

box(ax,  3.5, 34.3, 2.8, 0.6, 'Routine management\n& reassess in 4 wks', C_LTGRAY, tc='#37474F', fs=8)
box(ax, 24.5, 34.3, 2.8, 0.6, 'URGENT REFERRAL\nto oncology / specialist', C_CANCER, fs=8, bold=True)

arrow(ax, 14, 33.97, 14, 33.2)

# ══════════════════════════════════════════════════════
# PRESENTING SYMPTOM TRIAGE  (wide row)
# ══════════════════════════════════════════════════════
box(ax, 14, 32.85, 9, 0.65,
    'IDENTIFY PRESENTING SYMPTOM CLUSTER  →  Route to cancer-specific pathway',
    C_SYMPTOM, fs=9, bold=True)

# 6 symptom branches  x-positions
SX = [2.2, 6.4, 10.6, 14.8, 19.0, 23.2]
SY_TOP = 32.0
SY_BOX = 31.5

symptom_labels = [
    'Breast lump /\nnipple discharge',
    'Cough · haemo-\nptysis · wt loss',
    'Rectal bleed ·\nbowel change',
    'Post-coital /\nvaginal bleed',
    'Neck / thyroid\nswelling',
    'Oral ulcer >\n2 weeks / LUTS'
]
sym_colors = [C_SYMPTOM]*6

for sx, sl in zip(SX, symptom_labels):
    arrow(ax, 14, 32.52, sx, SY_TOP+0.03)
    box(ax, sx, SY_BOX, 3.6, 0.8, sl, C_SYMPTOM, fs=8, bold=True)

# ══════════════════════════════════════════════════════
# COLUMN HEADERS (cancer labels)
# ══════════════════════════════════════════════════════
cancer_names = ['BREAST\nCANCER', 'LUNG\nCANCER', 'COLORECTAL\nCANCER',
                'CERVICAL\nCANCER', 'THYROID\nCANCER', 'ORAL / PROSTATE\nCANCER']

CY_LABEL = 30.6
for sx, cn in zip(SX, cancer_names):
    arrow(ax, sx, SY_BOX - 0.4, sx, CY_LABEL + 0.35)
    box(ax, sx, CY_LABEL, 3.6, 0.65, cn, C_CANCER, fs=8.5, bold=True)

# ══════════════════════════════════════════════════════
# HISTORY & EXAMINATION  row
# ══════════════════════════════════════════════════════
hx_data = [
    'Age, family hx\nBRCA1/2\nHRT use',
    'Smoking hx\nOccupational\nexposure',
    'Family hx, IBD\nDiet, age >50\nFAP / Lynch',
    'Sexual hx, HPV\nvaccine status\nPap smear hx',
    'Radiation hx\nFamily hx MEN\nRapid growth',
    'Tobacco/alcohol\nBetel nut\nPSA if male'
]

HY = 29.3
for sx, hd in zip(SX, hx_data):
    arrow(ax, sx, CY_LABEL-0.33, sx, HY+0.42)
    box(ax, sx, HY, 3.6, 0.78, 'HISTORY\n' + hd, C_LTGRAY, tc='#1A237E', fs=7.5, bold=False, ec='#90A4AE')

# ══════════════════════════════════════════════════════
# PHYSICAL EXAMINATION
# ══════════════════════════════════════════════════════
pe_data = [
    'Lump: size/\nfixity/skin\nAxilla nodes',
    'Chest auscult.\nSupraclav.\nnodes',
    'PR exam\nAbdominal\nmass, ascites',
    'Speculum +\nBimanual exam\nCervix inspect.',
    'Thyroid size\nfixity, LN\nVoice change',
    'Oral inspection\nDRE (prostate)\nNeck nodes'
]

PE_Y = 28.05
for sx, pd in zip(SX, pe_data):
    arrow(ax, sx, HY-0.39, sx, PE_Y+0.42)
    box(ax, sx, PE_Y, 3.6, 0.78, 'EXAMINATION\n' + pd, C_LTGRAY, tc='#1A237E', fs=7.5, ec='#90A4AE')

# ══════════════════════════════════════════════════════
# FIRST-LINE INVESTIGATIONS
# ══════════════════════════════════════════════════════
inv1_data = [
    'Mammogram\n+ USG breast\nFNAC / CNB',
    'CXR\nCT chest\nSputum cytol.',
    'Colonoscopy\n(gold std)\nCEA level',
    'Pap smear\nColposcopy\n+ biopsy',
    'TSH + fT4\nNeck USG\nFNAC nodule',
    'OPG / biopsy\nPSA (prostate)\nNeck USG'
]

INV1_Y = 26.75
for sx, iv in zip(SX, inv1_data):
    arrow(ax, sx, PE_Y-0.39, sx, INV1_Y+0.45)
    box(ax, sx, INV1_Y, 3.6, 0.84, '1st LINE TESTS\n' + iv, C_INVEST, fs=7.5)

# ══════════════════════════════════════════════════════
# BIOPSY / HISTOLOGY CONFIRMATION
# ══════════════════════════════════════════════════════
bx_data = [
    'Core needle bx\nER/PR/HER2\nGrade',
    'Bronchoscopy bx\nCT-guided FNA\nMol: EGFR/ALK',
    'Colonoscopic bx\nMSI/KRAS/BRAF\nstatus',
    'Cervical punch\nbiopsy\nHPV typing',
    'FNAC (Bethesda)\nCore bx if\nindeterminate',
    'Incisional bx\nGleason score\n(prostate)'
]

BX_Y = 25.35
for sx, bx in zip(SX, bx_data):
    arrow(ax, sx, INV1_Y-0.42, sx, BX_Y+0.45)
    box(ax, sx, BX_Y, 3.6, 0.84, 'BIOPSY / HISTOLOGY\n' + bx, C_INVEST, fs=7.5)

# ══════════════════════════════════════════════════════
# STAGING INVESTIGATIONS
# ══════════════════════════════════════════════════════
stg_data = [
    'CT CAP\nBone scan\nPET (if HER2+)',
    'PET-CT\nBrain MRI\nLFTs/LDH',
    'CT abd/pelvis\nCXR\nLFTs',
    'MRI pelvis\nCystoscopy\nIVU (FIGO)',
    'CT neck/chest\nCalcitonin\n(medullary)',
    'PET/CT\nPSA + bone\nscan if PSA>20'
]

STG_Y = 23.95
for sx, sg in zip(SX, stg_data):
    arrow(ax, sx, BX_Y-0.42, sx, STG_Y+0.45)
    box(ax, sx, STG_Y, 3.6, 0.84, 'STAGING\n' + sg, C_STAGE, fs=7.5)

# ══════════════════════════════════════════════════════
# TNM STAGE GROUP  (merges all 6 columns)
# ══════════════════════════════════════════════════════
TNM_Y = 22.7
for sx in SX:
    arrow(ax, sx, STG_Y-0.42, sx, TNM_Y+0.38)

box(ax, 14, TNM_Y, 25.5, 0.7,
    'ASSIGN TNM STAGE GROUP   |   Stage I  •  Stage II  •  Stage III  •  Stage IV',
    C_STAGE, fs=10, bold=True)

arrow(ax, 14, TNM_Y-0.35, 14, 21.9)

# ══════════════════════════════════════════════════════
# STAGE DECISION DIAMOND
# ══════════════════════════════════════════════════════
diamond(ax, 14, 21.45, 5, 0.85, 'Stage?', C_STAGE, fs=9)

# branches
arrow(ax, 11.5, 21.45,  5.2, 21.45, color=C_TREAT, label='Stage I–II')
arrow(ax, 16.5, 21.45, 23.2, 21.45, color=C_CANCER, label='Stage III–IV')
arrow(ax, 14, 21.02, 14, 20.3,  color='#546E7A', label='')

# ══════════════════════════════════════════════════════
# TREATMENT BLOCKS  (3 columns)
# ══════════════════════════════════════════════════════
# Early stage (left)
box(ax, 3.8, 21.45, 5.5, 0.72,
    'EARLY STAGE  (I – II)\nSurgery ± Adjuvant chemo/RT\nCurative intent',
    C_TREAT, fs=8.5, bold=True)

# Advanced stage (right)
box(ax, 24.2, 21.45, 5.5, 0.72,
    'ADVANCED STAGE  (III – IV)\nNeoadjuvant / Palliative\nMultidisciplinary approach',
    C_CANCER, fs=8.5, bold=True)

# Tumour board (centre-below diamond)
box(ax, 14, 20.0, 12, 0.62,
    'MULTIDISCIPLINARY TUMOUR BOARD  (Surgery · Oncology · Radiology · Pathology · Palliative Care)',
    '#37474F', fs=8.5, bold=True)

arrow(ax, 14, 19.69, 14, 18.95)

# ══════════════════════════════════════════════════════
# TREATMENT MODALITIES  (5 boxes in a row)
# ══════════════════════════════════════════════════════
tx_labels = [
    'SURGERY\n• Wide excision\n• Lymph node dissection\n• Reconstructive options',
    'RADIOTHERAPY\n• External beam (EBRT)\n• Brachytherapy\n• Stereotactic (SBRT)',
    'CHEMOTHERAPY\n• Neoadjuvant / Adjuvant\n• Palliative intent\n• Regimen by cancer type',
    'TARGETED /\nIMMUNOTHERAPY\n• Anti-HER2 (trastuzumab)\n• EGFR / ALK inhibitors\n• Checkpoint inhibitors',
    'HORMONAL\nTHERAPY\n• ER+ breast (tamoxifen)\n• Prostate (LHRH agonist)\n• Thyroid (TSH suppression)'
]
TX_X = [3.5, 8.2, 14.0, 19.8, 24.5]
TX_Y = 18.3
for tx, tl in zip(TX_X, tx_labels):
    arrow(ax, 14, 18.95, tx, TX_Y+0.55)
    box(ax, tx, TX_Y, 4.4, 1.0, tl, C_TREAT, fs=7.5)

arrow(ax, 14, TX_Y-0.5, 14, 17.15)

# ══════════════════════════════════════════════════════
# RESPONSE ASSESSMENT
# ══════════════════════════════════════════════════════
box(ax, 14, 16.85, 14, 0.55,
    'RESPONSE ASSESSMENT  (repeat imaging · tumour markers · clinical exam)  after 2–3 cycles / post-surgery',
    '#455A64', fs=8.5, bold=True)

# 3 outcomes
RES_X = [5.5, 14, 22.5]
RES_LABELS = ['COMPLETE\nRESPONSE', 'PARTIAL / STABLE\nDISEASE', 'PROGRESSIVE\nDISEASE']
RES_COLORS = [C_TREAT, C_INVEST, C_CANCER]
RES_Y = 15.85

for rx, rl, rc in zip(RES_X, RES_LABELS, RES_COLORS):
    arrow(ax, 14, 16.58, rx, RES_Y+0.3)
    box(ax, rx, RES_Y, 4.2, 0.56, rl, rc, fs=8, bold=True)

# sub-actions
sub_actions = [
    'Surveillance\nprogram\n(scheduled follow-up)',
    'Continue ±\nmodify regimen\n(dose / add agent)',
    'Switch to 2nd line\nPalliative / BSC\nClinical trial'
]
SUB_Y = 14.75
for rx, sa, rc in zip(RES_X, sub_actions, RES_COLORS):
    arrow(ax, rx, RES_Y-0.28, rx, SUB_Y+0.38)
    box(ax, rx, SUB_Y, 4.2, 0.72, sa, rc, fs=8, ec=rc)

arrow(ax, 14, SUB_Y-0.36, 14, 13.8)

# ══════════════════════════════════════════════════════
# FOLLOW-UP / SURVEILLANCE
# ══════════════════════════════════════════════════════
box(ax, 14, 13.55, 25, 0.48,
    'FOLLOW-UP & SURVEILLANCE  –  Tumour markers · Imaging · Clinical exam · Toxicity monitoring · Psychosocial support',
    '#37474F', fs=8.5, bold=True)

# ══════════════════════════════════════════════════════
# CANCER-SPECIFIC QUICK-REF TABLE  (bottom section)
# ══════════════════════════════════════════════════════
box(ax, 14, 12.85, 27, 0.45,
    'CANCER-SPECIFIC QUICK REFERENCE',
    C_TITLE, fs=9, bold=True)

# Table header row
cols = ['CANCER', 'KEY SYMPTOM', 'RISK FACTOR #1', '1st LINE TEST', 'TUMOUR MARKER', '5-YR SURVIVAL (Stage I)']
col_x = [1.5, 5.5, 9.8, 13.8, 18.0, 22.5]
COL_W = [2.6, 3.5, 3.8, 3.6, 3.6, 4.5]
HDR_Y = 12.25
for cx, cw, cl in zip(col_x, COL_W, cols):
    box(ax, cx, HDR_Y, cw, 0.45, cl, '#37474F', fs=8, bold=True, radius=0.1)

# Table rows
rows = [
    ['Breast',       'Painless lump',         'Age / BRCA mutation',  'Mammogram + CNB',  'CA 15-3',   '~99%'],
    ['Lung',         'Cough + haemoptysis',   'Cigarette smoking',    'CT chest',         'CEA / NSE', '~60%'],
    ['Colorectal',   'Rectal bleed / anaemia','Diet + age >50',       'Colonoscopy',      'CEA',       '~90%'],
    ['Cervical',     'Post-coital bleed',     'HPV 16 / 18',         'Pap smear',        'SCC Ag',    '~92%'],
    ['Thyroid',      'Neck nodule',           'Radiation exposure',   'USG + FNAC',       'Thyroglobulin','~99%'],
    ['Oral',         'Non-healing ulcer',     'Tobacco + alcohol',   'Biopsy',           'SCC Ag',    '~83%'],
    ['Prostate',     'LUTS / bone pain',      'Age + family history', 'PSA + DRE + bx',  'PSA',       '~99%'],
]

ROW_COLORS = ['#E3F2FD','#E8F5E9','#FFF3E0','#FCE4EC','#EDE7F6','#E0F7FA','#F9FBE7']
ROW_H = 0.44
row_y = HDR_Y - ROW_H * 0.6
for i, (row, rc) in enumerate(zip(rows, ROW_COLORS)):
    row_y -= ROW_H
    for cx, cw, cell in zip(col_x, COL_W, row):
        box(ax, cx, row_y, cw, ROW_H-0.03, cell, rc, tc='#1A237E', fs=8, radius=0.08, ec='#CFD8DC', lw=0.8)

# ══════════════════════════════════════════════════════
# LEGEND  (bottom)
# ══════════════════════════════════════════════════════
legend_items = [
    (C_SYMPTOM, 'Symptom / Clinical step'),
    (C_INVEST,  'Investigation'),
    (C_CANCER,  'Cancer diagnosis / Advanced'),
    (C_STAGE,   'Staging'),
    (C_TREAT,   'Treatment / Good outcome'),
    (C_REFER,   'Red flags / Urgent action'),
]
LX, LY = 1.0, 0.55
for i, (lc, lt) in enumerate(legend_items):
    bx = LX + i * 4.5
    ax.add_patch(FancyBboxPatch((bx, LY), 0.55, 0.35, boxstyle="round,pad=0.05",
                                facecolor=lc, edgecolor=lc, zorder=4))
    ax.text(bx + 0.7, LY + 0.175, lt, fontsize=8, va='center', color='#37474F', zorder=4)

ax.text(14, 0.2, 'Orris Medical Reference  |  Based on AJCC 8th Edition, Harrison\'s 22e, Robbins Pathology, Schwartz Surgery',
        ha='center', va='center', fontsize=7.5, color='#90A4AE')

plt.tight_layout(pad=0.2)
plt.savefig('/tmp/workspace/cancer-flowchart/cancer_diagnostic_algorithm.png',
            dpi=180, bbox_inches='tight', facecolor='#F8F9FA')
print("Saved successfully")