download_path = "./norman2019/"

download_norman_2019(download_path)

matrix = read_norman_2019(download_path)

# List of cell barcodes. The barcodes in this list are stored in the same order
# as cells are in the count matrix.
cell_barcodes = pd.read_csv(
    os.path.join(download_path, "GSE133344_filtered_barcodes.tsv.gz"),
    sep="\t",
    header=None,
    names=["cell_barcode"],
)

# IDs/names of the gene features.
gene_list = pd.read_csv(
    os.path.join(download_path, "GSE133344_filtered_genes.tsv.gz"),
    sep="\t",
    header=None,
    names=["gene_id", "gene_name"],
)

# Dataframe where each row corresponds to a cell, and each column corresponds
# to a gene feature.
matrix = pd.DataFrame(
    matrix.transpose().todense(),
    columns=gene_list["gene_id"],
    index=cell_barcodes["cell_barcode"],
    dtype="int32",
)

# Dataframe mapping cell barcodes to metadata about that cell (e.g. which CRISPR
# guides were applied to that cell). Unfortunately, this list has a different
# ordering from the count matrix, so we have to be careful combining the metadata
# and count data.
cell_identities = pd.read_csv(
    os.path.join(download_path, "GSE133344_filtered_cell_identities.csv.gz")
).set_index("cell_barcode")

# This merge call reorders our metadata dataframe to match the ordering in the
# count matrix. Some cells in `cell_barcodes` do not have metadata associated with
# them, and their metadata values will be filled in as NaN.
aligned_metadata = pd.merge(
    cell_barcodes,
    cell_identities,
    left_on="cell_barcode",
    right_index=True,
    how="left",
).set_index("cell_barcode")

adata = AnnData(matrix)
adata.obs = aligned_metadata

# Filter out any cells that don't have metadata values.
rows_without_nans = [
    index for index, row in adata.obs.iterrows() if not row.isnull().any()
]
adata = adata[rows_without_nans, :]

# Remove these as suggested by the authors. See lines referring to
# NegCtrl1_NegCtrl0 in GI_generate_populations.ipynb in the Norman 2019 paper's
# Github repo https://github.com/thomasmaxwellnorman/Perturbseq_GI/
adata = adata[adata.obs["guide_identity"] != "NegCtrl1_NegCtrl0__NegCtrl1_NegCtrl0"]

# We create a new metadata column with cleaner representations of CRISPR guide
# identities. The original format is <Guide1>_<Guide2>__<Guide1>_<Guide2>_<number>
adata.obs["guide_merged"] = adata.obs["guide_identity"]

control_regex = re.compile(r"NegCtrl(.*)_NegCtrl(.*)+NegCtrl(.*)_NegCtrl(.*)")
for i in adata.obs["guide_merged"].unique():
    if control_regex.match(i):
        # For any cells that only had control guides, we don't care about the
        # specific IDs of the guides. Here we relabel them just as "ctrl".
        adata.obs["guide_merged"].replace(i, "ctrl", inplace=True)
    else:
        # Otherwise, we reformat the guide label to be <Guide1>+<Guide2>. If Guide1
        # or Guide2 was a control, we replace it with "ctrl".
        split = i.split("__")[0]
        split = split.split("_")
        for j, string in enumerate(split):
            if "NegCtrl" in split[j]:
                split[j] = "ctrl"
        adata.obs["guide_merged"].replace(i, f"{split[0]}+{split[1]}", inplace=True)

guides_to_programs = {}
guides_to_programs.update(dict.fromkeys(G1_CYCLE, "G1 cell cycle arrest"))
guides_to_programs.update(dict.fromkeys(ERYTHROID, "Erythroid"))
guides_to_programs.update(dict.fromkeys(PIONEER_FACTORS, "Pioneer factors"))
guides_to_programs.update(
    dict.fromkeys(GRANULOCYTE_APOPTOSIS, "Granulocyte/apoptosis")
)
guides_to_programs.update(dict.fromkeys(PRO_GROWTH, "Pro-growth"))
guides_to_programs.update(dict.fromkeys(MEGAKARYOCYTE, "Megakaryocyte"))
guides_to_programs.update(dict.fromkeys(["ctrl"], "Ctrl"))

adata.obs["gene_program"] = [guides_to_programs[x] if x in guides_to_programs else "N/A" for x in adata.obs["guide_merged"]]
adata.obs["good_coverage"] = adata.obs["good_coverage"].astype(bool)