from typing import Iterable, List, Optional, Tuple
class Sonia(object):
    def __init__(self,
                 features: List[Iterable[str]] = [],
                 data_seqs: List[Iterable[str]] = [],
                 gen_seqs: List[Iterable[str]] = [],
                 load_dir: Optional[str] = None,
                 pgen_model: Optional[str] = None, # pgen_model keyword parameter here
                 max_depth: int = 25,
                 max_L: int = 30,
                 load_seqs: bool = True,
                 l2_reg: float = 0.,
                 l1_reg: float = 0.,
                 min_energy_clip: int = -5,
                 max_energy_clip: int = 10,
                 seed: Optional[int] = None,
                 vj: bool = False,
                 custom_pgen_model: Optional[str] = None,
                 processes: Optional[int] = None,
                 objective: str = 'BCE',
                 gamma: int = 1,
                 include_joint_genes: bool = False,
                 joint_vjl: bool = False,
                 include_aminoacids: bool =True
                ) -> None:
        # Other initializations unchanged.
        if load_dir is None and pgen_model is None:
            raise ValueError('Both load_dir and pgen_model cannot be None.')
        # Use pgen_model in the SONIA/SoNNia model folder.
        if pgen_model is None:
            self.pgen_model = load_dir
        else:
            self.pgen_model = pgen_model

import olga.generation_probability as generation_probability
import olga.sequence_generation as sequence_generation
import olga.load_model as olga_load_model
import olga.generation_probability as pgen
import olga.sequence_generation as seq_gen

DEFAULT_CHAIN_TYPES = {'humanTRA': 'human_T_alpha', 'human_T_alpha': 'human_T_alpha',
                       'humanTRB': 'human_T_beta', 'human_T_beta': 'human_T_beta',
                       'humanIGH': 'human_B_heavy', 'human_B_heavy': 'human_B_heavy',
                       'humanIGK': 'human_B_kappa', 'human_B_kappa': 'human_B_kappa',
                       'humanIGL': 'human_B_lambda', 'human_B_lambda': 'human_B_lambda',
                       'mouseTRB': 'mouse_T_beta', 'mouse_T_beta': 'mouse_T_beta',
                       'mouseTRA': 'mouse_T_alpha','mouse_T_alpha':'mouse_T_alpha'}

def define_pgen_model(pgen_model: str,
                      vj: bool = False,
                      return_files: bool = False
                     ):
    if pgen_model in DEFAULT_CHAIN_TYPES:
        pgen_model = DEFAULT_CHAIN_TYPES[pgen_model]
        filedir = os.path.dirname(os.path.abspath(__file__))
        main_folder = os.path.join(filedir, 'default_models', pgen_model)
    elif os.path.isdir(pgen_model):
        main_folder = pgen_model
    else:
        options = f'{list(DEFAULT_CHAIN_TYPES.keys())[::2]}'[1:-1]
        raise ValueError('pgen_model is neither a directory that exists '
                         f'nor one of the default options ({options}). '
                         'Try using one of the default options or an existing '
                         'directory containing the pgen model.')

    pgen_files = ('model_params.txt', 'model_marginals.txt',
                  'V_gene_CDR3_anchors.csv', 'J_gene_CDR3_anchors.csv')
    files_in_dir = set(os.listdir(pgen_model))
    missing_files = set(pgen_files) - files_in_dir

    if len(missing_files) > 0:
        missing_files = f'{missing_files}'[1:-1]
        raise RuntimeError('The pgen model cannot be loaded. The following files '
                           f'are missing: {missing_files}.')

    params_file_name = os.path.join(main_folder, pgen_files[0])
    marginals_file_name = os.path.join(main_folder, pgen_files[1])
    V_anchor_pos_file = os.path.join(main_folder, pgen_files[2])
    J_anchor_pos_file = os.path.join(main_folder, pgen_files[3])

    if vj:
        model_str = 'VJ'
    else:
        model_str = 'VDJ'

    genomic_data = getattr(olga_load_model, f'GenomicData{model_str}')()
    genomic_data.load_igor_genomic_data(params_file_name,
                                        V_anchor_pos_file,
                                        J_anchor_pos_file)
    generative_model = getattr(olga_load_model, f'GenerativeModel{model_str}')()
    generative_model.load_and_process_igor_model(marginals_file_name)
    pgen_model = getattr(pgen, f'GenerationProbability{model_str}')(generative_model, genomic_data)
    seqgen_model = getattr(seq_gen, f'SequenceGeneration{model_str}')(generative_model, genomic_data)

    to_return = (genomic_data, generative_model, pgen_model, seqgen_model)

    if return_files:
        return to_return + (params_file_name, marginals_file_name, V_anchor_pos_file, J_anchor_pos_file)
    return to_return