# Load libraries
import pandas as pd
import scanpy as sc
import numpy as np
import scanpy.external as sce



class MakeAnndata:
    # This class is used to make the anndata with spatial coordinates from CURIO.
    def __init__(self, sample_path, barcodes_path):
        # sample_path = the path to the h5ad file generated by CURIO pipeline.
        # barcodes_path = the path to the bead barcode whitelist corresponding to sample.
        self.sample = sc.read_h5ad(sample_path)
        self.barcodes = pd.read_csv(barcodes_path, delimiter="\t", names=["barcode", "x", "y"])
        # the data will be the finished anndata.
        self.data = self._shared_barcodes()

    def _process_whitelist(self):
        # Function to preprocess bead barcode whitelist.
        print(self.barcodes[["x", "y"]].to_numpy())
        barcodes = self.barcodes
        barcodes.index = barcodes["barcode"]
        return barcodes

    def _process_h5ad(self):
        # Function to preprocess h5ad file from CURIO.
        sc.pp.filter_cells(self.sample, min_counts=1)
        data = self.sample
        data.obs_names = [(''.join(x.split("_"))) for x in data.obs_names]
        return data

    def _shared_barcodes(self):
        # Function to generate finale anndata file with spatial positions.
        barcodes = self._process_whitelist()
        data = self._process_h5ad()

        shared_barcodes = list(set(data.obs_names).intersection(set(barcodes["barcode"])))
        data = data[shared_barcodes]
        barcodes = barcodes.reindex(data.obs_names)
        data.obs = data.obs.join(barcodes)
        data.obsm["spatial"] = barcodes[["x", "y"]].to_numpy()
        return data