# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/package/genotype.ipynb (unless otherwise specified).

__all__ = ['get_geno_one_snp', 'index_search', 'BgenFileObject', 'get_genetic_file_path', 'load_genetic_file',
           'parameters']

# Cell
from .data_catalog import get_catalog, get_config, get_parameters
from pydantic import BaseModel
import dask.dataframe as dd
from dask.delayed import Delayed
from bgen_reader import open_bgen, read_bgen
import numpy as np
import pandas as pd
from typing import List, Union
from fastcore.meta import delegates
import logging

# Cell

def get_geno_one_snp(row,  high_lim=0.9, low_lim=0.3, NA_val = np.nan):
    geno_1, geno_2, geno_3 = row
    homo_ref_cond = (geno_1 >=high_lim) & (geno_2 < low_lim) & (geno_3 < low_lim)
    het_cond = (geno_2 >= high_lim) & (geno_1 < low_lim) & (geno_3 < low_lim)
    homo_alt_cond = (geno_3 >= high_lim) & (geno_1 < low_lim) & (geno_2 < low_lim)
    geno_df =  np.select([homo_ref_cond, het_cond, homo_alt_cond],
             [0., 1., 2.],
             default = NA_val)
    return geno_df

def index_search(list_to_search, query_list):
    sorter = np.argsort(list_to_search)
    index = sorter[np.searchsorted(list_to_search, query_list, sorter=sorter)]
    return index

# Cell
class BgenFileObject():
#     variants: dd.DataFrame
#     samples: pd.Series
#     genotype: List[Delayed]
    bgen_reader_obj: open_bgen

    def __init__(self, bgen_reader_obj):
        self.bgen_reader_obj = bgen_reader_obj
        self.samples = np.vstack(np.char.split(self.bgen_reader_obj.samples, sep = " "))[:,0]

    @property
    def ids(self):
        return self.bgen_reader_obj.ids

    def __repr__(self):
        return str(self.__class__) + f" {self.samples.shape[0]} samples"

    def get_variant_index(self,ids=None):
        variant_index = index_search(self.ids, ids) if ids is not None else None
        return variant_index

    def get_sample_index(self, sample_ids=None):
        sample_index = index_search(self.samples, sample_ids) if sample_ids is not None else None
        return sample_index

    def get_probs(self, sample_ids=None, variant_ids=None):
        variant_index = self.get_variant_index(variant_ids)
        sample_index = self.get_sample_index(sample_ids)

        return self.bgen_reader_obj.read((sample_index, variant_index))

    @delegates(get_probs)
    def get_geno_each_sample(self,prob_to_geno_func:Union["max", "stringent"]= "stringent", high_lim=0.9, low_lim=0.3, NA_val=np.nan, one_hot_encoded=True, **kwargs):
        probs = self.get_probs(**kwargs)
        if prob_to_geno_func == "max":
            geno = np.nanargmax(probs, axis=2).astype(float)

        elif prob_to_geno_func == "stringent":
            geno = np.apply_along_axis(get_geno_one_snp, axis=2, arr=probs, high_lim=high_lim, low_lim=low_lim, NA_val=NA_val)


        geno_no_nan = np.nan_to_num(geno, nan=3).astype(int)

        if one_hot_encoded:
            geno = np.identity(4)[geno_no_nan]
        return geno


    def get_allele_ids(self, rsids = None, variant_index = None):
        if variant_index is None:
            variant_index = self.get_variant_index(rsids)
        df = pd.DataFrame([allele_str.split(",") for allele_str in self.bgen_reader_obj.allele_ids[variant_index]], columns = ["allele_1", "allele_2"])

        if rsids is not None:
            df.index = rsids
        return df

    def get_variant_combinations(self, rsids = None, variant_index = None):
        if variant_index is None:
            variant_index = np.argwhere(np.isin(self.bgen_reader_obj.rsids, rsids)).reshape(-1,) if rsids is not None else None
        geno_df = self.get_allele_ids(rsids, variant_index)
        geno_df = get_possible_geno_combinations(geno_df, "allele_1", "allele_2")
        return geno_df

# Cell
parameters = get_parameters()

def get_genetic_file_path(chrom):
    bgen_path = parameters["genetic_file_path_template"].format(chrom_number=chrom)
    sample_path = parameters["sample_file_template"].format(chrom_number=chrom)
    return (bgen_path,sample_path)

def load_genetic_file(chrom):
    bgen_path, sample_path = get_genetic_file_path(chrom)
    return BgenFileObject(bgen_reader_obj = open_bgen(filepath=bgen_path, samples_filepath=sample_path))