import numpy as np
from scipy.special import betaln

def nal_to_beta(f,c):
    w=c/(1-c)
    return w*f, w*(1-f)

def hellinger2(a1,b1,a2,b2):
    return 1-np.exp(betaln((a1+a2)/2,(b1+b2)/2)-0.5*(betaln(a1,b1)+betaln(a2,b2)))

def rao_distance(f1,c1,f2,c2):
    a1,b1=nal_to_beta(f1,c1)
    a2,b2=nal_to_beta(f2,c2)
    h2=hellinger2(a1,b1,a2,b2)
    return 2*np.arcsinh(np.sqrt(h2/(1-h2+1e-15)))

def select_premises(goal_stv, candidates, top_k=3):
    """Rank candidate beliefs by Rao proximity to goal belief."""
    gf,gc=goal_stv
    scored=[(name,stv,rao_distance(gf,gc,stv[0],stv[1])) for name,stv in candidates]
    scored.sort(key=lambda x:x[2])
    return scored[:top_k]

# Test: goal is cat->entity stv(0.8,0.5), pick best premises
goal=("cat->entity",(0.8,0.5))
candidates=[("cat->animal",(0.9,0.8)),("cat->pet",(0.85,0.7)),("fish->swim",(0.95,0.9)),("cat->danger",(0.1,0.6)),("bird->fly",(0.7,0.3)),("cat->mammal",(0.88,0.75))]
print("Goal:",goal[0],"stv",goal[1])
print("Premise ranking by Rao distance:")
for name,stv,d in select_premises(goal[1],[c for c in candidates]):
    print(f"  {name} stv{stv} Rao={d:.4f}")