from ion import *
from kandinsky import *
from math import cos, pi


def remplir(energies, ne, deg=None, ev=False):
    "Range/remplit les OM, donne HO/BV. energies=coeff b (grand=bas) ou eV si ev=True."
    if deg is None:
        deg = [1] * len(energies)
    orbs = []
    for i in range(len(energies)):
        for _ in range(deg[i]):
            orbs.append(energies[i])
    orbs.sort(reverse=(not ev))
    n = len(orbs); occ = [0] * n; reste = ne
    for i in range(n):
        met = min(2, reste); occ[i] = met; reste -= met
        if reste <= 0:
            break
    ho = -1; bv = -1
    for i in range(n):
        if occ[i] > 0:
            ho = i
    for i in range(n):
        if occ[i] < 2:
            bv = i; break
    for i in range(n):
        tag = ""
        if i == ho:
            tag += " HO"
        if i == bv:
            tag += " BV"
        e = round(orbs[i], 3)
        print((str(e) + "eV") if ev else ("a+ " + str(e) + "b"), occ[i], "e", tag)
    lab = (lambda x: str(round(x, 3)) + "eV") if ev else (lambda x: "a+" + str(round(x, 3)) + "b")
    print("HO", lab(orbs[ho]), "| BV", lab(orbs[bv]))
    return orbs, occ, ho, bv


def ecart(mho, mbv):
    "Ecart HO-BV en |beta| = |mHO-mBV|. Petit ecart = interaction forte."
    e = abs(mho - mbv)
    print("ecart HO-BV=", round(e, 3), "|b|")
    return e


def couple(hoA, bvA, hoB, bvB):
    "2 partenaires: le plus petit ecart HO-BV domine (donne le donneur)."
    e1 = abs(hoA - bvB); e2 = abs(hoB - bvA)
    print("HO(A)-BV(B)=", round(e1, 3), "|b|")
    print("HO(B)-BV(A)=", round(e2, 3), "|b|")
    if e1 < e2:
        print("-> HO(A)-BV(B) domine (A donneur)")
    else:
        print("-> HO(B)-BV(A) domine (B donneur)")


def regio(coeffs):
    "Controle orbitalaire: attaque sur le plus gros |coefficient|."
    b = 0
    for i in range(len(coeffs)):
        if abs(coeffs[i]) > abs(coeffs[b]):
            b = i
    for i in range(len(coeffs)):
        print("at", i + 1, "|c|=", round(abs(coeffs[i]), 3), "<<" if i == b else "")
    print("-> attaque at", b + 1)


def charge_max(ch):
    "Controle de charge: nucleophile sur l'atome le plus positif."
    b = 0
    for i in range(len(ch)):
        if ch[i] > ch[b]:
            b = i
    for i in range(len(ch)):
        print("at", i + 1, "q=", round(ch[i], 3), "<<" if i == b else "")
    print("-> nucleophile sur at", b + 1)


def recouvrement(s0, ang):
    "Recouvrement hors axe: S=S0*cos(angle). cos90=0 => interaction nulle."
    S = s0 * cos(ang * pi / 180.0)
    print("S=S0*cos(", ang, ")=", round(S, 4))
    if abs(S) < 1e-9:
        print("-> NUL, pas d'interaction")
    return S


def pick(titre, opts):
    sel = 0
    while True:
        fill_rect(0, 0, 320, 222, color(255, 255, 255))
        draw_string(titre, 2, 2, color(255, 255, 255), color(60, 90, 200))
        for i in range(len(opts)):
            y = 26 + i * 20
            if i == sel:
                fill_rect(0, y, 320, 18, color(60, 120, 240))
                draw_string(opts[i], 6, y, color(255, 255, 255), color(60, 120, 240))
            else:
                draw_string(opts[i], 6, y)
        k = -1
        while k == -1:
            if keydown(KEY_UP):
                k = 0
            elif keydown(KEY_DOWN):
                k = 1
            elif keydown(KEY_BACK):
                k = 2
            elif keydown(KEY_OK) or keydown(KEY_EXE):
                k = 3
        while keydown(KEY_UP) or keydown(KEY_DOWN) or keydown(KEY_BACK) or keydown(KEY_OK) or keydown(KEY_EXE):
            pass
        if k == 0:
            sel = (sel - 1) % len(opts)
        elif k == 1:
            sel = (sel + 1) % len(opts)
        elif k == 2:
            return -1
        else:
            return sel


def _e(m):
    return float(input(m))


def menu():
    while True:
        c = pick("FRONTIERES", ["Remplissage HO/BV", "Ecart HO-BV", "Choix couple", "Regio orbitalaire", "Regio charge", "Recouvrement cos", "Console"])
        if c == -1 or c == 6:
            print("ex: remplir([1.414,0,-1.414],2)")
            return
        if c == 0:
            print("Q: HO/BV, remplir un diagramme d'OM")
            u = int(_e("1=coeff beta, 2=eV? "))
            k = int(_e("Nb niveaux? ")); E = []; D = []
            for i in range(k):
                E.append(_e("E" + str(i + 1) + "? "))
                D.append(int(_e("  degen(1 si simple)? ")))
            remplir(E, int(_e("Nb e-? ")), D, u == 2)
        elif c == 1:
            print("Q: force de l'interaction (ecart HO-BV)")
            ecart(_e("m HO? "), _e("m BV? "))
        elif c == 2:
            print("Q: quel couple? demande normale/inverse (Diels-Alder)")
            couple(_e("HO A? "), _e("BV A? "), _e("HO B? "), _e("BV B? "))
        elif c == 3:
            print("Q: regiochimie, controle orbitalaire (gros coeff)")
            k = int(_e("Nb atomes? ")); L = []
            for i in range(k):
                L.append(_e("c" + str(i + 1) + "? "))
            regio(L)
        elif c == 4:
            print("Q: regiochimie, controle de charge (plus positif)")
            k = int(_e("Nb atomes? ")); L = []
            for i in range(k):
                L.append(_e("q" + str(i + 1) + "? "))
            charge_max(L)
        elif c == 5:
            print("Q: recouvrement nul? (symetrie / cosinus)")
            recouvrement(_e("S0? "), _e("angle deg? "))
        input("--OK--")


menu()