# /// script
# requires-python = ">=3.10"
# dependencies = ["matplotlib", "requests"]
# ///

import csv
import io
import json
from math import log1p
from pathlib import Path

import matplotlib.pyplot as plt
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry


def main(base: str) -> None:
    session = requests.Session()
    session.mount("https://", HTTPAdapter(max_retries=Retry(total=3, backoff_factor=1)))

    def download(name: str) -> requests.Response:
        response = session.get(f"{base}/{name}", timeout=60)
        response.raise_for_status()
        return response

    programs = download("programs.json").json()
    scores = {}
    for population, panel in programs.items():
        scores[population] = {}
        for kind, prefix, key in [
            ("references", "", "sample"),
            ("patient", "patient-", "preparation"),
        ]:
            rows = csv.DictReader(
                io.StringIO(download(f"{prefix}{population}-gene-means.csv").text)
            )
            scores[population][kind] = [
                {
                    "sample": row[key],
                    "scores": {
                        name: sum(log1p(float(row[g])) for g in genes) / len(genes)
                        for name, genes in panel.items()
                    },
                }
                for row in rows
            ]
    Path("plot-scores.json").write_text(json.dumps(scores, indent=2))

    plots = [
        ("macrophage-c1q", "macrophages", "C1Q TAM", "Inhibitory signaling"),
        ("macrophage-ifn", "macrophages", "IFN response", "MHC-II presentation"),
        ("t-cells", "t_cells", "Cytotoxicity", "Exhaustion-associated"),
        ("tumor-candidates", "tumor_candidates", "MHC-I presentation", "IFN response"),
    ]
    plt.style.use("dark_background")
    for filename, population, x, y in plots:
        reference = scores[population]["references"]
        patient = [r for r in scores[population]["patient"] if r["sample"].endswith(" unsorted")]
        points = reference + patient
        max_x = max(p["scores"][x] for p in points) * 1.16
        max_y = max(p["scores"][y] for p in points) * 1.16
        fig, ax = plt.subplots(figsize=(6.4, 3.9))
        ax.scatter(
            [r["scores"][x] for r in reference],
            [r["scores"][y] for r in reference],
            c="#b5b7c6",
            alpha=0.7,
            s=20,
        )
        ax.scatter(
            [r["scores"][x] for r in patient], [r["scores"][y] for r in patient], c="#ff6b6b", s=49
        )
        for row in patient:
            ax.annotate(
                row["sample"][:2],
                (row["scores"][x], row["scores"][y]),
                xytext=(7, -12 if row["sample"].startswith("T3") else 7),
                textcoords="offset points",
                color="#ff8d8d",
            )
        ax.set(
            xlim=(0, max_x),
            ylim=(0, max_y),
            xlabel=f"{x}\nLog counts/10k",
            ylabel=f"{y}\nLog counts/10k",
        )
        ticks = [0, 0.25, 0.5, 0.75, 1]
        ax.set_xticks([max_x * t for t in ticks], labels=[f"{max_x * t:.1f}" for t in ticks])
        ax.set_yticks([max_y * t for t in ticks], labels=[f"{max_y * t:.1f}" for t in ticks])
        ax.spines[["top", "right", "bottom", "left"]].set_visible(False)
        ax.grid(alpha=0.04)
        fig.tight_layout()
        fig.savefig(f"{filename}.png", dpi=150)
        plt.close(fig)
        print(f"{filename}.png: {len(reference)} references, {len(patient)} patient samples")


if __name__ == "__main__":
    main("https://glass.bio/blog/osteosarcoma-reference")