import io

import numpy as np

import anndata as anndata

import scipy as s

import scanpy as sc

import matplotlib.pyplot as plt

import seaborn

import pandas as pd

import csv

import os

from dfply import *

# from IPython.display import display

from re import search

from pathlib import Path

import argparse

#########

## I/O ##

#########

io = {}

host = os.uname()[1]

if search("ricard", host):

    io["basedir"] = "/Users/ricard/data/gastrulation_multiome_10x/test"

    io["cellranger_output"] = "/Users/ricard/data/gastrulation_multiome_10x/original"

    io["gene_metadata"] = "/Users/ricard/data/ensembl/mouse/v87/BioMart/all_genes/Mmusculus_genes_BioMart.87.txt"

    io["gene_metadata_gtf"] = "/Users/ricard/data/ensembl/mouse/v98/GTF/Mus_musculus.GRCm38.98.gtf.gz"

elif search("Workstation", host):

    io["basedir"]="/home/lijingyu/gastrulation/data/gastrulation_multiome_10x"                   

    io["gene_metadata"]='/home/lijingyu/gastrulation/data/ensembl/mouse/v87/BioMart/mRNA/Mmusculus_genes_BioMart.87.txt'

elif search("BI2404M", host):

    io["basedir"] = "/Users/argelagr/data/gastrulation_multiome_10x/test"

    io["gene_metadata"] = "/Users/argelagr/data/ensembl/mouse/v87/BioMart/all_genes/Mmusculus_genes_BioMart.87.txt"

elif search("pebble|headstone", host):

    io["basedir"] = "/bi/group/reik/ricard/data/gastrulation_multiome_10x/test"

    io["cellranger_output"] = "/bi/group/reik/ricard/data/gastrulation_multiome_10x/original"

    io["gene_metadata"] = "/bi/group/reik/ricard/data/ensembl/mouse/v87/BioMart/all_genes/Mmusculus_genes_BioMart.87.txt"

    io["gene_metadata_gtf"] = "/bi/group/reik/ricard/data/ensembl/mouse/v98/GTF/Mus_musculus.GRCm38.98.gtf.gz"

else:

    print("Computer not recognised"); exit()

io["metadata"] = io["basedir"] + "/sample_metadata.txt.gz"

io["anndata"] = io["basedir"] + "/processed/rna/anndata.h5ad"

io["anndata_scvelo"] = io["basedir"] + "/processed/rna/velocyto/anndata_scvelo.h5ad"

#############

## Options ##

#############

opts = {}

opts["samples"] = [

    "E7.5_rep1",

    "E7.5_rep2",

    "E7.75_rep1",

    "E8.0_rep1",

    "E8.0_rep2",

    "E8.5_rep1",

    "E8.5_rep2",

    "E8.75_rep1",

    "E8.75_rep2",

    "E8.5_CRISPR_T_KO",

    "E8.5_CRISPR_T_WT"

]

opts["stages"] = [

    "E7.5",

    "E7.75",

    "E8.0",

    "E8.5",

    "E8.75"

]

opts["celltypes"] = [

    "Epiblast",

    "Primitive_Streak",

    "Caudal_epiblast",

    "PGC",

    "Anterior_Primitive_Streak",

    "Notochord",

    "Def._endoderm",

    "Gut",

    "Nascent_mesoderm",

    "Mixed_mesoderm",

    "Intermediate_mesoderm",

    "Caudal_Mesoderm",

    "Paraxial_mesoderm",

    "Somitic_mesoderm",

    "Pharyngeal_mesoderm",

    "Cardiomyocytes",

    "Allantois",

    "ExE_mesoderm",

    "Mesenchyme",

    "Haematoendothelial_progenitors",

    "Endothelium",

    "Blood_progenitors_1",

    "Blood_progenitors_2",

    "Erythroid1",

    "Erythroid2",

    "Erythroid3",

    "NMP",

    "Rostral_neurectoderm",

    "Caudal_neurectoderm",

    "Neural_crest",

    "Forebrain_Midbrain_Hindbrain",

    "Spinal_cord",

    "Surface_ectoderm",

    "Visceral_endoderm",

    "ExE_endoderm",

    "ExE_ectoderm",

    "Parietal_endoderm"

]

opts["celltype_colors"] = {

    "Epiblast" : "#635547",

    "Primitive_Streak" : "#DABE99",

    "Caudal_epiblast" : "#9e6762",

    "PGC" : "#FACB12",

    "Anterior_Primitive_Streak" : "#c19f70",

    "Notochord" : "#0F4A9C",

    "Def._endoderm" : "#F397C0",

    "Gut" : "#EF5A9D",

    "Nascent_mesoderm" : "#C594BF",

    "Mixed_mesoderm" : "#DFCDE4",

    "Intermediate_mesoderm" : "#139992",

    "Caudal_Mesoderm" : "#3F84AA",

    "Paraxial_mesoderm" : "#8DB5CE",

    "Somitic_mesoderm" : "#005579",

    "Pharyngeal_mesoderm" : "#C9EBFB",

    "Cardiomyocytes" : "#B51D8D",

    "Allantois" : "#532C8A",

    "ExE_mesoderm" : "#8870ad",

    "Mesenchyme" : "#cc7818",

    "Haematoendothelial_progenitors" : "#FBBE92",

    "Endothelium" : "#ff891c",

    "Blood_progenitors" : "#c9a997",

    "Blood_progenitors_1" : "#f9decf",

    "Blood_progenitors_2" : "#c9a997",

    "Erythroid" : "#EF4E22",

    "Erythroid1" : "#C72228",

    "Erythroid2" : "#f79083",

    "Erythroid3" : "#EF4E22",

    "NMP" : "#8EC792",

    "Neurectoderm" : "#65A83E",

    "Rostral_neurectoderm" : "#65A83E",

    "Caudal_neurectoderm" : "#354E23",

    "Neural_crest" : "#C3C388",

    "Forebrain_Midbrain_Hindbrain" : "#647a4f",

    "Spinal_cord" : "#CDE088",

    "Surface_ectoderm" : "#f7f79e",

    "Visceral_endoderm" : "#F6BFCB",

    "ExE_endoderm" : "#7F6874",

    "ExE_ectoderm" : "#989898",

    "Parietal_endoderm" : "#1A1A1A"

}

# opts["stages_colors"] = {

#   'E6.5':"#D53E4F",

#   'E6.75':"#F46D43",

#   'E7.0':"#FDAE61",

#   'E7.5':"#FFFFBF",

#   'E7.25':"#FEE08B",

#   'E7.75':"#E6F598",

#   'E8.0':"#ABDDA4",

#   'E8.5':"#3288BD",

#   'E8.25':"#66C2A5",

#   'mixed_gastrulation': "#A9A9A9"  

# }

opts["stage_colors"] = {

    'E7.5':"#440154FF",

    'E7.75':"#E6F598",

    'E8.0':"#31688EFF",

    'E8.5':"#35B779FF",

    'E8.75':"#FDE725FF"

}