from SigProfilerAssignment import decomposition as decomp


def decompose_fit(
    samples,
    output,
    signatures=None,
    signature_database=None,
    nnls_add_penalty=0.05,
    nnls_remove_penalty=0.01,
    initial_remove_penalty=0.05,
    genome_build="GRCh37",
    cosmic_version=3.4,
    make_plots=True,
    collapse_to_SBS96=True,
    connected_sigs=True,
    verbose=False,
    devopts=None,
    new_signature_thresh_hold=0.8,
    exclude_signature_subgroups=None,
    exome=False,
    input_type="matrix",
    context_type="96",
    export_probabilities=True,
    export_probabilities_per_mutation=False,
    volume=None,
):
    decomp.spa_analyze(
        samples=samples,
        output=output,
        signatures=signatures,
        signature_database=signature_database,
        nnls_add_penalty=nnls_add_penalty,
        nnls_remove_penalty=nnls_remove_penalty,
        initial_remove_penalty=initial_remove_penalty,
        genome_build=genome_build,
        cosmic_version=cosmic_version,
        make_plots=make_plots,
        collapse_to_SBS96=collapse_to_SBS96,
        connected_sigs=connected_sigs,
        verbose=verbose,
        decompose_fit_option=True,
        denovo_refit_option=False,
        cosmic_fit_option=False,
        devopts=devopts,
        new_signature_thresh_hold=new_signature_thresh_hold,
        exclude_signature_subgroups=exclude_signature_subgroups,
        exome=exome,
        input_type=input_type,
        context_type=context_type,
        export_probabilities=export_probabilities,
        export_probabilities_per_mutation=export_probabilities_per_mutation,
        volume=volume,
    )


def denovo_fit(
    samples,
    output,
    signatures=None,
    signature_database=None,
    nnls_add_penalty=0.05,
    nnls_remove_penalty=0.01,
    initial_remove_penalty=0.05,
    genome_build="GRCh37",
    cosmic_version=3.4,
    make_plots=True,
    collapse_to_SBS96=True,
    connected_sigs=True,
    verbose=False,
    devopts=None,
    new_signature_thresh_hold=0.8,
    exome=False,
    input_type="matrix",
    context_type="96",
    export_probabilities=True,
    export_probabilities_per_mutation=False,
    volume=None,
):
    decomp.spa_analyze(
        samples=samples,
        output=output,
        signatures=signatures,
        signature_database=signature_database,
        nnls_add_penalty=nnls_add_penalty,
        nnls_remove_penalty=nnls_remove_penalty,
        initial_remove_penalty=initial_remove_penalty,
        genome_build=genome_build,
        cosmic_version=cosmic_version,
        new_signature_thresh_hold=new_signature_thresh_hold,
        make_plots=make_plots,
        collapse_to_SBS96=collapse_to_SBS96,
        connected_sigs=connected_sigs,
        verbose=verbose,
        decompose_fit_option=False,
        denovo_refit_option=True,
        cosmic_fit_option=False,
        devopts=devopts,
        exome=exome,
        input_type=input_type,
        context_type=context_type,
        export_probabilities=export_probabilities,
        export_probabilities_per_mutation=export_probabilities_per_mutation,
        volume=volume,
    )


def cosmic_fit(
    samples,
    output,
    signatures=None,
    signature_database=None,
    nnls_add_penalty=0.05,
    nnls_remove_penalty=0.01,
    initial_remove_penalty=0.05,
    genome_build="GRCh37",
    cosmic_version=3.4,
    make_plots=True,
    collapse_to_SBS96=True,
    connected_sigs=True,
    verbose=False,
    devopts=None,
    exclude_signature_subgroups=None,
    exome=False,
    input_type="matrix",
    context_type="96",
    export_probabilities=True,
    export_probabilities_per_mutation=False,
    sample_reconstruction_plots=False,
    volume=None,
):
    decomp.spa_analyze(
        samples=samples,
        output=output,
        signatures=signatures,
        signature_database=signature_database,
        nnls_add_penalty=nnls_add_penalty,
        nnls_remove_penalty=nnls_remove_penalty,
        initial_remove_penalty=initial_remove_penalty,
        genome_build=genome_build,
        cosmic_version=cosmic_version,
        make_plots=make_plots,
        collapse_to_SBS96=collapse_to_SBS96,
        connected_sigs=connected_sigs,
        verbose=verbose,
        decompose_fit_option=False,
        denovo_refit_option=False,
        cosmic_fit_option=True,
        devopts=devopts,
        exclude_signature_subgroups=exclude_signature_subgroups,
        exome=exome,
        input_type=input_type,
        context_type=context_type,
        export_probabilities=export_probabilities,
        export_probabilities_per_mutation=export_probabilities_per_mutation,
        sample_reconstruction_plots=sample_reconstruction_plots,
        volume=volume,
    )
