"""
Generate sample_sequencing_plan.xlsx
Labels every sample as Already sequenced / To be sequenced / Skip
at the chosen concentration threshold, with colour-coded rows.
"""

import numpy as np
import pandas as pd
from openpyxl import load_workbook
from openpyxl.styles import PatternFill

THRESHOLD  = 0.25
OUT_FILE   = "sample_sequencing_plan.xlsx"

PILOT_IDS  = {("A", n) for n in [1, 2, 3, 4, 5, 6]} | \
             {("B", n) for n in [1, 2, 15, 16, 17, 18]}

FILLS = {
    "Already sequenced": PatternFill("solid", fgColor="C6EFCE"),  # green
    "To be sequenced":   PatternFill("solid", fgColor="FFEB9C"),  # yellow
    "Skip":              PatternFill("solid", fgColor="FFC7CE"),  # red/pink
}

# ── load data ─────────────────────────────────────────────────────────────────
patients = pd.read_excel("patients.xlsx", usecols=["Number", "Exf"])
patients = patients.dropna(subset=["Number"])
patients["Number"] = patients["Number"].astype(int)
patients["Exf"]    = patients["Exf"].astype(int)

def parse_conc(v):
    if isinstance(v, str): return np.nan
    try: return float(v)
    except: return np.nan

qa = pd.read_excel("BEA25P077_RP.xlsx", sheet_name="A-samples")
qb = pd.read_excel("BEA25P077_RP.xlsx", sheet_name="B-samples")
qa["conc_val"] = qa["conc (ng/ul)"].apply(parse_conc)
qb["conc_val"] = qb["Conc [ng/ul]"].apply(parse_conc)
qa["Number"] = pd.to_numeric(qa["ID"].str.extract(r"^(\d+)A")[0], errors="coerce")
qb["Number"] = pd.to_numeric(qb["ID"].str.extract(r"^(\d+)B")[0], errors="coerce")
qa = qa.dropna(subset=["Number"]); qa["Number"] = qa["Number"].astype(int); qa["tissue"] = "A"
qb = qb.dropna(subset=["Number"]); qb["Number"] = qb["Number"].astype(int); qb["tissue"] = "B"
qa = qa.merge(patients, on="Number", how="left")
qb = qb.merge(patients, on="Number", how="left")

# ── build plan ────────────────────────────────────────────────────────────────
rows = []
for _, r in pd.concat([qa, qb], ignore_index=True).iterrows():
    tissue, num, conc = r["tissue"], r["Number"], r["conc_val"]
    if (tissue, num) in PILOT_IDS:
        status = "Already sequenced"
    elif pd.isna(conc) or conc < THRESHOLD:
        status = "Skip"
    else:
        status = "To be sequenced"
    rows.append({
        "Sample":       f"{num}{tissue}_G2",
        "Patient":      num,
        "Tissue":       tissue,
        "Conc (ng/ul)": conc,
        "Status":       status,
    })

df = pd.DataFrame(rows).sort_values(["Tissue", "Patient"])

# ── write Excel ───────────────────────────────────────────────────────────────
df.to_excel(OUT_FILE, index=False)

wb = load_workbook(OUT_FILE)
ws = wb.active

# auto-width columns
for col_cells in ws.columns:
    width = max(len(str(c.value)) if c.value is not None else 0 for c in col_cells)
    ws.column_dimensions[col_cells[0].column_letter].width = width + 4

# colour every row based on Status (Status is the last column)
status_col = df.columns.get_loc("Status") + 1   # 1-indexed
for row in ws.iter_rows(min_row=2):
    status = row[status_col - 1].value
    fill = FILLS.get(status)
    if fill:
        for cell in row:
            cell.fill = fill

wb.save(OUT_FILE)
print(f"Saved {OUT_FILE}  ({len(df)} samples)")
print(df["Status"].value_counts().to_string())