from ion import *
from kandinsky import *
from time import sleep


def _jac(a, n):
    v = [[1.0 if i == j else 0.0 for j in range(n)] for i in range(n)]
    for _ in range(100):
        o = 0.0
        for p in range(n):
            for q in range(p + 1, n):
                o += a[p][q] * a[p][q]
        if o < 1e-20:
            break
        for p in range(n):
            for q in range(p + 1, n):
                if abs(a[p][q]) < 1e-15:
                    continue
                th = (a[q][q] - a[p][p]) / (2.0 * a[p][q])
                t = (1.0 if th >= 0 else -1.0) / (abs(th) + (th * th + 1.0) ** 0.5)
                c = 1.0 / (t * t + 1.0) ** 0.5
                s = t * c
                for i in range(n):
                    x = a[i][p]; y = a[i][q]
                    a[i][p] = c * x - s * y; a[i][q] = s * x + c * y
                for i in range(n):
                    x = a[p][i]; y = a[q][i]
                    a[p][i] = c * x - s * y; a[q][i] = s * x + c * y
                for i in range(n):
                    x = v[i][p]; y = v[i][q]
                    v[i][p] = c * x - s * y; v[i][q] = s * x + c * y
    return [a[i][i] for i in range(n)], v


def _mul(a, b, n):
    r = [[0.0] * n for _ in range(n)]
    for i in range(n):
        for j in range(n):
            s = 0.0
            for k in range(n):
                s += a[i][k] * b[k][j]
            r[i][j] = s
    return r


def _charpoly(B, n):
    M = [[0.0] * n for _ in range(n)]; co = [1.0]
    for k in range(1, n + 1):
        AM = _mul(B, M, n)
        for i in range(n):
            AM[i][i] += co[k - 1]
        M = AM; BM = _mul(B, M, n); tr = 0.0
        for i in range(n):
            tr += BM[i][i]
        co.append(-tr / k)
    return co


def _poly(co):
    n = len(co) - 1; t = []
    for i in range(len(co)):
        p = n - i; c = round(co[i], 3)
        if abs(c) < 1e-6:
            continue
        if p == 0:
            t.append(str(c))
        elif p == 1:
            t.append(str(c) + "x")
        else:
            t.append(str(c) + "x^" + str(p))
    return (" + ".join(t)).replace("+ -", "- ") + " = 0"


def _f(v):
    return str(int(v)) if v == int(v) else str(round(v, 2))


def huckel(liaisons, ne, n=None, hetero=None, N=None):
    "liaisons=[[i,j]..] voisins, ne e-pi, hetero={at:(h,k)}, N={at:e-pi}. E=a+m*b, b<0."
    if n is None:
        n = 0
        for b in liaisons:
            n = max(n, b[0], b[1])
    if hetero is None:
        hetero = {}
    if N is None:
        N = {}
    A = [[0.0] * n for _ in range(n)]
    for at in hetero:
        A[at - 1][at - 1] = hetero[at][0]
    for b in liaisons:
        k = 1.0
        if b[0] in hetero:
            k = hetero[b[0]][1]
        if b[1] in hetero:
            k = hetero[b[1]][1]
        A[b[0] - 1][b[1] - 1] = k; A[b[1] - 1][b[0] - 1] = k
    eig, vec = _jac([r[:] for r in A], n)
    idx = sorted(range(n), key=lambda i: -eig[i])
    m = [eig[i] for i in idx]
    C = [[vec[a][i] for a in range(n)] for i in idx]
    occ = [0.0] * n; reste = ne; o = 0
    while o < n and reste > 0:
        g = [o]
        while g[-1] + 1 < n and abs(m[g[-1] + 1] - m[o]) < 1e-6:
            g.append(g[-1] + 1)
        met = min(2 * len(g), reste)
        for gi in g:
            occ[gi] = met / len(g)
        reste -= met; o = g[-1] + 1
    ho = -1; bv = -1
    for i in range(n):
        if occ[i] > 1e-9:
            ho = i
    for i in range(n):
        if occ[i] < 2 - 1e-9:
            bv = i; break
    print("HUCKEL", n, "at,", ne, "e-pi (E=a+m*b, b<0)")
    print("Matrice (x diag, k voisins):")
    for i in range(n):
        s = ""
        for j in range(n):
            if i == j:
                d = A[i][i]
                s += " x" if abs(d) < 1e-9 else (" x+" + _f(d))
            else:
                s += " " + _f(A[i][j])
        print(s)
    print("det=0 :", _poly(_charpoly([[-A[p][q] for q in range(n)] for p in range(n)], n)))
    print("racines x=-m :", [round(-mm, 3) for mm in m])
    print("(E=a-x*b=a+m*b)")
    for i in range(n):
        nat = "L" if m[i] > 1e-6 else ("NL" if abs(m[i]) < 1e-6 else "AL")
        oc = round(occ[i], 2) if occ[i] > 1e-9 else 0
        print("phi" + str(i + 1), "a+", round(m[i], 3), "b", nat, oc, "e")
    print("HO=phi" + str(ho + 1), "a+", round(m[ho], 3), "b")
    print("BV=phi" + str(bv + 1), "a+", round(m[bv], 3), "b")
    sm = 0.0
    for i in range(n):
        sm += occ[i] * m[i]
    print("E_pi=", ne, "a +", round(sm, 3), "b")
    print("--coeffs (at1..n)--")
    for i in range(n):
        s = "phi" + str(i + 1) + ":"
        for a in range(n):
            s += " " + str(round(C[i][a], 3))
        print(s)
    print("--q / charge d--")
    tot = 0.0
    for a in range(n):
        q = 0.0
        for i in range(n):
            q += occ[i] * C[i][a] * C[i][a]
        Na = N.get(a + 1, 1); d = Na - q; tot += d
        print("at", a + 1, "q=", round(q, 3), "d=", round(d, 3))
    print("sum d=", round(tot, 3))
    return m, C, occ


def deloc(epi, eloc):
    "E_deloc = coeff_b(E_pi delocalise) - coeff_b(E_pi localise), en beta."
    d = epi - eloc
    print("E_deloc=(", round(d, 3), ")b  (b<0: negatif=stabilisant)")
    return d


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)
        while not (keydown(KEY_UP) or keydown(KEY_DOWN) or keydown(KEY_OK) or keydown(KEY_EXE) or keydown(KEY_BACK)):
            pass
        if keydown(KEY_UP):
            sel = (sel - 1) % len(opts)
        elif keydown(KEY_DOWN):
            sel = (sel + 1) % len(opts)
        elif keydown(KEY_BACK):
            return -1
        else:
            while keydown(KEY_OK) or keydown(KEY_EXE):
                pass
            return sel
        sleep(0.16)


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


def menu():
    while True:
        c = pick("HUCKEL", ["Nouvelle molecule", "E delocalisation", "Console"])
        if c == -1 or c == 2:
            print("ex: huckel([[1,2],[2,3]],2)")
            return
        if c == 0:
            print("Que le systeme pi (H,sp3 exclus)")
            print("atomes 1..n ; cycle: liaison n-1")
            n = int(_e("Nb atomes pi? "))
            nl = int(_e("Nb liaisons? "))
            L = []
            for i in range(nl):
                print("liaison " + str(i + 1) + " entre:")
                a = int(_e(" atome? "))
                b = int(_e(" voisin? "))
                L.append([a, b])
            ne = int(_e("Nb e-pi? "))
            nh = int(_e("Nb hetero(0 sinon)? "))
            if nh:
                print("h:sur alpha k:sur beta N:e-pi")
            H = {}; N = {}
            for i in range(nh):
                at = int(_e("hetero atome? "))
                H[at] = (_e("  h? "), _e("  k? "))
                N[at] = int(_e("  N e-pi? "))
            huckel(L, ne, n, H if H else None, N if N else None)
            input("--OK--")
        elif c == 1:
            deloc(_e("coeff b E_pi? "), _e("coeff b E_loc? "))
            input("--OK--")


menu()