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

__all__ = ['QueryDataframe', 'parameters', 'catalog_all', 'catalog_all', 'read_csv_compressed', 'get_GWAS_result_link',
           'heritability_Neale', 'display_cols', 'quality_heritability_phenos', 'quality_heritability_phenos',
           'icd10_pheno_matrix', 'icd10_primary_cols', 'icd10_pheno_matrix', 'upsample_pheno', 'get_phenotype',
           'get_GWAS_snps_for_trait']

# Cell
from ..data_catalog import get_catalog, get_parameters
import combinatorial_gwas
from pathlib import Path
import pandas as pd
from dataclasses import dataclass
from functools import partial
import numpy as np
from typing import List, Union
from fastcore.utils import partialler
import logging

# Cell

@pd.api.extensions.register_dataframe_accessor("pheno")
@dataclass
class QueryDataframe():
    df: pd.DataFrame

    def query(self, **column_dict:dict):
        query_str = " and ".join([f"({col} {cond})" for col, cond in column_dict.items()])
        return self.df.query(query_str)

# Cell
parameters = get_parameters()
parameters

# Cell
catalog_all = get_catalog()
catalog_all = catalog_all.reload()
catalog_all.list()

# Cell
read_csv_compressed= partialler(pd.read_csv, sep="\t", compression= "gzip")
get_GWAS_result_link = partialler(parameters['template_gwas_result_file_link'].format)


# Cell
heritability_Neale = catalog_all.load("heritability_trait_level_summary")
heritability_Neale.head()

# Cell
display_cols = ['description', 'h2_liability', 'h2_sig', 'confidence', 'n_cases', 'n_controls', 'prevalence']

# Cell
quality_heritability_phenos = heritability_Neale.pheno.query(h2_sig = "in ['z7', 'z4']", source= " == 'icd10'", confidence= "in ['medium', 'high']").sort_values("h2_liability", ascending = False)
quality_heritability_phenos = quality_heritability_phenos.set_index("phenotype")
quality_heritability_phenos.head()[display_cols]

# Cell
logging.warning("Loading ICD phenotype matrix, this might take a while")
icd10_pheno_matrix = catalog_all.load("ICD10_pheno_matrix")


#get the first 3 character of ICD code
icd10_primary_cols = icd10_pheno_matrix.columns[icd10_pheno_matrix.columns.str.contains("primary")]
icd10_pheno_matrix = icd10_pheno_matrix.astype(str).apply(lambda x: x.str.slice(0,3))

logging.warning("Finished loading ICD10 matrix")

# Cell

def upsample_pheno(pheno_df, balance_pheno, max_samples, random_state):
    weights = pheno_df[balance_pheno].replace(pheno_df[balance_pheno].value_counts(normalize=True).to_dict())
    pheno_df_upsampled = pheno_df.sample(max_samples, replace=True, weights = 1/weights, random_state=random_state)
    return pheno_df_upsampled


def get_phenotype(icd10_codes: Union[str, List[str]] ="I84", samples:np.array=None, max_samples:int = None, balance_pheno: str = None, random_state=42):
    """
    if samples argument is provided from genetic file, then find common set of samples and output ordered phenotype

    if `max_samples` is provided, then over-sample the data so that we have `max_samples/2` for both cases and controls
    """
    icd10_codes = [icd10_codes] if not isinstance(icd10_codes, list) else icd10_codes
    pheno_df_list = [icd10_pheno_matrix[icd10_primary_cols].isin([icd10_code]).any(axis=1).astype(int) for icd10_code in icd10_codes]
    pheno_df = pd.concat(pheno_df_list, axis=1)
    pheno_df.columns = icd10_codes

    if samples is not None:
        geno_pheno_sample_index_mask = np.isin(samples.astype(int), pheno_df.index)
        pheno_geno_samples_common_set = samples[geno_pheno_sample_index_mask].astype(int)
        pheno_df_ordered = pheno_df.loc[list(pheno_geno_samples_common_set), :]
        pheno_df_ordered = pheno_df_ordered.loc[~pheno_df_ordered.index.duplicated(keep="first"),:]
        sample_index = np.argwhere(geno_pheno_sample_index_mask).reshape(-1)

        if max_samples is not None:
            if balance_pheno is None:
                raise ValueError("Need to specify `balance_pheno` param: phenotype to balance during subsampling")
            pheno_df_ordered = upsample_pheno(pheno_df_ordered, balance_pheno, max_samples, random_state)

            sorter = np.argsort(samples.astype(int))
            sample_index = sorter[np.searchsorted(samples.astype(int), pheno_df_ordered.index, sorter=sorter)]

        assert np.allclose(samples[sample_index].astype(int), pheno_df_ordered.index), "sample mismatch between genotype and phenotype, something wrong with the `get_phenotype` function!"

        pheno_df_ordered.index = pheno_df_ordered.index.astype(str)
        return pheno_df_ordered
    return pheno_df

# Cell

def get_GWAS_snps_for_trait(phenotype_code= "I84", chromosome:Union[int, List[int]] = 21, sort_val_cols_list: List[str] = ["pval"], ascending_bool_list: List[bool] = [True], id_only= True):
    chromosome = [chromosome] if not isinstance(chromosome, list) else chromosome
    chromosome_str = [f"{single_chromosome}:" for single_chromosome in chromosome]
    gwas_result_df = read_csv_compressed(get_GWAS_result_link(phenotype_code=phenotype_code))
    gwas_result_df = gwas_result_df.loc[gwas_result_df["variant"].str.startswith(tuple(chromosome_str)), :]
    gwas_result_df = gwas_result_df.reset_index(drop=True).reset_index().rename(columns = {"index":"position_rank"})
    gwas_result_df = gwas_result_df.sort_values(sort_val_cols_list, ascending = ascending_bool_list)
    variant_id_df = gwas_result_df["variant"].str.split(":",expand=True)
    variant_id_df["chr1_4"] =variant_id_df[[1,2,3]].apply("_".join, axis=1)
    variant_id_df[1] = variant_id_df[1].astype(int)
    gwas_result_df[["chr", "position", "major_allele"]] = variant_id_df[[0, 1, 2]]
    gwas_result_df["full_id"] =  variant_id_df[[0, "chr1_4"]].apply(":".join, axis=1)

    if id_only:
        return gwas_result_df["full_id"].values
    else:
        return gwas_result_df
