from main import set_publication_params, configure_dataset, fit_SCGLUE
import models
import matplotlib.pyplot as plt
import scanpy as sc
import pandas as pd
from matplotlib import rcParams
import anndata as ad
from pathlib import Path

import networkx as nx
from itertools import chain
set_publication_params()
rcParams["figure.figsize"] = (4, 4)

rna = ad.read_h5ad("data/rna-pp.h5ad")
atac = ad.read_h5ad("data/atac-pp.h5ad")
guidance = nx.read_graphml("data/guidance.graphml.gz")
print("successfully read the data")
configure_dataset(rna, "NB", use_highly_variable=True, use_layer="counts", use_rep="X_pca")
configure_dataset(atac, "NB", use_highly_variable=True, use_rep="X_lsi")
print("successfully configure dataset")
guidance_hvf = guidance.subgraph(chain(rna.var.query("highly_variable").index, atac.var.query("highly_variable").index)).copy()

glue = fit_SCGLUE({"rna": rna, "atac": atac}, guidance_hvf, model=models.PairedSCGLUEModel, fit_kws={"directory": "paired"})
rna.obsm["X_glue"] = glue.encode_data("rna", rna)
atac.obsm["X_glue"] = glue.encode_data("atac", atac)

rna.write_h5ad("rna-emb.h5ad")
atac.write_h5ad("atac-emb.h5ad")
nx.write_gml(guidance_hvf, "guidance-hvf.graphml.gz")

combined = ad.concat([rna, atac])
sc.pp.neighbors(combined, use_rep="X_glue", metric="cosine")
sc.tl.umap(combined)
sc.pl.umap(combined, color=["cell_type", "domain"], wspace=0.65)
plt.savefig("cluster.png")
graph = guidance_hvf
feature_embeddings = glue.encode_graph(graph)
feature_embeddings = pd.DataFrame(feature_embeddings, index=glue.vertices)
print(feature_embeddings.iloc[:5, :5])