from numpy import *
import mymath
import readandwrite
import various
import datastruct


def setup_cconst(modestoconst, cartgrad):
    Xconstmat = array(1*modestoconst)
    Xconstmat = mymath.orthonormalize_mat(Xconstmat)
    Xconstmat = mymath.remove_zrows(Xconstmat)
    for i in range(len(modestoconst)):
        cartgrad = cartgrad-sum(cartgrad*Xconstmat[i])*Xconstmat[i]
    return cartgrad


def null_cartesians(Bmat, modestoconst):
    modestoconst = array(modestoconst)
    for i in range(len(modestoconst)):
        if sum(abs(modestoconst[i])) == 1.:
            for j in range(len(modestoconst[i])):
                if abs(modestoconst[i][j]) == 1.:
                    Bmat[:, j] = 0.00
    return Bmat


def filterout_cconst(modestoconst, newc, cartesian):
    modestoconst = array(modestoconst)
    modestoconst = mymath.orthonormalize_mat(modestoconst)
    # newc=various.change_format(newc)
    # cartesian=various.change_format(cartesian)
    dcart = newc-cartesian
    sdcart = mymath.vector_size(dcart)
    for i in range(len(modestoconst)):
        smodestoconst = mymath.vector_size(modestoconst[i])
        dcart = dcart-modestoconst[i] * \
            inner(dcart, modestoconst[i])/(smodestoconst**2)
    newc = cartesian+dcart
    # newc=various.format_change(newc)
    return newc


def get_iconst(coords, crt, lattmat, atomictags):
    CONSTRAINTS = []
    ircdata = []
    try:
        consttags, constwhat, constwhere, constcoefs, conststat, complextype =\
            readandwrite.read_cwhatwhere('ICOORD', crt, lattmat)

    except IOError:
        numconst = 0
        constcoefs = []
        conststat = []
    else:
        numconst = len(constwhat)
        print('GEOMETRICAL CONSTRAINTS DETECTED')
    if numconst > 0:
        dim = len(coords)
        if constcoefs == []:
            for i in range(numconst):
                constwhattags = give_whattags(
                    constwhat[i], consttags[i], atomictags)
                coords.append(datastruct.Complextype('simple', [1], consttags[i], constwhat[i],
                                                     constwhattags, constwhere[i], 0.0, array([]), 'Keep'))
                CONSTRAINTS.append(i+dim)
        else:
            for i in range(len(constcoefs)):
                if len(constcoefs[i]) != len(consttags):
                    print('incorect definition of complex constrained coordinate!!!')
                else:
                    if len(constcoefs[i]) == 1:
                        constwhattags = give_whattags(
                            constwhat[i], consttags[i], atomictags)
                        coords.append(datastruct.Complextype('simple', constcoefs[i], consttags[i], constwhat[i],
                                                             constwhattags, constwhere[i], 0.0, array([]), conststat[i]))
                        CONSTRAINTS.append(i+dim)
                    else:
                        coords.append(datastruct.Complextype(complextype[i], constcoefs[i], consttags, constwhat,
                                                             [None], constwhere, 0.0, array([]), conststat[i]))
                        CONSTRAINTS.append(i+dim)

    for i in range(len(coords)):
        if coords[i].dtyp == 'is' or coords[i].dtyp == 'iz':
            ircdata = read_irccar('IRCCAR')

    return coords, CONSTRAINTS, conststat, ircdata


def read_irccar(f):
    ircdata = []
    try:
        f = open(f, 'r')
    except FileNotFoundError:
        pass
    else:
        i = 0
        dimX = 0
        for line in f.readlines():
            line = line.split()
            if i == 0:
                npoints = int(line[0])
            else:
                dummy = []
                for i in range(len(line)):
                    dummy.append(float(line[i]))
                if dimX == 0:
                    dimX = len(dummy)
                    if dimX == 0:
                        print("koko")
                else:
                    if len(dummy) != dimX:
                        print("koko")
                ircdata.append(dummy)
            i += 1
        ircdata = array(ircdata)
        f.close()
    return ircdata


def give_whattags(what, tags, atomictags):
    whattags = [None]
    if tags == 'R' or tags == 'A' or tags == 'T' or tags == 'tV' or tags == 'IR1'\
            or tags == 'IR6' or tags == 'RatioR' or tags == 'X' or \
            tags == 'Y' or tags == 'Z' or tags == 'fX' or tags == 'fY' or tags == 'fZ':
        whattags = []
        for i in range(len(what)):
            whattags.append(atomictags[what[i]])
    return whattags


def tell_iconst(coords):
    CONSTRAINTS = []
    conststat = []
    for i in range(len(coords)):
        if coords[i].status != 'free':
            CONSTRAINTS.append(i)
            conststat.append(coords[i].status)
    return CONSTRAINTS, conststat


def setup_iconst(constrains, Bmatp):
    Cmatp = zeros((len(Bmatp), len(Bmatp[0])), float)
    # print "b",len(Bmatp),len(Bmatp[0])
    # print "c",len(Cmatp),len(Cmatp[0])
    # print constrains
    for i in constrains:
        Cmatp[i] = Bmatp[i]
    Cmatp = mymath.orthonormalize_mat(Cmatp)
    oBmatp = dot(Bmatp, transpose(Cmatp))
    for i in range(len(oBmatp)):
        oBmatp[i][i] = 0
    oBmatp = dot(oBmatp, Cmatp)
    Bmatp = Bmatp-oBmatp
    return Bmatp


def con_con(coords, CONSTRAINTS, conststat):
    newCONSTRAINTS = []
    newconststat = []
    for i in range(len(coords)):
        dum = 0
        for j in CONSTRAINTS:
            if i == j:
                dum = 1
                break
        if dum == 0:
            newCONSTRAINTS.append(i)
            newconststat.append('Keep')
    return newCONSTRAINTS, newconststat
