from .analysis import (
    phewas,
    comorbidity_strength, 
    binomial_test, 
    comorbidity_network, 
    disease_trajectory)
from .data_management import DiseaseNetworkData
import numpy as np
import os
from .utility import n_process_check,threshold_check,validate_method_specific_kwargs,matching_var_check
from .utility import correction_method_check
import time
import pandas as pd

def disease_network_pipeline(
    data:DiseaseNetworkData, 
    n_process:int, 
    n_threshold_phewas:int, 
    n_threshold_comorbidity:int,
    output_dir:str,
    project_prefix:str,
    keep_positive_associations:bool=False,
    save_intermediate_data:bool=False,
    system_exl:list=None, 
    pipeline_mode:str='v1',
    method:str='RPCN',
    covariates:list=None,
    matching_var_dict:dict={'sex':'exact'},
    matching_n:int=2,
    min_interval_days:int=0,
    max_interval_days:int=np.inf,
    enforce_temporal_order:bool=False,
    correction:str='bonferroni',
    cutoff=0.05,
    **kwargs) -> pd.DataFrame:
    """
    This is a pipeline function to run the whole disease network analysis.

    Parameters
    ----------
    data : DiseaseNetworkData
        The DiseaseNetworkData object.

    n_process : int
        Specifies the number of parallel processes to use for the disease network analysis.
        Multiprocessing is enabled when `n_process` is set to a value greater than one.

    n_threshold_phewas : int
        The minimum number of cases within the exposed group required for a phecode to be included in the PheWAS analysis.
        This parameter will be passed to the phewas function.
        See the phewas function for more information.

    n_threshold_comorbidity : int
        The minimum number of individuals in the exposed group in which a disease pair must co-occur (temporal or non-temporal) to be included in the comorbidity strength estimation.
        This parameter will be passed to the comorbidity_strength function.
        See the comorbidity_strength function for more information.
    
    output_dir : str
        Directory path to store output files generated by the pipeline.

    project_prefix : str
        Prefix for naming output files and intermediate data.

    keep_positive_associations : bool, default=False
        If True, retains only diseases with hazard ratio (HR) > 1 from the PheWAS analysis,
        and disease pairs with positive correlations from the comorbidity strength estimation.

    save_intermediate_data : bool, default=False
        If True, intermediate DiseaseNetworkData objects created by the `DiNetxify.DiseaseNetworkData.disease_pair` function are saved to disk.
    
    system_exl : list, default=None
        List of phecode systems to exclude from the analysis. 
        List of eligible phecode systems: 
            circulatory system; congenital anomalies; dermatologic; digestive; 
            endocrine/metabolic; genitourinary; hematopoietic; infectious diseases; injuries & poisonings; 
            mental disorders; musculoskeletal; neoplasms; neurological; pregnancy complications; 
            respiratory; sense organs; symptoms; others.
        This parameter will be passed to the phewas function.
        See the phewas function for more information.
    
    pipeline_mode : str, default='v1'
        Specifies the analysis order. Available modes:
            - 'v1': PheWAS → comorbidity strength → binomial test → (comorbidity network analysis/disease trajectory analysis)
            - 'v2': PheWAS → comorbidity strength → comorbidity network analysis → binomial test → disease trajectory analysis
        In 'v1', the binomial test does not depend on results from the comorbidity network analysis; thus, disease trajectory and comorbidity network analyses can be conducted independently.
        In 'v2', the binomial test is performed only on disease pairs identified as significant by the comorbidity network analysis, making the disease trajectory analysis dependent on these results.

    method : str, default='RPCN'
        The method to use for the comorbidity network and disease trajectory analysis.
        Available methods are:
            - 'RPCN: Regularized Partial Correlation Network.
            - 'PCN_PCA: Partial Correlation Network with PCA.
            - 'CN': Correlation Network.
        This parameter will be passed to the comorbidity_network and disease_trajectory function.
        See the these two functions for more information.
    
    covariates : list, default=None
        List of covariates to adjust for in the PheWAS, comorbidity network and disease trajectory analysis.
        To include the required variable sex as a covariate, always use 'sex' instead of its original column name.
        For other covariates specified in the 'DiNetxify.DiseaseNetworkData.phenotype_data()' function, use their original column names.
        This parameter will be passed to the phewas, comorbidity_network and disease_trajectory function.
        See these three functions for more information.
    
    matching_var_dict : dict, default={'sex':'exact'}
        Specifies the matching variables and the criteria used for incidence density sampling.
        For categorical and binary variables, the matching criteria should always be 'exact'.
        For continuous variables, provide a scalar greater than 0 as the matching criterion, indicating the maximum allowed difference when matching.
        To include the required variable sex as a matching variable, always use 'sex' instead of its original column name.
        This parameter will be passed to the disease_trajectory function.
        See the disease_trajectory function for more information.
    
    matching_n : int, default=2
        Specifies the maximum number of matched controls for each case.
        This parameter will be passed to the disease_trajectory function.
        See the disease_trajectory function for more information.
    
    min_interval_days : int/float, default=0
        Minimum required time interval (in days) between diagnosis dates when constructing temporal D1 → D2 disease pair for each individual.
        This parameter will be passed to the DiNetxify.DiseaseNetworkData.disease_pair function.
        See the disease_pair function for more information.

    max_interval_days : int/float, default=np.inf
        Maximum allowed time interval (in days) between diagnosis dates when constructing temporal and non-temporal D1-D2 disease pair for each individual.
        This parameter will be passed to the DiNetxify.DiseaseNetworkData.disease_pair function.
        See the disease_pair function for more information.
    
    enforce_temporal_order : bool, default=False
        If set to True, exclude individuals with non-temporal D1-D2 pair when performing the binomial test;
        also applies the specified minimum and maximum time intervals when performing disease trajectory analysis.
        See 'enforce_temporal_order' parameter in binomial_test function and 'enforce_time_interval' parameter in disease_trajectory function.
    
    correction : str, default='bonferroni'
        Method for p-value correction from the statsmodels.stats.multitest.multipletests.
        Available methods are:
        none : no correction
        bonferroni : one-step correction
        sidak : one-step correction
        holm-sidak : step down method using Sidak adjustments
        holm : step-down method using Bonferroni adjustments
        simes-hochberg : step-up method (independent)
        hommel : closed method based on Simes tests (non-negative)
        fdr_bh : Benjamini/Hochberg (non-negative)
        fdr_by : Benjamini/Yekutieli (negative)
        fdr_tsbh : two stage fdr correction (non-negative)
        fdr_tsbky : two stage fdr correction (non-negative)
        See https://www.statsmodels.org/dev/generated/statsmodels.stats.multitest.multipletests.html for more details.
    
    cutoff : float, default=0.05
        The significance threshold for adjusted p-values.
    
    **kwargs
        RPCN Method Parameters:
            alpha : non-negative scalar
                The weight multiplying the l1 penalty term for other diseases covariates. 
                Ignored if 'auto_penalty' is enabled.
            auto_penalty : bool, default=True
                If 'True', automatically determines the best 'alpha' based on model AIC value.
            alpha_range : tuple, default=(1,15)
                When 'auto_penalty' is True, search the optimal 'alpha' in this range.
            scaling_factor : positive scalar, default=1
                The scaling factor for the alpha when 'auto_penalty' is True.

        PCN_PCA Method Parameters:
            n_PC : int, default=5
                Fixed number of principal components to include in each model.
            explained_variance : float
                Cumulative explained variance threshold to determine the number of principal components. 
                Overrides 'n_PC' if specified.
    
    Returns
    -------
    tuple
        A tuple of DataFrames for result of PheWAS, comorbidity strength estimation, comorbidity network analysis, binomial test, and disease trajectory analysis.
    """
    
    # --------check each parameter--------
    #data type check
    if not isinstance(data,DiseaseNetworkData):
        raise TypeError("The input 'data' must be a DiseaseNetworkData object.")
    #attribute check
    data_attrs = ['phenotype_df', 'diagnosis', 'n_diagnosis', 'history']
    for attr in data_attrs:
        if getattr(data, attr) is None:
            raise ValueError(f"Attribute '{attr}' is empty.")
    #n_process check
    n_process, _ = n_process_check(n_process,'PheWAS')
    #check threshold
    n_exposed = data.get_attribute('phenotype_statistics')['n_exposed']
    n_threshold_phewas = threshold_check(None,n_threshold_phewas,n_exposed)
    n_threshold_comorbidity = threshold_check(None,n_threshold_comorbidity,n_exposed)
    if n_threshold_comorbidity > n_threshold_phewas:
        raise ValueError("The 'n_threshold_comorbidity' must be less than or equal to 'n_threshold_phewas'.")
    #output_dir check
    if not isinstance(output_dir,str):
        raise TypeError("The provided input 'output_dir' must be a string.")
    if not os.path.exists(output_dir):
        raise ValueError("The provided output directory does not exist.")
    #project_prefix check
    if not isinstance(project_prefix,str):
        raise TypeError("The provided input 'project_prefix' must be a string.")
    #keep_positive_associations check
    if not isinstance(keep_positive_associations,bool):
        raise TypeError("The provided input 'keep_positive_associations' must be a bool.")
    #save_intermediate_data check
    if not isinstance(save_intermediate_data,bool):
        raise TypeError("The provided input 'save_intermediate_data' must be a bool.")
    #check system_exl
    # phecode_info = data.phecode_info
    # phecode_lst_all = filter_phecodes(phecode_info, None, system_exl, None, None)
    # print(f'A total of {len(phecode_lst_all)} phecodes included in the PheWAS analysis.')
    #pipeline_mode check
    if pipeline_mode not in ['v1','v2']:
        raise ValueError("Invalid pipeline_mode. Must be either 'v1' or 'v2'.")
    #method and kwargs check
    if kwargs:
        parameter_dict = validate_method_specific_kwargs(method,**kwargs)
        if kwargs:
            raise ValueError(f"Unexpected keyword arguments provided: {kwargs.keys()}")
    else:
        parameter_dict = {}
    #check covariates
    # covariates = covariates_check(covariates,data.get_attribute('phenotype_info'))
    #check matching_var_dict
    matching_var_check(matching_var_dict,data.get_attribute('phenotype_info'))
    #check matching_n
    if not isinstance(matching_n,int) or matching_n<1:
        raise ValueError("The input 'matching_n' must be a positive integer.")
    #check min and max interval 
    if not isinstance(min_interval_days, (float,int)):
        raise TypeError("The provided input 'min_interval_days' must be a int")
    if min_interval_days<0:
        raise ValueError("The provided input 'min_interval_days' is not valide.")
    if not isinstance(max_interval_days, (float,int)):
        raise TypeError("The provided input 'max_interval_days' must be a int")
    if max_interval_days<0 or max_interval_days<=min_interval_days:
        raise ValueError("The provided input 'max_interval_days' is not valide.")
    #check enforce_temporal_order
    if not isinstance(enforce_temporal_order,bool):
        raise TypeError("The provided input 'enforce_temporal_order' must be a bool.")
    #check correction
    correction_method_check(correction,cutoff)

    # time used
    time_start = time.time()

    # --------run the pipeline--------
    #run PheWAS
    phewas_log_file = os.path.join(output_dir,f'{project_prefix}_phewas.log')
    phewas_result_file = os.path.join(output_dir,f'{project_prefix}_phewas_result.csv')
    phewas_result = phewas(data=data, n_threshold=n_threshold_phewas,
                            n_process=n_process, system_exl=system_exl, covariates=covariates, 
                            correction=correction, cutoff=cutoff, log_file=phewas_log_file)
    phewas_result.to_csv(phewas_result_file,index=False)
    #filter the positive associations if needed
    if keep_positive_associations:
        if data.study_design == "exposed-only cohort":
            pass
        else:
            phewas_result = phewas_result[phewas_result['phewas_coef']>=0]
    #number of significant phecodes, saved it to the results_dict
    n_sig_phecodes = len(phewas_result[phewas_result['phewas_p_significance']==True])
    print(f'Number of significant diseases identified in the PheWAS analysis: {n_sig_phecodes}.')
    #generate intermediate data
    data.disease_pair(phewas_result=phewas_result,
                      min_interval_days=min_interval_days,
                      max_interval_days=max_interval_days,
                      force=True)
    if save_intermediate_data:
        data_intermediate_file = os.path.join(output_dir,f'{project_prefix}_data_withtra')
        data.save(data_intermediate_file)
    
    #run comorbidity strength estimation
    comorbidity_log_file = os.path.join(output_dir,f'{project_prefix}_comorbidity_strength.log')
    com_strength_result_file = os.path.join(output_dir,f'{project_prefix}_comorbidity_strength_result.csv')
    com_strength_result = comorbidity_strength(data=data, n_threshold=n_threshold_comorbidity,n_process=n_process, 
                                               correction_phi=correction, correction_RR=correction,
                                               cutoff_phi=cutoff, cutoff_RR=cutoff,
                                               log_file=comorbidity_log_file)
    com_strength_result.to_csv(com_strength_result_file,index=False)
    if keep_positive_associations:
        com_strength_result = com_strength_result[(com_strength_result['phi'] > 0) & (com_strength_result['RR'] > 1)]
    #number of significant comorbidity pairs, saved it to the results_dict
    n_sig_com_strength = len(com_strength_result[(com_strength_result['RR_p_significance']==True) & 
                                                (com_strength_result['phi_p_significance']==True)])
    print(f'Number of significant disease pairs identified in the comorbidity strength estimation: {n_sig_com_strength}.')
    
    #mode v1
    if pipeline_mode == 'v1':
        #run binomial test
        binomial_log_file = os.path.join(output_dir,f'{project_prefix}_binomial.log')
        binomial_result_file = os.path.join(output_dir,f'{project_prefix}_binomial_result.csv')
        binomial_result = binomial_test(data=data, comorbidity_strength_result=com_strength_result,
                                        n_process=1, correction=correction, cutoff=cutoff,
                                        enforce_temporal_order=enforce_temporal_order,
                                        log_file=binomial_log_file)
        binomial_result.to_csv(binomial_result_file,index=False)
        #number of temporal pairs
        n_temporal_pairs = len(binomial_result[binomial_result['binomial_p_significance']==True])
        print(f'Number of significant disease pairs identified in the binomial test: {n_temporal_pairs}.')
        #run comorbidity network analysis
        com_network_log_file = os.path.join(output_dir,f'{project_prefix}_comorbidity_network.log')
        com_network_result_file = os.path.join(output_dir,f'{project_prefix}_comorbidity_network_result.csv')
        if parameter_dict:
            com_network_result = comorbidity_network(data=data, comorbidity_strength_result=com_strength_result,
                                                    n_process=n_process, method=method, 
                                                    covariates=covariates, correction=correction, cutoff=cutoff,
                                                    log_file=com_network_log_file, **parameter_dict)
        else:
            com_network_result = comorbidity_network(data=data, comorbidity_strength_result=com_strength_result,
                                                    n_process=n_process, method=method, 
                                                    covariates=covariates, correction=correction, cutoff=cutoff,
                                                    log_file=com_network_log_file)
            
        com_network_result.to_csv(com_network_result_file,index=False)
        #number of significant comorbidity network pairs
        n_sig_com_network = len(com_network_result[com_network_result['comorbidity_p_significance']==True])
        print(f'Number of significant disease pairs identified in the comorbidity network analysis: {n_sig_com_network}.')
    elif pipeline_mode == 'v2':
        #run comorbidity network analysis
        com_network_log_file = os.path.join(output_dir,f'{project_prefix}_comorbidity_network.log')
        com_network_result_file = os.path.join(output_dir,f'{project_prefix}_comorbidity_network_result.csv')
        if parameter_dict:
            com_network_result = comorbidity_network(data=data, comorbidity_strength_result=com_strength_result,
                                                    n_process=n_process, method=method, 
                                                    covariates=covariates, correction=correction, cutoff=cutoff,
                                                    log_file=com_network_log_file, **parameter_dict)
        else:
            com_network_result = comorbidity_network(data=data, comorbidity_strength_result=com_strength_result,
                                                    n_process=n_process, method=method, 
                                                    covariates=covariates, correction=correction, cutoff=cutoff,
                                                    log_file=com_network_log_file)
        com_network_result.to_csv(com_network_result_file,index=False)
        #number of significant comorbidity network pairs
        n_sig_com_network = len(com_network_result[com_network_result['comorbidity_p_significance']==True])
        print(f'Number of significant disease pairs identified in the comorbidity network analysis: {n_sig_com_network}.')
        #run binomial test
        binomial_log_file = os.path.join(output_dir,f'{project_prefix}_binomial.log')
        binomial_result_file = os.path.join(output_dir,f'{project_prefix}_binomial_result.csv')
        binomial_result = binomial_test(data=data, comorbidity_strength_result=com_strength_result,
                                        comorbidity_network_result=com_network_result,
                                        n_process=1, correction=correction, cutoff=cutoff,
                                        enforce_temporal_order=enforce_temporal_order,
                                        log_file=binomial_log_file)
        binomial_result.to_csv(binomial_result_file,index=False)
        #number of temporal pairs
        n_temporal_pairs = len(binomial_result[binomial_result['binomial_p_significance']==True])
        print(f'Number of significant disease pairs identified in the binomial test: {n_temporal_pairs}.')

    #run disease trajectory analysis
    # first get the final covariates list for disease trajectory analysis, 
    # by exluding potential categorical variables used for matching
    # covariates_trajectory = covariates_check(covariates,data.get_attribute('phenotype_info'),matching_var_dict,exclude=True)
    #analysis
    trajectory_log_file = os.path.join(output_dir,f'{project_prefix}_trajectory.log')
    trajectory_result_file = os.path.join(output_dir,f'{project_prefix}_trajectory_result.csv')
    if parameter_dict:
        trajectory_result = disease_trajectory(data=data, comorbidity_strength_result=com_strength_result,
                                            binomial_test_result=binomial_result,
                                            n_process=n_process, method=method, 
                                            covariates=covariates,
                                            matching_var_dict=matching_var_dict, matching_n=matching_n,
                                            correction=correction, cutoff=cutoff,
                                            log_file=trajectory_log_file, **parameter_dict)
    else:
        trajectory_result = disease_trajectory(data=data, comorbidity_strength_result=com_strength_result,
                                            binomial_test_result=binomial_result,
                                            n_process=n_process, method=method, 
                                            covariates=covariates,
                                            matching_var_dict=matching_var_dict, matching_n=matching_n,
                                            correction=correction, cutoff=cutoff,
                                            log_file=trajectory_log_file)
        
    trajectory_result.to_csv(trajectory_result_file,index=False)
    #number of significant disease trajectories
    n_sig_trajectory = len(trajectory_result[trajectory_result['trajectory_p_significance']==True])
    print(f'Number of significant disease pairs identified in the disease trajectory analysis: {n_sig_trajectory}.')

    #calculate time used
    time_end = time.time()
    time_used = round((time_end-time_start)/60,2)
    print(f'The disease network analysis pipeline has been completed, total time: {time_used} mins.')
    
    #return all the results df in order
    return phewas_result, com_strength_result, com_network_result, binomial_result, trajectory_result

    
