import os
from numpy import log10
from cogent3 import LoadSeqs, LoadTree, DNA, LoadTable
from cogent3.evolve.models import CNFGTR
from cogent3.maths.stats import chisqprob

submodel = CNFGTR()

def do_lr_test(lnL_alt, lnL_null, nfp_alt, nfp_null):
    LR = 2 * (lnL_alt - lnL_null)
    LR = LR if LR >= 0 else 0
    df = nfp_alt - nfp_null
    pval = chisqprob(LR, df)
    return pval

def eval_alignment(path):
    # append the gene name, the p-vals from the two LRTs,
    # omega from the Human edge and the length from the Human edge

    results = []
    gene_name = os.path.basename(path).split(".")[0]
    results.append(gene_name)

    aln = LoadSeqs(path, moltype=DNA)
    tree = LoadTree(tip_names=aln.names)
    lf = submodel.make_likelihood_function(tree)

    lf.set_alignment(aln)

    ##### First LRT calc
    #define the neutral null
    lf.set_param_rule('omega', value=1.0, is_constant=True)
    lf.optimise()
    neutral_null_nfp = lf.get_num_free_params()
    neutral_null_lnL = lf.get_log_likelihood()

    #define the neutral alt
    lf.set_param_rule('omega', value=1.0, is_constant=False)
    lf.optimise()
    neutral_alt_nfp = lf.get_num_free_params()
    neutral_alt_lnL = lf.get_log_likelihood()

    results.append(do_lr_test(neutral_alt_lnL, neutral_null_lnL, neutral_alt_nfp, neutral_null_nfp))

    ##### Second LRT calc
    #define the Human alternate
    lf.set_param_rule('omega', is_constant=True, edge="Human")
    human_alt_nfp = lf.get_num_free_params()
    human_alt_lnL = lf.get_log_likelihood()

    results.append(do_lr_test(human_alt_lnL, neutral_null_lnL, human_alt_nfp, neutral_null_nfp))
    print(do_lr_test(human_alt_lnL, neutral_null_lnL, human_alt_nfp, neutral_null_nfp))


    results.append(lf.get_param_value("omega", edge="Human"))
    results.append(lf.get_param_value("length", edge="Human"))

    return results

paths = get_one2one_align_paths()
result = eval_alignment(paths[0])

headers = ['Name', 'pval(Neutral)', 'pval(Human diff)', 'Omega from Human edge', 'Length from Human edge']
tbl = LoadTable(header= headers, rows= [result])
tbl