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

__all__ = ['get_info_dict', 'catalog', 'simulation_I83_queries_pheno_dict', 'snp_filter', 'DataSplitParams',
           'chromosome_datasource', 'default_data_split_params']

# Cell
from .simulation import SimulatedPheno, SNPInfoUnit
from .genotype import load_genetic_file
from .data_catalog import get_catalog, get_config, get_parameters
from .phenotypes import get_phenotype, get_GWAS_snps_for_trait, upsample_pheno

from typing import List, Union, Literal
import numpy as np
import logging
from tqdm.auto import tqdm
import pandas as pd
from sklearn.model_selection import train_test_split
from dataclasses import dataclass

# Cell


def get_info_dict(data, balance_pheno, name="test"):
    return_dict = {f"{name}_data": data.shape[0],
                            f"{name}_data_n_controls": data[balance_pheno].value_counts().to_dict()[0],
                            f"{name}_data_n_cases": data[balance_pheno].value_counts().to_dict()[1],
                           f"{name}_data_frac_controls": data[balance_pheno].value_counts(normalize=True).to_dict()[0],
                           f"{name}_data_frac_cases": data[balance_pheno].value_counts(normalize=True).to_dict()[1]
                           }
    return return_dict

# Cell
catalog = get_catalog()
simulation_I83_queries_pheno_dict = catalog.load('simulation_I83_queries_pheno_dict')


# Cell

class snp_filter:
    SORT_PVALUE = 'pval'
    SORT_BETA = 'beta'
    def __init__(self, phenotype:str, sort:str, threshold=1e-5):
        self.phenotype = phenotype
        self.sort = sort
        self.threshold = threshold

@dataclass
class DataSplitParams:
    validation_split: float
    test_split: float
    tune_split: float
    validation_tune_split: float

default_data_split_params = DataSplitParams(validation_split=0.2, test_split=0.1, tune_split=0.01, validation_tune_split = 0.2)

class chromosome_datasource:
    def __init__(self, chromosomes:List[int] =list(range(1, 23)), snp_filters:List[snp_filter] =[snp_filter('I84', snp_filter.SORT_PVALUE)], samples: np.array =None, max_samples:int = 100_000, random_state:int =42, balance_pheno:str="I84", data_split_params:DataSplitParams = default_data_split_params):
        self.genome_files = {chromosome: load_genetic_file(chromosome) for chromosome in tqdm(chromosomes, "Loading genotype file(s)")}
        self.chromosomes = chromosomes
        self.phenotypes = [snp.phenotype for snp in snp_filters]
        self.data_split_params = data_split_params
        self.balance_pheno = balance_pheno
        self.max_samples = max_samples
        self.random_state = random_state

        if samples is None:
            self.samples = self.genome_files[chromosomes[0]].samples

        logging.warning("creating phenotype matrix")

        self.pheno_df_ordered = get_phenotype(self. phenotypes, samples = self.samples, max_samples = None)
        self.data_dict, self.data_split_info_df = self.split_data(self.pheno_df_ordered, self.balance_pheno, self.data_split_params, self.random_state, train_upsampled_max_samples = self.max_samples )
        logging.warning("finished creating phenotype matrix")

        self.all_snps_dict = {}
        self.snp_filters = snp_filters

        logging.warning(f"Loading SNPs list for {len(chromosomes)} chromosomes: {chromosomes=}")
        self.snp_full_info, self.snp_trait_info = self.load_snps_for_traits_all_chroms()
        logging.warning(f"Found {self.snp_full_info.shape[0]} SNPs associated with traits {self. phenotypes} for {chromosomes=}")


    def split_data(self, pheno_df, balance_pheno, data_split_params_obj, random_state, train_upsampled_max_samples):
        train_original, test = train_test_split( pheno_df, test_size=data_split_params_obj.test_split, stratify = pheno_df[balance_pheno], random_state = random_state)
        train_remaining, validation = train_test_split(train_original, test_size=data_split_params_obj.validation_split, stratify = train_original[balance_pheno], random_state = random_state)
        train_before_upsample, tune = train_test_split(train_remaining, test_size=data_split_params_obj.tune_split, stratify = train_remaining[balance_pheno], random_state = random_state)
        train_tune, validation_tune = train_test_split(tune, test_size=data_split_params_obj.validation_tune_split, stratify = tune[balance_pheno], random_state = random_state)
        if train_upsampled_max_samples == None:
            train_final = train_before_upsample
        else:
            train_tune = upsample_pheno(pheno_df = train_tune, balance_pheno = balance_pheno, max_samples = int(train_upsampled_max_samples * ((1 - data_split_params_obj.validation_tune_split)/4)), random_state = random_state)
            train_final = upsample_pheno(pheno_df = train_before_upsample, balance_pheno = balance_pheno, max_samples = train_upsampled_max_samples, random_state = random_state)

        data_dict = {"train_original": train_original,
                    "test": test,
                    "train_before_upsample": train_before_upsample,
                    "validation": validation,
                    "tune": tune,
                    "validation_tune": validation_tune,
                    "train_tune": train_tune,
                    "train": train_final}

        info_dict = {}
        for name, data in data_dict.items():
            info_dict = {**info_dict, **get_info_dict(data,balance_pheno = balance_pheno, name=name )}
        info_df = pd.DataFrame.from_dict(info_dict, orient="index", columns = [balance_pheno])

        return data_dict, info_df

    def load_snps_for_traits_all_chroms(self):
        all_SNPs_df = []
        pbar = tqdm(self.snp_filters)
        for snp in pbar:
            pbar.set_description(f"Searching GWAS result for trait: {snp.phenotype}")
            trait_snps_df = get_GWAS_snps_for_trait(snp.phenotype, chromosome=self.chromosomes, id_only=False, sort_val_cols_list=snp.sort, ascending_bool_list=[snp.sort == snp_filter.SORT_PVALUE]).query(f'{snp.sort} {"<" if snp.sort == snp_filter.SORT_PVALUE else ">"} {snp.threshold}')#.sort_values('position')['full_id'].values
            trait_snps_df[["chr", "position"]] = trait_snps_df[["chr", "position"]].astype(int)
            trait_snps_df["trait"] =  snp.phenotype
            trait_snps_df["sort_cols"] = snp.sort
            trait_snps_df["threshold"] = snp.threshold
            all_SNPs_df.append(trait_snps_df)

        snp_full_info = pd.concat(all_SNPs_df).sort_values(["chr", "position"])
        snp_trait_info = snp_full_info.groupby("full_id")["trait"].agg(["unique", "nunique"])
        return snp_full_info, snp_trait_info

    def get_geno_matrix_specific_chrom(self, sample_id_subset, chrom):
        chrom_specific_variant_ids = self.snp_full_info.drop_duplicates("full_id").query(f"chr == {chrom}").sort_values("position")["full_id"].values
        logging.warning(f"Loading {len(chrom_specific_variant_ids)} SNPs for chromosome {chrom}, {sample_id_subset.shape[0]} people")
        genos = self.genome_files[chrom].get_geno_each_sample(prob_to_geno_func = "max", sample_ids=sample_id_subset, variant_ids=chrom_specific_variant_ids)
        return genos

    def get_sample_id_in_split(self, sample_slice, split:Literal["train", "validation", "train_tune", "validation_tune", "test"]):
        split_df = self.data_dict.get(split)
        split_df_subset = split_df.iloc[sample_slice, :]
        sample_id_subset = split_df_subset.index.values
        return split_df_subset, sample_id_subset

    def get_X(self, sample_id_subset, pbar):
        all_genos_all_chrom = []
        for chrom in pbar:
            genos = self.get_geno_matrix_specific_chrom(sample_id_subset, chrom)
            all_genos_all_chrom.append(genos)
        X = np.hstack(all_genos_all_chrom)
        return X

    def get_data(self, sample_slice: slice, split:Literal["train", "validation", "train_tune", "validation_tune", "test"]):
        import multiprocessing as mp

        split_df_subset, sample_id_subset = self.get_sample_id_in_split(sample_slice, split)
        pbar = tqdm(self.snp_full_info.chr.unique(), f"Loading genotype data from {split=}")
        X = self.get_X(sample_id_subset, pbar)
        y = split_df_subset.values.reshape(-1)
        return X, y


    def get_simulated_data(self, sample_slice: slice, split:Literal["train", "validation", "train_tune", "validation_tune", "test"], snp_pair: List):
        split_df_subset, sample_id_subset = self.get_sample_id_in_split(sample_slice, split)
        pbar = tqdm(self.snp_full_info.chr.unique(), f"Loading genotype data from {split=}")
        X = self.get_X(sample_id_subset, pbar)
        try:
            y = simulation_I83_queries_pheno_dict[(snp_pair[0], snp_pair[1])]
        except KeyError:
            y = simulation_I83_queries_pheno_dict[(snp_pair[1], snp_pair[0])]
        y = {k: v.pheno_col.loc[sample_id_subset] for k, v in y.items()}

        return X, y