# Predict water positions
import time
import sys # To use sys.argv

import mdtraj
import numpy as np
import math

def BWHB(x1, x2, print_dist=False): # Bridge Water Hydrogen Bond
    d = np.sqrt(((x1*10-x2*10)**2).sum())
    if print_dist:
        print(f'{d:.3f}', end=' \t')
    return (1 / (1 + (d/2.6)**6)) / 0.58

def rotation3D(v, axis, degrees):
    """
    Return the rotation matrix associated with counterclockwise rotation about
    the given axis by theta radians. Using the Euler-Rodrigues formula:
    https://stackoverflow.com/questions/6802577/rotation-of-3d-vector
    """
    axis = np.asarray(axis)
    axis = axis / math.sqrt(np.dot(axis, axis))
    theta = degrees * math.pi / 180
    a = math.cos(theta / 2.0)
    b, c, d = -axis * math.sin(theta / 2.0)
    aa, bb, cc, dd = a * a, b * b, c * c, d * d
    bc, ad, ac, ab, bd, cd = b * c, a * d, a * c, a * b, b * d, c * d
    rot_mat = np.array([[aa + bb - cc - dd, 2 * (bc + ad), 2 * (bd - ac)],
                        [2 * (bc - ad), aa + cc - bb - dd, 2 * (cd + ab)],
                        [2 * (bd + ac), 2 * (cd - ab), aa + dd - bb - cc]])

    return np.dot(rot_mat, v)

def calcAngle(A, B, C):
    vec1 = np.array(A) - np.array(B)
    vec2 = np.array(C) - np.array(B)
    return np.arccos(np.dot(vec1, vec2) / np.linalg.norm(vec1) / np.linalg.norm(vec2))*180/ np.pi, vec1, vec2

def prepVec(v, length):
    v = v / np.linalg.norm(v)
    v.flatten()
    return v * length

def outputMockPDB(fname, p, e='X'):
    with open(fname,'w') as f:
#     f.write(f'{len(DOL)}\n\n')
        for ii in range(len(p)):
    #         ATOM      1  N   PRO A   1       8.316  21.206  21.530  1.00 17.44           N
            f.write(f'ATOM    {ii+1:3d}  {e}   {e*3} A   1    {p[ii][0]*10:8.3f}{p[ii][1]*10:8.3f}{p[ii][2]*10:8.3f}  0.00  0.00            \n')
        f.write('END\n')

def outputMolecularMockPDB(fname, p, e=None, resname='HOH'):
    assert len(e) == len(p[0]), "Input: p as a N_mol * N_atom * 3 matrix, e as a list of elements with length N_atom"
    resname += ' '*3  # Fill in resnames
    e2 = [x + ' '*3 for x in e]
    with open(fname,'w') as f:
#     f.write(f'{len(DOL)}\n\n')
        for jj in range(len(p)):
            for ii in range(len(p[jj])):
    #         ATOM      1  N   PRO A   1       8.316  21.206  21.530  1.00 17.44           N
                f.write(f'ATOM    {jj*len(p[jj])+ii+1:3d}  {e2[ii][:3]} {resname[:3]} A{jj:4d}    {p[jj][ii][0]*10:8.3f}{p[jj][ii][1]*10:8.3f}{p[jj][ii][2]*10:8.3f}  0.00  0.00            \n')
        f.write('END\n')

def calcWater(inItem, outFile=None, outResname='HOH', printTiming=False, printResult=False, processFrame=0, returnNumWater=False, returnHBscore=False):

    # processFrame takes a number, a list of numbers, or -1 for all frames

    if printTiming:
        t0 = time.time()

    if type(inItem) == str:
        comp = mdtraj.load(inItem)

    if type(inItem) == mdtraj.Trajectory:
        comp = inItem

    if printTiming:
        t1 = time.time()
        print(f'Load pdb: {t1-t0:.3f} s')
        t0 = time.time()

    protein_residue_list = [
        'ALA', 'ARG', 'ASN', 'ASP', 'CYS',
        'GLN', 'GLU', 'GLY', 'HIS', 'ILE',
        'LEU', 'LYS', 'MET', 'PHE', 'PRO',
        'SER', 'THR', 'TYR', 'VAL', 'TRP',
        'ACE', 'NME', 'HID', 'HIE', 'HIP', 'WAT', 'HOH', 'TIP3'] # For additional caps and protonation states of HIS

    atom_vdW_radii = { # in nanometers
        'H': 0.107,
        'C': 0.170,
        'N': 0.155,
        'O': 0.152,
        'F': 0.147,
        'S': 0.180,
        'P': 0.180,
        'Cl': 0.175,
        'Br': 0.185,
        'I': 0.198
    }

    atom_list = {}
    bond_to_symbol = {}
    bond_to_index = {}
    atom_type = {}
    atom_in_residue = {}
    NOS_index = []
    H_index = []
    ALL_vdW_radii = []

    for atom in comp.top.atoms:
    #     print(atom.index, atom)
        atom_list[atom.index] = atom.element.symbol
        atom_in_residue[atom.index] = atom.residue.name
        bond_to_symbol[atom.index] = []
        bond_to_index[atom.index] = []
        atom_type[atom.index] = []
        ALL_vdW_radii.append(atom_vdW_radii[atom.element.symbol])
        if atom.element.symbol in ['N','O','S']:
            NOS_index.append(atom.index)
        elif atom.element.symbol == 'H':
            H_index.append(atom.index)

    NOS_index = np.array(NOS_index)
    H_index = np.array(H_index)
    ALL_vdW_radii = np.array(ALL_vdW_radii)

    for bond in comp.top.bonds:
        bond_to_symbol[bond.atom1.index].append(bond.atom2.element.symbol)
        bond_to_symbol[bond.atom2.index].append(bond.atom1.element.symbol)
        bond_to_index[bond.atom1.index].append(bond.atom2.index)
        bond_to_index[bond.atom2.index].append(bond.atom1.index)

    for ii in range(len(atom_list)):
        bond_to_symbol[ii] = np.array(bond_to_symbol[ii])
        bond_to_index[ii] = np.array(bond_to_index[ii])
        sort_order = np.argsort(bond_to_symbol[ii])
        if len(sort_order) > 1:
    #         print(sort_order)
            bond_to_symbol[ii] = bond_to_symbol[ii][sort_order]
            bond_to_index[ii] = bond_to_index[ii][sort_order]
    #     print(ii, atom_list[ii], bond_to_symbol[ii], bond_to_index[ii])
        bond_to_symbol[ii] = list(bond_to_symbol[ii])
        bond_to_index[ii] = list(bond_to_index[ii])

    # Determine its atom type from bonds
    for ii in range(len(atom_list)):
        if atom_in_residue[ii] in protein_residue_list: # this is protein atom, treat with protein rules
            if atom_list[ii] == 'H':
                if bond_to_symbol[ii] == ['O']:        # This is a hydroxyl H, generate donor
                    atom_type[ii] = 'Hydroxyl H'
                    atom_type[ii] = 'OH'
                elif bond_to_symbol[ii] == ['N']:      # This is amino H, generate donor
                    atom_type[ii] = 'Amino H'
                    atom_type[ii] = 'NH'
                elif bond_to_symbol[ii] == ['C']:      # This is aliphatic H
                    atom_type[ii] = 'Aliphatic H'
                    atom_type[ii] = 'CH'
                elif bond_to_symbol[ii] == ['S']:      # This is thiol H, generate donor
                    atom_type[ii] = 'Thiol H'
                    atom_type[ii] = 'SH'
            elif atom_list[ii] == 'O':
                if bond_to_symbol[ii] == ['C']:        # This is carbonyl, generate two acceptor
                    atom_type[ii] = 'Carbonyl O'
                    atom_type[ii] = 'CO'
                    if bond_to_symbol[bond_to_index[ii][0]] == ['C','O','O']: # Carboxylate
                        atom_type[ii] = 'Carboxylate O'
                        atom_type[ii] = 'cCO'
                elif bond_to_symbol[ii] == ['C','H']:    # This is hydroxyl O, generate two acceptor, one donor
                    atom_type[ii] = 'Hydroxyl O'
                    atom_type[ii] = 'CHO'
                elif bond_to_symbol[ii] == ['H','H']:    # This is water O, generate two acceptor, two donor
                    atom_type[ii] = 'Water O'
                    atom_type[ii] = 'HHO'
            elif atom_list[ii] == 'N':
                if bond_to_symbol[ii] == ['C','C','H']:    # This is either the amide N (one donor), NME N, HIS N
                    atom_type[ii] = 'Amino N 1 donor'
                    atom_type[ii] = 'CCHN'
                elif bond_to_symbol[ii] == ['C','H','H']:    # This is ASN N, GLN N, or ARG N. Generate 2 donors each
                    atom_type[ii] = 'Amino N 2'
                    atom_type[ii] = 'CCHN'
                elif bond_to_symbol[ii] == ['C','C']:    # This is imine N (in Histidine and heterocycle cpds), one acceptor
                    atom_type[ii] = 'Imine N'
                    atom_type[ii] = 'CCN'
                elif bond_to_symbol[ii] == ['C','C','C']: # This is proline N. It has no H-bond
                    atom_type[ii] = 'Proline N'
                    atom_type[ii] = 'CCCN'
            elif atom_list[ii] == 'S':
                if bond_to_symbol[ii] == ['C','H']:    # This is thiol S, maybe generate acceptors??
                    atom_type[ii] = 'Thiol S'
                    atom_type[ii] = 'CHS'
                elif bond_to_symbol[ii] == ['C','C']:    # This is thioester S, maybe generate acceptors??
                    atom_type[ii] = 'Thioester S'
                    atom_type[ii] = 'CCS'
    #         if atom_list[ii] != 'C' and atom_type[ii] == []:
    #             print(ii, atom_list[ii], atom_type[ii], bond_to_symbol[ii])
        else: # Ligand residues
            if atom_list[ii] == 'H':
                if bond_to_symbol[ii] == ['O']:        # This is a hydroxyl H, generate donor
                    atom_type[ii] = 'Hydroxyl H'
                    atom_type[ii] = 'OH'
                elif bond_to_symbol[ii] == ['N']:      # This is amino H, generate donor
                    atom_type[ii] = 'Amino H'
                    atom_type[ii] = 'NH'
                elif bond_to_symbol[ii] == ['C']:      # This is aliphatic H
                    atom_type[ii] = 'Aliphatic H'
                    atom_type[ii] = 'CH'
                elif bond_to_symbol[ii] == ['S']:      # This is thiol H, generate donor
                    atom_type[ii] = 'Thiol H'
                    atom_type[ii] = 'SH'
                elif bond_to_symbol[ii] == ['P']:      # This is phosphine H, generate donor
                    atom_type[ii] = 'Phosphine H'
                    atom_type[ii] = 'PH'

            elif atom_list[ii] == 'O':
                if bond_to_symbol[ii] == ['C']:        # This is carbonyl, generate two acceptor
                    atom_type[ii] = 'Carbonyl O'
                    atom_type[ii] = 'CO'
                    if bond_to_symbol[bond_to_index[ii][0]] == ['C','O','O']: # Carboxylate group
                        atom_type[ii] = 'Carboxylate O'
                        atom_type[ii] = 'cCO'
                if bond_to_symbol[ii] == ['N']:        # This is nitro, generate one acceptor
                    atom_type[ii] = 'Nitro O'
                    atom_type[ii] = 'NO'
                elif bond_to_symbol[ii] == ['S']:        # This is carbonyl, generate two acceptor
                    atom_type[ii] = 'Sulfonyl / Sulfoxide O'
                    atom_type[ii] = 'SO'
                elif bond_to_symbol[ii] == ['P']:        # This is carbonyl, generate two acceptor
                    atom_type[ii] = 'Phosphine oxide / Phosphone O'
                    atom_type[ii] = 'PO'
                elif bond_to_symbol[ii] == ['C','H']:    # This is hydroxyl O, generate two acceptor, one donor
                    atom_type[ii] = 'Hydroxyl O'
                    atom_type[ii] = 'CHO'
                elif bond_to_symbol[ii] == ['C','C']:    # This is ether, generate one acceptor
                    atom_type[ii] = 'Ether O'
                    atom_type[ii] = 'CCO'
                elif bond_to_symbol[ii] == ['H','H']:    # This is water O, generate two acceptor, two donor
                    atom_type[ii] = 'Water O'
                    atom_type[ii] = 'HHO'
                elif bond_to_symbol[ii] == ['H','O']:    # This is water O, generate two acceptor, two donor
                    atom_type[ii] = 'Peroxide O'
                    atom_type[ii] = 'HOO'
                elif bond_to_symbol[ii] == ['C','O']:    # This is water O, generate two acceptor, two donor
                    atom_type[ii] = 'Peroxide O'
                    atom_type[ii] = 'COO'
            elif atom_list[ii] == 'N':
                if bond_to_symbol[ii] == ['C']:          # Cyanate / nitrile. Nitrile is an acceptor but cyanate is not
                    atom_type[ii] = 'Nitrile N'          # Default is nitrile unless the connecting C bonds to an O
                    atom_type[ii] = 'nCN'
                    for temp in bond_to_index[ii]:
                        if bond_to_symbol[temp] == 'O':
                            atom_type[ii] = 'Cyanate N'  # in which case this is an cyanate
                            atom_type[ii] = 'cCN'

                if bond_to_symbol[ii] == ['N']:          # Azide
                    atom_type[ii] = 'Azide N, acceptor'
                    atom_type[ii] = 'NN'
                if bond_to_symbol[ii] == ['C','C']:      # Isonitrile / sp2 nitrogen
                    atom_type[ii] = 'sp2 N acceptor'              # Default is sp2 nitroten
                    atom_type[ii] = 'aCCN'
                    for temp in bond_to_index[ii]:
                        if len(bond_to_symbol[temp]) == 1:
                            atom_type[ii] = 'Isonitrile N'  # If there's an sp carbon, than this is an isonitrile N
                            atom_type[ii] = 'iCCN'

                if bond_to_symbol[ii] == ['C','H']:      # Primary ketimine / aldimine
                    atom_type[ii] = 'Ketimine / aldimine N'
                    atom_type[ii] = 'CHN'
                if bond_to_symbol[ii] == ['C','N']:      # Azo / Azide
                    atom_type[ii] = 'Azo / Azide starting N'
                    atom_type[ii] = 'CNN'
                if bond_to_symbol[ii] == ['C','O']:      # Nitroso
                    atom_type[ii] = 'Nitroso / Oxime N'
                    atom_type[ii] = 'CON'
                if bond_to_symbol[ii] == ['N','N']:      # Azide
                    atom_type[ii] = 'Azide N, base of acceptor'
                    atom_type[ii] = 'NNN'
                if bond_to_symbol[ii] == ['O','O']:      # Nitrite
                    atom_type[ii] = 'Nitrite N'
                    atom_type[ii] = 'OON'
                if bond_to_symbol[ii] == ['C','C','C']:      #
                    atom_type[ii] = 'sp3 N'
                    atom_type[ii] = 'aCCCN'
                    for temp in bond_to_index[ii]:
                        if len(bond_to_symbol[temp]) == 3:
                            atom_type[ii] = 'sp2 N'
                            atom_type[ii] = 'iCCCN'

                if bond_to_symbol[ii] == ['C','C','H']:      #
                    atom_type[ii] = 'sp3 N'
                    atom_type[ii] = 'adCCHN'
                    for temp in bond_to_index[ii]:
                        if len(bond_to_symbol[temp]) == 3:
                            atom_type[ii] = 'sp2 N donor'
                            atom_type[ii] = 'dCCHN'
                if bond_to_symbol[ii] == ['C','H','H']:      #
                    atom_type[ii] = 'sp3 N'
                    atom_type[ii] = 'adCHHN'
                    for temp in bond_to_index[ii]:
                        if len(bond_to_symbol[temp]) == 3:
                            atom_type[ii] = 'sp2 N donors'
                            atom_type[ii] = 'dCHHN'

                if bond_to_symbol[ii] == ['H','H','H']:      #
                    atom_type[ii] = 'Ammonia N, donors + acceptors'
                    atom_type[ii] = 'HHHN'
                if bond_to_symbol[ii] == ['C','O','O']:      #
                    atom_type[ii] = 'Nitro N'
                    atom_type[ii] = 'COON'
                if bond_to_symbol[ii] == ['O','O','O']:      #
                    atom_type[ii] = 'Nitrate / Nitric acid'
                    atom_type[ii] = 'OOON'
                if bond_to_symbol[ii] == ['C','C','C','C']:      #
                    atom_type[ii] = '4 deg ammonium'
                    atom_type[ii] = 'CCCCN'
                if bond_to_symbol[ii] == ['H','H','H','H']:      #
                    atom_type[ii] = 'Ammonium'
                    atom_type[ii] = 'HHHHN'

            elif atom_list[ii] == 'S':
                if bond_to_symbol[ii] == ['C']:        # This is thioketone, generate two acceptor
                    atom_type[ii] = 'Thioketone S / Isothiocyanite S / Thial S '
                    atom_type[ii] = 'CS'
                elif bond_to_symbol[ii] == ['C','H']:    # This is thiol S, maybe generate acceptors??
                    atom_type[ii] = 'Thiol S'
                    atom_type[ii] = 'CHS'
                elif bond_to_symbol[ii] == ['C','C']:    # This is thioester S, maybe generate acceptors??
                    atom_type[ii] = 'Thioether S / Thiocyanate S'
                    atom_type[ii] = 'CCS'
                elif bond_to_symbol[ii] == ['C','S']:    # Disulfide
                    atom_type[ii] = 'Disulfide S'
                    atom_type[ii] = 'CSS'
                elif bond_to_symbol[ii] == ['C','C','O']: # Sulfoxide
                    atom_type[ii] = 'Sulfoxide S'
                    atom_type[ii] = 'CCOS'
                elif bond_to_symbol[ii] == ['C','C','O','O']: # Sulfone
                    atom_type[ii] = 'Sulfone S'
                    atom_type[ii] = 'CCOOS'
                elif bond_to_symbol[ii] == ['C','O','O']: # Sulfinic acid
                    atom_type[ii] = 'Sulfinic acid S'
                    atom_type[ii] = 'COOS'
                elif bond_to_symbol[ii] == ['C','O','O','O']: # Sulfonic acid / Ester
                    atom_type[ii] = 'Sulfonic acid / ester S'
                    atom_type[ii] = 'COOOS'
            elif atom_list[ii] == 'P':
                if bond_to_symbol[ii] == ['C','C','C']:
                    atom_type[ii] = 'CCCP'
                if bond_to_symbol[ii] == ['C','C','H']:
                    atom_type[ii] = 'CCCP'
                if bond_to_symbol[ii] == ['C','H','H']:
                    atom_type[ii] = 'CCCP'
                if bond_to_symbol[ii] == ['H','O','O','O']:
                    atom_type[ii] = 'HOOOP'
                if bond_to_symbol[ii] == ['O','O','O','O']:
                    atom_type[ii] = 'OOOOP'

    if printTiming:
        t1 = time.time()
        print(f'Parse atom types: {t1-t0:.3f} s')
        t0 = time.time()

    if processFrame == -1:
        frames = range(len(comp))
    elif type(processFrame) == int:
        frames = [processFrame]
    elif type(processFrame) == list:
        frames = processFrame

    for frame in frames:
        NOS_coordinate = []
        H_coordinate = []
        ALL_coordinate = comp.xyz[frame]
        NOS_coordinate = comp.xyz[frame][NOS_index]
    #    print(NOS_coordinate.shape)
        H_coordinate = comp.xyz[frame][H_index]
    #    NOS_coordinate = np.array(NOS_coordinate)
    #    H_coordinate = np.array(H_coordinate)

        # Donor generation
        DOL = [] # Donor Oxygen Location
        DOL_index = []
        for ii in range(len(atom_list)):
            if atom_type[ii] in ['OH', 'NH', 'SH']:
                vec = comp.xyz[frame][ii] - comp.xyz[frame][bond_to_index[ii][0]]
                vec /= np.linalg.norm(vec)
                vec = vec.flatten()
                if atom_type[ii] == 'OH':
                    vec *= 0.18
                elif atom_type[ii] == 'NH':
                    vec *= 0.2
                elif atom_type[ii] == 'SH':
                    vec *= 0.245  # Temporary
                DOL.append(comp.xyz[frame][ii] + vec)
                DOL_index.append(ii)
        # print(np.array(DOL))
        DOL = np.array(DOL)
        DOL_index = np.array(DOL_index)
        # # Write out a test xyz
        # with open('test_donor.pdb','w') as f:
        # #     f.write(f'{len(DOL)}\n\n')
        #     for ii in range(len(DOL)):
        # #         ATOM      1  N   PRO A   1       8.316  21.206  21.530  1.00 17.44           N
        #         f.write(f'ATOM    {ii+1:3d}  X   XXX A{ii:4d}    {DOL[ii][0]*10:8.3f}{DOL[ii][1]*10:8.3f}{DOL[ii][2]*10:8.3f}  0.00  0.00            \n')
        #     f.write('END\n')
        #

        # Acceptor generation
        counter = 0
        AHL = [] # Acceptor hydrogen Location
        AHL_index = []
        for ii in range(len(atom_list)):
            if atom_type[ii] in ['CO', 'cCO']: # Carbonyl
        #         counter += 1
                vec1 = comp.xyz[frame][ii] - comp.xyz[frame][bond_to_index[ii][0]]
                # print(ii)
                # print(bond_to_index[bond_to_index[ii][0]])
                for temp in bond_to_index[bond_to_index[ii][0]]:
                    if temp != ii:
                        vec2 = comp.xyz[frame][temp] - comp.xyz[frame][bond_to_index[ii][0]]
                        rot_axis = np.cross(vec1.flatten(), vec2.flatten())
                        # print(rot_axis)
                        break
                vec = rotation3D(vec1, rot_axis, 60)
                AHL.append(comp.xyz[frame][ii] + prepVec(vec, 0.194))
                AHL_index.append(ii)
                vec = rotation3D(vec1, rot_axis, -60)
                AHL.append(comp.xyz[frame][ii] + prepVec(vec, 0.194))
                AHL_index.append(ii)
            if atom_type[ii] in ['NO']: # Nitro, very similar to CO
        #         counter += 1
                vec1 = comp.xyz[frame][ii] - comp.xyz[frame][bond_to_index[ii][0]]
                # print(ii)
                # print(bond_to_index[bond_to_index[ii][0]])
                for temp in bond_to_index[bond_to_index[ii][0]]:
                    if temp != ii:
                        vec2 = comp.xyz[frame][temp] - comp.xyz[frame][bond_to_index[ii][0]]
                        rot_axis = np.cross(vec1.flatten(), vec2.flatten())
                        # print(rot_axis)
                        break
                vec = rotation3D(vec1, rot_axis, 60)
                AHL.append(comp.xyz[frame][ii] + prepVec(vec, 0.233))
                AHL_index.append(ii)
                vec = rotation3D(vec1, rot_axis, -60)
                AHL.append(comp.xyz[frame][ii] + prepVec(vec, 0.233))
                AHL_index.append(ii)
        #     if atom_type[ii] in ['cCO']: # Carboxylate
        #         vec = comp.xyz[frame][bond_to_index[ii][0]] - comp.xyz[frame][bond_to_index[bond_to_index[ii][0]][0]]
        #         AHL.append(comp.xyz[frame][ii] + prepVec(vec, 0.192))
        #         AHL_index.append(ii)
            if atom_type[ii] in ['CHO']: # Hydroxyl
                vec1 = comp.xyz[frame][bond_to_index[ii][1]] - comp.xyz[frame][ii] # OH
                vec2 = comp.xyz[frame][ii] - comp.xyz[frame][bond_to_index[ii][0]] # CH
                vec = rotation3D(vec1, vec2, 180)
                AHL.append(comp.xyz[frame][ii] + prepVec(vec, 0.203))
                AHL_index.append(ii)

            if (atom_type[ii] in ['aCCN','CCN']): # mostly histidine and nitrogens in heterocycles
                counter += 1
        #         vec = comp.xyz[frame][ii] - comp.xyz[frame][bond_to_index[ii]]
                # print(ii)
                # print(bond_to_index[bond_to_index[ii][0]])
                avgCC = np.mean(comp.xyz[frame][bond_to_index[ii]],axis=0)
                vec = comp.xyz[frame][ii] - avgCC
                AHL.append(comp.xyz[frame][ii] + prepVec(vec, 0.21))
                AHL_index.append(ii)
            if atom_type[ii] in ['CCO']: # Ether
                counter += 1
        #         vec = comp.xyz[frame][ii] - comp.xyz[frame][bond_to_index[ii]]
                # print(ii)
                # print(bond_to_index[bond_to_index[ii][0]])
                avgCC = np.mean(comp.xyz[frame][bond_to_index[ii]],axis=0)
                vec = comp.xyz[frame][ii] - avgCC
                AHL.append(comp.xyz[frame][ii] + prepVec(vec, 0.21))
                AHL_index.append(ii)
            if atom_type[ii] in ['SO']: # Sulfoxide
                vec = comp.xyz[frame][ii] - comp.xyz[frame][bond_to_index[ii][0]]
                AHL.append(comp.xyz[frame][ii] + prepVec(vec, 0.196))
                AHL_index.append(ii)
            if atom_type[ii] in ['PO']: # Phosphine
                vec = comp.xyz[frame][ii] - comp.xyz[frame][bond_to_index[ii][0]]
                AHL.append(comp.xyz[frame][ii] + prepVec(vec, 0.188))
                AHL_index.append(ii)
            if atom_type[ii] in ['nCN']: # Nitrile
                vec = comp.xyz[frame][ii] - comp.xyz[frame][bond_to_index[ii][0]]
                AHL.append(comp.xyz[frame][ii] + prepVec(vec, 0.21))
                AHL_index.append(ii)
            if atom_type[ii] in ['aCCCN', 'adCCHN', 'adCHHN', 'HHHN']: # Various acceptors for amines
                avgCC = np.mean(comp.xyz[frame][bond_to_index[ii]],axis=0)
                vec = comp.xyz[frame][ii] - avgCC
                AHL.append(comp.xyz[frame][ii] + prepVec(vec, 0.213))
                AHL_index.append(ii)

        AHL = np.array(AHL)
        AHL_index = np.array(AHL_index)

        # # Write out a test xyz
        # with open('test_acceptor.pdb','w') as f:
        # #     f.write(f'{len(DOL)}\n\n')
        #     for ii in range(len(AHL)):
        # #         ATOM      1  N   PRO A   1       8.316  21.206  21.530  1.00 17.44           N
        #         f.write(f'ATOM    {ii+1:3d}  Z   ZZZ A{ii:4d}    {AHL[ii][0]*10:8.3f}{AHL[ii][1]*10:8.3f}{AHL[ii][2]*10:8.3f}  0.00  0.00            \n')
        #     f.write('END\n')
        #

        if printTiming:
            t1 = time.time()
            print(f'Generate donors / acceptors: {t1-t0:.3f} s')
            t0 = time.time()

        # Calculate DOL -> N, O, S distance
        DOL_NOS_dist = np.min(np.sqrt(((DOL[:,:,None] - NOS_coordinate[:,:,None].T)**2).sum(1)),axis=1)
        # print(DOL_NOS_dist.shape)
        AHL_H_dist = np.min(np.sqrt(((AHL[:,:,None] - H_coordinate[:,:,None].T)**2).sum(1)),axis=1)
        # print(AHL_H_dist.shape)
        # import matplotlib.pyplot as plt
        # plt.hist(DOL_NOS_dist, bins=30)
        # plt.hist(AHL_H_dist, bins=30)
        # np.nonzero(AHL_H_dist > 0.1)[0]
        AHL_inter = AHL[np.nonzero(AHL_H_dist > 0.12)[0]]
        AHL_inter_index = AHL_index[np.nonzero(AHL_H_dist > 0.12)[0]] # Atom indices of acceptor atoms
        DOL_inter = DOL[np.nonzero(DOL_NOS_dist > 0.12)[0]]
        DOL_inter_index = DOL_index[np.nonzero(DOL_NOS_dist > 0.12)[0]]
        # outputMockPDB('test_AHL_inter.pdb',AHL_inter)
        # outputMockPDB('test_DOL_inter.pdb',DOL_inter)

        # Calculate inter-donor/acceptor distances
        DOL_DOL = np.sqrt(((DOL_inter[:,:,None] - DOL_inter[:,:,None].T)**2).sum(1))
        np.fill_diagonal(DOL_DOL, 1)
        AHL_DOL = np.sqrt(((AHL_inter[:,:,None] - DOL_inter[:,:,None].T)**2).sum(1))
        # np.fill_diagonal(AHL_DOL, 1)
        AHL_AHL = np.sqrt(((AHL_inter[:,:,None] - AHL_inter[:,:,None].T)**2).sum(1))
        np.fill_diagonal(AHL_AHL, 1)

        # Qualify donor/acceptors based on distances
        DOL_DOL_q = np.where(DOL_DOL < 0.14)
        AHL_DOL_q = np.where(AHL_DOL < 0.28)
        AHL_AHL_q = np.where(AHL_AHL < 0.28)

        if printTiming:
            t1 = time.time()
            print(f'Qualify donors / acceptors: {t1-t0:.3f} s')
            t0 = time.time()

        # Place water molecules. Note this is ordered - AHL+AHL takes priority.
        # This will happen in redundant water elimination and hydrogen placements

        temp_water = []
        temp_water_index = []
        temp_water_type = []
        for ii in range(len(AHL_AHL_q[0])):
            temp_water.append((AHL_inter[AHL_AHL_q[0][ii]]+AHL_inter[AHL_AHL_q[1][ii]])/2)
            temp_water_index.append([AHL_inter_index[AHL_AHL_q[0][ii]], AHL_inter_index[AHL_AHL_q[1][ii]]])
            temp_water_type.append(0)
        for ii in range(len(AHL_DOL_q[0])):
            temp_water.append((AHL_inter[AHL_DOL_q[0][ii]]+DOL_inter[AHL_DOL_q[1][ii]])/2)
            temp_water_index.append([AHL_inter_index[AHL_DOL_q[0][ii]], DOL_inter_index[AHL_DOL_q[1][ii]]])
            temp_water_type.append(1)
        for ii in range(len(DOL_DOL_q[0])):
            temp_water.append((DOL_inter[DOL_DOL_q[0][ii]]+DOL_inter[DOL_DOL_q[1][ii]])/2)
            temp_water_index.append([DOL_inter_index[DOL_DOL_q[0][ii]], DOL_inter_index[DOL_DOL_q[1][ii]]])
            temp_water_type.append(2)
        temp_water = np.array(temp_water)
        temp_water_index = np.array(temp_water_index)
        temp_water_type = np.array(temp_water_type)

        # outputMockPDB('temp_water.pdb',temp_water,e='O')

        if printTiming:
            t1 = time.time()
            print(f'Place template oxygens: {t1-t0:.3f} s')
            t0 = time.time()

        # Test for waters that are too close
        ALL_WAT = np.sqrt(((ALL_coordinate[:,:,None] - temp_water[:,:,None].T)**2).sum(1)) - ALL_vdW_radii[:,None] - 0.06

        qualified_water = temp_water[np.where(np.all(ALL_WAT > 0, axis=0))]
        qualified_water_index = temp_water_index[np.where(np.all(ALL_WAT > 0, axis=0))]
        qualified_water_type = temp_water_type[np.where(np.all(ALL_WAT > 0, axis=0))]

        # outputMockPDB('qualified_water.pdb',qualified_water,e='O')

        if printTiming:
            t1 = time.time()
            print(f'Place qualified oxygens: {t1-t0:.3f} s')
            t0 = time.time()

        # Orient each qualifying water
        full_water = []
        full_water_index = []
        for ii in range(len(qualified_water)):
            this_water = np.array([qualified_water[ii], [0,0,0], [0,0,0]])
            if qualified_water_type[ii] == 0: # A two donor water - H's of water point to acceptors
                this_water[1] = this_water[0] + prepVec(comp.xyz[frame][qualified_water_index[ii][0]] - this_water[0], 0.09572) # Acceptor 1
                this_water[2] = this_water[0] + prepVec(comp.xyz[frame][qualified_water_index[ii][1]] - this_water[0], 0.09572) # Acceptor 2
                # Adjust angles to 104.5 deg
                thisAngle, vec1, vec2 = calcAngle(comp.xyz[frame][qualified_water_index[ii][0]], this_water[0], comp.xyz[frame][qualified_water_index[ii][1]])
                thisAxis = np.cross(vec1, vec2)
                this_water[1] = this_water[0] + rotation3D(this_water[1] - this_water[0], thisAxis, (thisAngle - 104.5)/2)
                this_water[2] = this_water[0] + rotation3D(this_water[2] - this_water[0], thisAxis, -(thisAngle - 104.5)/2)
            if qualified_water_type[ii] == 1: # A acceptor / donor water
                this_water[1] = this_water[0] + prepVec(comp.xyz[frame][qualified_water_index[ii][0]] - this_water[0], 0.09572) # Acceptor 1
                # Create second hydrogen and rotate that away (so the lone pair faces the donor)
                thisAngle, vec1, vec2 = calcAngle(comp.xyz[frame][qualified_water_index[ii][0]], this_water[0], comp.xyz[frame][qualified_water_index[ii][1]])
                thisAxis = np.cross(vec1, vec2)
                this_water[2] = this_water[0] + rotation3D(this_water[1] - this_water[0], thisAxis, 104.5)
                this_water[2] = this_water[0] + rotation3D(this_water[2] - this_water[0], this_water[1] - this_water[0], 120)
                #parallel_NH = comp.xyz[frame][qualified_water_index[ii][1]] - comp.xyz[frame][bond_to_index[qualified_water_index[ii][1]]]
                #this_water[2] = this_water[0] + prepVec(parallel_NH, 0.09572) # Acceptor 2
            full_water.append(this_water)
            full_water_index.append(qualified_water_index[ii])
        # print(full_water_index)

        # outputMolecularMockPDB('Full_water.pdb', full_water, e=['O','H','H'])

        if printTiming:
            t1 = time.time()
            print(f'Place full waters: {t1-t0:.3f} s')
            t0 = time.time()

        # Geometry check
        geo_water = []
        geo_water_O = []
        geo_water_index = []
        lowerLength = {'O':0.161, 'N': 0.171, 'S':0.229}
        upperLength = {'O':0.230, 'N': 0.240, 'S':0.312}
        for ii in range(len(full_water)):
            # print(f'WATER {ii}')
            if qualified_water_type[ii] == 0: # Test two distances and two angles
                testLength = np.linalg.norm(full_water[ii][1] - comp.xyz[frame][qualified_water_index[ii][0]])
                # print(testLength)
        #         if atom_list[qualified_water_index[ii][0]] in ['O','S','N']:
                if (testLength < lowerLength[atom_list[qualified_water_index[ii][0]]]) or \
                   (testLength > upperLength[atom_list[qualified_water_index[ii][0]]]):
                    # print(f'Water {ii} failed due to distance between H1 and acceptor')
                    continue
                testLength = np.linalg.norm(full_water[ii][2] - comp.xyz[frame][qualified_water_index[ii][1]])
                # print(testLength)
        #         if atom_list[qualified_water_index[ii][0]] in ['O','S','N']:
                if (testLength < lowerLength[atom_list[qualified_water_index[ii][1]]]) or \
                   (testLength > upperLength[atom_list[qualified_water_index[ii][1]]]):
                    # print(f'Water {ii} failed due to distance between H2 and acceptor')
                    continue
                testAngle, _, _ = calcAngle(full_water[ii][0], full_water[ii][1], comp.xyz[frame][qualified_water_index[ii][0]])
                # print(testAngle)
                if testAngle < 120:
                    # print(f'Water {ii} failed due to angle between O, H1 and acceptor')
                    continue
                testAngle, _, _ = calcAngle(full_water[ii][0], full_water[ii][2], comp.xyz[frame][qualified_water_index[ii][1]])
                # print(testAngle)
                if testAngle < 120:
                    # print(f'Water {ii} failed due to angle between O, H2 and acceptor')
                    continue
            if qualified_water_type[ii] == 1:
                # Test 1. distance between OH and acceptor,
                #      2. angle between (donor, donor H, O), and
                #      3. angle between (accpetor expected H, acceptor, H)
                testLength = np.linalg.norm(full_water[ii][1] - comp.xyz[frame][qualified_water_index[ii][0]])
                # print(testLength)

                if (testLength < lowerLength[atom_list[qualified_water_index[ii][0]]]) or \
                   (testLength > upperLength[atom_list[qualified_water_index[ii][0]]]):
                    # print(f'Water {ii} failed due to distance between H1 and acceptor')
                    continue
                # print(f'Qualified H, {qualified_water_index[ii][1]} (H) bonds to {bond_to_index[qualified_water_index[ii][1]][0]}')
                testAngle, _, _ = calcAngle(comp.xyz[frame][bond_to_index[qualified_water_index[ii][1]][0]],
                                            comp.xyz[frame][qualified_water_index[ii][1]],
                                            full_water[ii][0])
                if testAngle < 120:
                    # print(f'Water {ii} failed due to angle between donor, donor H and water O')
                    continue
                for jj in np.where(AHL_index == qualified_water_index[ii][0])[0]:
                    testAngle, _, _ = calcAngle(AHL[jj],
                                                comp.xyz[frame][qualified_water_index[ii][0]],
                                                full_water[ii][1])
                    if testAngle < 45:
        #                 print(f'passed angle test')
                        break
        #             print(f'jjjj {testAngle}')
                else:
                    continue
            geo_water.append(full_water[ii])
            geo_water_O.append(full_water[ii][0])
            geo_water_index.append(full_water_index[ii])
        geo_water = np.array(geo_water)
        geo_water_O = np.array(geo_water_O)
        geo_water_index = np.array(geo_water_index)

        # print(geo_water_index)

        if printTiming:
            t1 = time.time()
            print(f'Water geometry check: {t1-t0:.3f} s')
            t0 = time.time()

        # outputMolecularMockPDB('Geo_water.pdb', geo_water, e=['O','H','H'])

        # Eliminate duplicate O
        accept_water = []
        GEO_GEO = np.sqrt(((geo_water_O[:,:,None] - geo_water_O[:,:,None].T)**2).sum(1))
        np.fill_diagonal(GEO_GEO, 1)
        # print(GEO_GEO)
        for ii in range(len(GEO_GEO)):
            if np.any(GEO_GEO[ii][:ii] < 0.249):
                GEO_GEO[ii] = 1
                GEO_GEO[:,ii] = 1
            else:
                accept_water.append(ii)
        # print(GEO_GEO)
        # print(accept_water)
        final_water = geo_water[accept_water]
        final_water_O = geo_water_O[accept_water]
        final_water_index = geo_water_index[accept_water]

        if returnHBscore:
            HB_score = 0
            for ii in range(len(final_water)):

                atom1 = comp.top.atom(final_water_index[ii][0])
                atom2 = comp.top.atom(final_water_index[ii][1])

                if atom1.residue.name == 'LIG' or atom2.residue.name == 'LIG':
                    if atom1.residue.name == 'LIG' and atom2.residue.name == 'LIG':
                        pass
                    else:
                        if atom1.element.symbol == 'H':
                            # print(atom1)
                            # print(comp.top.atom(bond_to_index[final_water_index[ii][0]][0]), end=' \t')
                            btom1_index = bond_to_index[final_water_index[ii][0]][0]
                        else:
                            # print(atom1, end=' \t')
                            btom1_index = final_water_index[ii][0]
                        HB_score += BWHB(ALL_coordinate[btom1_index], final_water_O[ii], print_dist=False)
                        # print(f'{HB(ALL_coordinate[btom1_index], final_water_O[ii], print_dist=False):.3f}')
                        if atom2.element.symbol == 'H':
                            # print(comp.top.atom(bond_to_index[final_water_index[ii][1]][0]), end=' \t')
                            btom2_index = bond_to_index[final_water_index[ii][1]][0]
                        else:
                            # print(atom2, end=' \t')
                            btom2_index = final_water_index[ii][1]
                        HB_score += BWHB(ALL_coordinate[btom2_index], final_water_O[ii], print_dist=False)
                        # print(f'{HB(ALL_coordinate[btom2_index], final_water_O[ii], print_dist=True):.3f}')
            try:
                HB_score_list.append(HB_score)
            except:
                HB_score_list = [HB_score]
        # print()

        # print(final_water_index)

        if printTiming:
            t1 = time.time()
            print(f'Final water determination: {t1-t0:.3f} s')
            t0 = time.time()

        if outFile is not None:
            outputMolecularMockPDB(outFile, final_water, e=['O','H1','H2'], resname=outResname)

        if printTiming:
            t1 = time.time()
            print(f'Output final file: {t1-t0:.3f} s')
            t0 = time.time()

        if printResult:
            print(f'For this input we can place {len(final_water)} waters')

        if returnNumWater:
            try:
                numWater.append(len(final_water))
            except:
                numWater = [len(final_water)]

    if returnNumWater:
        return numWater

    if returnHBscore:
        return HB_score_list


import math
import copy

import matplotlib.pyplot as plt
import numpy as np

import Solution as sl
import Problem
import ParetoUtil as pu

class MOPSO():
    def __init__(self, problem, popSize: int, repSize: int):
        super().__init__()
        self.problem = problem
        self.popSize = popSize
        self.repSize = int(repSize)
        self.numberOfHyperCubesPerRow = 5

    def Evolve(self, genertionCount: int, initPopulation: np.array, mutationRate, crossoverRate):
        # 1, 2, 3, 6
        population = copy.deepcopy(initPopulation)
        repository = None

        # 4 Add non-dominated front of population to repository
        repository = pu.fastNonDominatedSord(population)[0]

        # remove dominated ones from repository
        repository = pu.fastNonDominatedSord(repository)[0]

        # remove extra based on reverse hypercube crowdness
        if(np.size(repository)) > self.repSize:
            repository = MOPSO.sortBasedOnHyperCubeCrowd(
                repository, self.numberOfHyperCubesPerRow)
            repository = repository[0: self.repSize]

        for g in range(genertionCount):
            # calculate the velocity
            w = .4
            for i in range(self.popSize):
                R1, R2 = np.random.rand(), np.random.rand()
                h = MOPSO.getSolutionIndexFromCrowdCube(
                    repository, self.numberOfHyperCubesPerRow)
                population[i].velocity = (w * population[i].velocity) +\
                                         (R1 * (population[i].Pbest.gene - population[i].gene)) +\
                                         (repository[h].gene - population[i].gene)

            # calculate new location
            for i in range(self.popSize):
                population[i].gene = population[i].gene + \
                    population[i].velocity
                # Apply varialble constrain after position update
                population[i].ApplyVarialbleConstrains()
                # update Pbest
                if (population[i].Dominate(population[i].Pbest)):
                    population[i].Pbest=copy.deepcopy(population[i].gene)

            # Add non-dominated front of population to repository with duplicate remove
            nonDominatedPop = pu.fastNonDominatedSord(population)[0]
            for i in range(np.size(nonDominatedPop)):
                isRepeated = False
                for j in range(np.size(repository)):
                    if nonDominatedPop[i] is repository[j]: isRepeated = True
                if isRepeated == False:
                    repository=sl.solution.AddPopulations(repository, nonDominatedPop[i])

            # remove dominated ones from repository
            repository= pu.fastNonDominatedSord(repository)[0]

            # remove extra based on reverse hyperCube crowdness
            if(np.size(repository)) > self.repSize:
                repository=MOPSO.sortBasedOnHyperCubeCrowd(
                    repository, self.numberOfHyperCubesPerRow)
                repository=repository[0:self.repSize]

            print("MSPSO generation" + str(g))

        # get the final Pareto front
        return repository

    @staticmethod
    def assignSolutionToHyperCube(solution, numberOfHyperCubesPerRow):
        hyperCubeWidth=1 / numberOfHyperCubesPerRow
        f1Cordination=int(solution.GetObjective(0) / hyperCubeWidth)
        f2Cordination=int(solution.GetObjective(1) / hyperCubeWidth)

        return (f2Cordination * numberOfHyperCubesPerRow) + f1Cordination

    @staticmethod
    def assignPopulationToHyperCube(population, numberOfHyperCubesPerRow):
        popSize=np.size(population)
        hyperCubeLocation=np.empty([popSize], dtype=int)
        for i in range(popSize):
            hyperCubeLocation[i]=MOPSO.assignSolutionToHyperCube(
                population[i], numberOfHyperCubesPerRow)
        return hyperCubeLocation

    @staticmethod
    def getSolutionIndexFromCrowdCube(population, numberOfHyperCubesPerRow):
        hyperCubeLocation=MOPSO.assignPopulationToHyperCube(
            population, numberOfHyperCubesPerRow)
        HyperCubeCrowdnessProbability=np.empty(
            [np.size(hyperCubeLocation)], dtype=float)
        # assign crowdness Probability
        for i in range(np.size(hyperCubeLocation)):
            HyperCubeCrowdnessProbability[i]=np.size(
                np.where(hyperCubeLocation == hyperCubeLocation[i]))
        sumOfCrowdness=np.sum(HyperCubeCrowdnessProbability)
        HyperCubeCrowdnessProbability=HyperCubeCrowdnessProbability/sumOfCrowdness
        # simulate the roulette wheel selection
        return np.random.choice(np.size(hyperCubeLocation), 1, p=HyperCubeCrowdnessProbability)[0]

    @staticmethod
    def sortBasedOnHyperCubeCrowd(population, numberOfHyperCubesPerRow):
        hyperCubeLocation=MOPSO.assignPopulationToHyperCube(
            population, numberOfHyperCubesPerRow)
        # find crowdest location
        largestCrowd=0
        crowdestLocation=0
        for i in range(np.size(population)):
            population[i].fitness=np.size(np.where(hyperCubeLocation == hyperCubeLocation[i]))

        _list = np.ndarray.tolist(population)
        _list.sort(key=lambda x: x.fitness, reverse=False)
        return np.asarray(_list, sl.solution)


# -*- coding: utf-8 -*-
"""
Created on 2023-11-28 (Tue) 15:03:39

Planar Maximally Filtered Graph implementation in python

@author: I.Azuma
"""
import numpy as np
import pandas as pd
import time
import networkx as nx
from networkx.algorithms.planarity import check_planarity
from tqdm import tqdm
import matplotlib.pyplot as plt

class GraphHandler:
    def __init__(self):
        self.X = np.array([[],[]])
        self.labels = dict()
        self.n_node = None
        self.n_edge = None
        self.graph = None
        self.centrality = dict()
        self.params = dict()

    def set_data(self,df:pd.DataFrame):
        """
        set adjucency matrix dataframe

        Parameters
        ----------
        df: dataframe
            a dataframe of adjucency matrix

        """
        idx = list(df.index)
        self.n_node = len(idx)
        self.labels = dict(zip(list(range(self.n_node)),idx))
        self.X = df.values

    def set_graph(self,graph):
        """
        set networkx.Graph object

        Parameters
        ----------
        graph: networkx.Graph

        """
        self.graph = graph
        self.n_edge = graph.number_of_edges()
        self.n_node = graph.number_of_nodes()

    def adj2graph(self,X,update:bool=True):
        """
        convert adjacency matrix to graph object of networkx

        Parameters
        ----------
        X: np.array
            adjacency matrix

        update: bool
            whether the stored graph is replaced by the generated one

        """
        n = X.shape[0]
        if self.X.shape[0]==0:
            self.X = X
        idx = list(range(n))
        edge = []
        ap = edge.append
        for i in range(n - 1):
            for j in range(i + 1,n):
                ap((idx[i],idx[j],X[i][j]))
        g = nx.Graph()
        g.add_weighted_edges_from(edge)
        if update:
            self.set_graph(g)
        return g

    def edge2graph(self,edge:pd.DataFrame,update:bool=True):
        """
        construct a graph from edge dataframe

        Parameters
        ----------
        edge: pd.DataFrame
            should be the form as follows:
                1st column: source node (note, src and dest is exchangeable)
                2nd column: destination node
                3rd column: weight

        update: bool
            whether the stored graph is replaced by the generated one

        """
        x = edge.values
        x = [(v[0],v[1],v[2]) for v in x]
        g = nx.Graph()
        g.add_weighted_edges_from(x)
        if update:
            self.set_graph(g)
        return g

    def graph2edge(self,graph,sorting:bool=False):
        """
        convert a networkx.Graph object to a dataframe of edge

        Parameters
        ----------
        graph: networkx.Graph object

        sorting: bool
            whether the result is sorted or not

        """
        edge = []
        ap = edge.append
        for v in graph.edges(data=True):
            ap([v[0],v[1],v[2]['weight']])
        df = pd.DataFrame(edge,columns=['source','dest','weight'])
        if sorting:
            df = df.sort_values(by='weight',ascending=False)
        return df

    def calc_centrality(self,graph,centrality='betweenness'):
        """
        calculate centrality of the given graph

        Parameters
        ----------
        graph: networkx.Graph object

        centrality: str
            indicates centrality method
            'betweenness', 'pagerank', 'degree', or 'closeness'

        """
        if centrality=='betweenness':
            self.centrality = nx.betweenness_centrality(graph)
        elif centrality=='degree':
            self.centrality = nx.degree_centrality(graph)
        elif centrality=='closness':
            self.centrality = nx.closeness_centrality(graph)
        elif centrality=='pagerank':
            self.centrality = nx.pagerank(graph)
        else:
            raise KeyError('!! Wrong centrality: use betweenness, pagerank, degree, or closeness !!')
        self.params['centrality'] = centrality
        return self.centrality

class PMFG(GraphHandler):
    def __init__(self):
        super().__init__()

    def pmfg(self,graph=None,fdr:float=0.05,update:bool=True,boyer=True):
        """
        obtain PMFG from the given graph

        Parameters
        ----------
        graph: networkx.Graph
            a networkx.Graph object

        fdr: float
            indicates the threshold for termination based on false discovery rate
            should be not 0

        update: bool
            whether the stored graph is replaced by the generated one

        """
        sta = time.time()
        if graph is None:
            if self.graph is None:
                if self.X is None:
                    raise ValueError('!! Provide graph as an argument or adjucency matrix by set_data !!')
                else:
                    graph = self.adj2graph(self.X,update=False)
            else:
                graph = self.graph
        n_node = graph.number_of_nodes()
        n_fdr = int(1/fdr)
        self.params['fdr'] = fdr
        graph = self._sort_graph(graph)
        nmax = 3*(n_node - 2)
        g = nx.Graph()
        rejected = 0
        if boyer:
            print("Boyer's method")
            for e in tqdm(graph):
                g.add_edge(e['source'],e['dest'],weight=e['weight'])
                if check_planarity(g)[0]:
                    n_edge = g.number_of_edges()
                    if n_edge==nmax:
                        print('--- terminated (#edge={} reached the max) ---'.format(n_edge))
                        break
                    else:
                        rejected = 0
                else:
                    g.remove_edge(e['source'],e['dest'])
                    rejected += 1
                    if rejected > n_fdr*n_node: # equivalent to FDR
                        print('--- terminated (#rejected={} reached the FDR threshold) ---'.format(rejected))
                        break
        else:
            print("Left-Right method")
            for e in tqdm(graph):
                g.add_edge(e['source'],e['dest'],weight=e['weight'])
                if check_planarity(g)[0]:
                    n_edge = g.number_of_edges()
                    if n_edge==nmax:
                        print('--- terminated (#edge={} reached the max) ---'.format(n_edge))
                        break
                    else:
                        rejected = 0
                else:
                    g.remove_edge(e['source'],e['dest'])
                    rejected += 1
                    if rejected > n_fdr*n_node: # equivalent to FDR
                        print('--- terminated (#rejected={} reached the FDR threshold) ---'.format(rejected))
                        break
        end = time.time()
        h,mod = divmod(end - sta,3600)
        m,s = divmod(mod,60)
        print("{0} hr {1} min {2} sec".format(int(h),int(m),round(s,4)))
        if update:
            self.set_graph(g)
        return g

    def _sort_graph(self,graph):
        """
        sort the given graph object of networkx

        Parameters
        ----------
        graph: networkx.Graph

        """
        sorted_edge = []
        ap = sorted_edge.append
        for s,d,w in sorted(graph.edges(data=True),key=lambda x:-x[2]['weight']):
            ap({'source':s,'dest':d,'weight':w['weight']})
        return sorted_edge

class GraphViewer(GraphHandler):
    def __init__(self):
        super().__init__()
        self.pos = None
        self.plt_params = {'figsize':(12,8),'edge_color':'lightgrey','width':0.2,
                           'font_size':14,'node_color':'royalblue','node_size':300,
                           'alpha':0.8}

    def plot(self,graph=None,pos=None,labels:dict=None,n_label:int=5,output:str='',
             fix_size:bool=False,fix_color:bool=False,
             cmap:str='Blues',centrality:str='pagerank',size_params:dict=None,
             plt_params:dict=dict()):
        """
        plot the given graph
        relatively take a long time

        Parameters
        ----------
        graph: networkx.Graph object

        pos: dict
            indicates the positions of nodes calculated by calc_pos method

        labels: dict
            indicates the labels
            keys and values are node indices and corresponding labels, respectively
            like inverter

        n_label: int
            indicates the number of nodes to be labeled
            0 indicates all nodes are labeled

        output: str
            the path for image when saved

        centrality: str
            indicates centrality method
            'betweenness', 'pagerank', 'degree', or 'closeness'

        size_params: dict
            indicates the parameters for size preparation in prep_size method

        plt_params: dict
            indicates the parameters for plot in general

        """
        if graph is None:
            if self.graph is None:
                raise ValueError('!! Provide graph to be visualized !!')
        else:
            self.set_graph(graph)
        if pos is None:
            if self.pos is None:
                raise ValueError('!! No pos: provide or prepare pos by calc_pos before this process !!')
            else:
                pos = self.pos
        if len(self.centrality)==0:
            self.calc_centrality(self.graph,centrality)
        # size preparation
        if not fix_size:
            if size_params is not None:
                size = self._prep_size(**size_params)
            else:
                size = self._prep_size()
            size = list(size.values())
        # color preparation
        if not fix_color:
            cm = plt.get_cmap(cmap)
            node_color = list(self.centrality.values())
        # label preparation
        if labels is not None:
            if n_label==0:
                label_focus = labels
            else:
                sc = sorted(self.centrality.items(),key=lambda x: x[1],reverse=True)[:n_label]
                sc = [v[0] for v in sc]
                label_focus = {k:v for k,v in labels.items() if k in sc}
        # plot
        self.plt_params.update(plt_params)
        plt.figure(figsize=self.plt_params['figsize'])
        nx.draw_networkx_edges(self.graph,pos,edge_color=self.plt_params['edge_color'],
                               width=self.plt_params['width'])
        alpha = self.plt_params['alpha']
        if fix_size:
            if fix_color:
                nx.draw_networkx_nodes(self.graph,pos,node_size=self.plt_params['node_size'],
                                       node_color=self.plt_params['node_color'],alpha=alpha)
            else:
                nx.draw_networkx_nodes(self.graph,pos,node_size=self.plt_params['node_size'],
                                       node_color=node_color,cmap=cm,alpha=alpha)
        else:
            if fix_color:
                nx.draw_networkx_nodes(self.graph,pos,node_size=size,node_color=self.plt_params['node_color'],
                                       alpha=alpha)
            else:
                nx.draw_networkx_nodes(self.graph,pos,node_size=size,node_color=node_color,cmap=cm,alpha=alpha)
        if labels is not None:
            nx.draw_networkx_labels(self.graph,pos,label_focus,font_size=self.plt_params['font_size'])
        plt.axis('off')
        if len(output) > 0:
            plt.savefig(output)
        plt.show()

    def plot_cluster(self,graph=None,cluster:dict=None,method:str='pagerank',layout:str='spring'):
        """
        plot the given graph
        relatively take a long time

        Parameters
        ----------
        graph: networkx.Graph object

        layout: str
            determines the shape of network
            'spring' or 'kamada_kawai'

        cluster: dict
            indicate where each node belongs

        method: str
            indicates centrality method
            'betweenness', 'pagerank', 'degree', or 'closeness'

        """
        raise NotImplementedError

    def calc_pos(self,graph=None,layout:str='spring'):
        """
        set a graph object
        relatively take a long time

        Parameters
        ----------
        graph: networkx.Graph object

        layout: str
            determines the shape of network
            'spring' or 'kamada_kawai'

        """
        if graph is not None:
            self.graph = graph
        if layout=='spring':
            self.pos = nx.spring_layout(self.graph)
        elif layout=='kamada_kawai':
            self.pos = nx.kamada_kawai_layout(self.graph)
        else:
            raise KeyError('!! Wrong layout: use spring or kamada_kawai !!')

    def set_inverter(self,inverter:dict):
        """
        set the index inverter, from int indices to string one

        """
        self.idx_inverter = inverter

    def _prep_size(self,basesize:int=500,power:bool=False,power_val:float=1.5):
        """
        prepare node size according to centrality

        basesize: int
            indicate the base size of nodes

        power: bool
            whether node size is prepared by power scaling

        power_val: float
            the value of power for scaling

        """
        size = [basesize*v for v in self.centrality.values()]
        if power:
            size = [np.power(v,power_val) for v in size]
        return dict(zip(self.centrality.keys(),size))


#!/usr/bin/python3

"""
# This file is part of the Pop Icon Theme and is free software; you can
# redistribute it and/or modify it under  the terms of the GNU Lesser General
# Public License as published by the Free Software Foundation; version 3.
#
# This file is part of the Pop Icon Theme and is distributed in the hope that
# it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty
# of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License along with
# this program; if not, see <https://www.gnu.org/licenses/lgpl-3.0.txt>
"""

import argparse
from ast import parse
import os
import shutil
import subprocess
from pathlib import Path

from importlib_metadata import sys

THEMENAME:str = 'Pop'
BINDIR = Path('/usr/bin')
BASEDIR = Path(os.getcwd())
SRCDIR = BASEDIR / 'src'
THEMEDIR = BASEDIR / THEMENAME
SIZES = (
    '8x8',
    '16x16',
    '24x24',
    '32x32',
    '48x48',
    '64x64',
    '128x128',
    '256x256',
    '512x512'
)

## Rendering Functions
def render_fullcolor() -> None:
    print('  -- Rendering bitmap icons...')
    os.chdir(SRCDIR / 'fullcolor')
    try:
        subprocess.run('./render-fullcolor.py', check=True)
    except subprocess.CalledProcessError:
        print('Failed to render fullcolor icons. See output above.')
        sys.exit(1)

def render_symbolics() -> None:
    print('  -- Rendering symbolic icons...')
    os.chdir(SRCDIR / 'scalable')
    try:
        subprocess.run('./extract-symbolic-icons.rb', check=True)
    except subprocess.CalledProcessError:
        print('Failed to render symbolic icons. See output above.')
        sys.exit(1)

def cleanup_unoptomized_renders() -> None:
    print('  -- Cleaning up any unoptimized output')

def generate_symlinks() -> None:
    print('  -- Generating symbolic links...')
    os.chdir(SRCDIR / 'symlinks')
    try:
        subprocess.run('./generate-symlinks.sh', check=True)
    except subprocess.CalledProcessError:
        print('Failed to generate fullcolor symlinks. See output above.')
        sys.exit(1)
    os.chdir(SRCDIR / 'scalable')
    try:
        subprocess.run('./generate-symbolic-symlinks.sh', check=True)
    except subprocess.CalledProcessError:
        print('Failed to generate sylbolic symlinks. See output above.')
        sys.exit(1)

def render_cursors() -> None:
    print('  -- Rendering cursors...')
    cursors_dir = SRCDIR / 'cursors'
    template_dir = cursors_dir / 'templates'
    output_dir = cursors_dir / 'bitmaps'

    os.chdir(cursors_dir)
    if output_dir.exists():
        print('    -- Output dir exists, use --clean to re-render')
        return
    shutil.copytree(template_dir, output_dir)

    print('    -- Rendering cursor fullcolor')
    subprocess.run(['./render-cursors.py', '-n 0', 'source-cursors.svg'])

    print('    -- Generatig cursor files')
    subprocess.run('./x11-make.sh')
    subprocess.run('./w32-make.sh')

def install_metadata() -> None:
    print('  -- Installing theme Metadata...')
    for file in ('index.theme', 'cursor.theme'):
        file_path = BASEDIR / f'{file}.in'
        shutil.copyfile(file_path, f'{THEMEDIR}/{file}')

## Artifact Cleanup/Removal Functions
def clean_fullcolor() -> None:
    print('  -- Removing Fullcolor Icons')
    for size in SIZES:
        size_dir = THEMEDIR / size
        if size_dir.exists():
            print(f'    -- Removing {size_dir}')
            shutil.rmtree(size_dir)
        else:
            print(f'    ** Skipping {size_dir}')

def clean_symbolics() -> None:
    scalable_dir = THEMEDIR / 'scalable'
    if scalable_dir.exists():
        print('  -- Removing symbolic icons')
        shutil.rmtree(scalable_dir)

def clean_cursors() -> None:
    cursor_dir = THEMEDIR / 'cursors'
    cursors_dir = SRCDIR / 'cursors'
    template_dir = cursors_dir / 'templates'
    output_dir = cursors_dir / 'bitmaps'
    if cursor_dir.exists():
        print('  -- Removing cursors')
        shutil.rmtree(cursor_dir)
    if output_dir.exists():
        print('    -- Cleaning up old render')
        shutil.rmtree(output_dir)

def clean_metadata() -> None:
    print('  -- Removing Metadata')
    for file in ('index.theme', 'cursor.theme'):
        file_path = THEMEDIR / file
        try:
            file_path.unlink()
            print(f'    -- Removed {file}')
        except FileNotFoundError:
            print(f'    ** Skipping {file}')

def clean_dirs(**kwargs) -> None:
    print('\n-- Cleaning up previous renders')
    os.chdir(THEMEDIR)

    if kwargs['everything']:
        print('  -- Performing Full Cleanup')
        clean_fullcolor()
        clean_symbolics()
        clean_cursors()
        clean_metadata()
        return

    # Cleanup Fullcolors
    if kwargs['fullcolor']:
        clean_fullcolor()
    else:
        print('  ** Skipping removing fullcolor')

    # Cleanup Symbolics
    if kwargs['symbolics']:
        clean_symbolics()
    else:
        print('  ** Skipping removing symbolic icons')

    # Cleanup Cursors
    if kwargs['cursors']:
        clean_cursors()
    else:
        print('  ** Skipping removing Cursors')

    # Cleanup Metadata
    if kwargs['metadata']:
        clean_metadata()
    else:
        print('  ** Skipping removing Metadata')
    print('\n\n')

def do_render(args) -> None:
    if args.clean:
        clean_dirs(
            everything=args.all,
            fullcolor=args.fullcolor,
            symbolics=args.symbolics,
            cursors=args.cursors,
            metadata=args.metadata
        )

    print('\n--Rendering icons')

    if args.all:
        render_fullcolor()
        render_symbolics()
        render_cursors()
        generate_symlinks()
        install_metadata()
        return

    if args.fullcolor:
        render_fullcolor()
    if args.symbolics:
        render_symbolics()
    if args.cursors:
        render_cursors()
    if args.links:
        generate_symlinks()
    if args.metadata:
        install_metadata()

parser = argparse.ArgumentParser(description='Render icons for the Pop Icon Theme')

parser.add_argument(
    '-c',
    '--clean',
    action='store_true',
    help='Remove existing files before rendering (takes a long time to render)'
)

parser.add_argument(
    '-a',
    '--all',
    action='store_true',
    help='Render all items (Default)'
)
parser.add_argument(
    '-f',
    '--fullcolor',
    action='store_true',
    help='Render fullcolor icons'
)
parser.add_argument(
    '-s',
    '--symbolics',
    action='store_true',
    help='Render Symbolic Icons'
)
parser.add_argument(
    '-x',
    '--cursors',
    action='store_true',
    help='Render Cursors'
)
parser.add_argument(
    '-l',
    '--links',
    action='store_true',
    help='Generate Theme Symlinks'
)
parser.add_argument(
    '-m',
    '--metadata',
    action='store_true',
    help='Generate Metadata'
)

args = parser.parse_args()

if not True in (args.fullcolor,
                args.symbolics,
                args.cursors,
                args.links,
                args.metadata):
    args.all = True
else:
    args.all = False

do_render(args)


# -*- coding: utf-8 -*-

from PyQt5.QtWidgets import QTreeWidgetItem, QDialog, QTreeWidgetItemIterator

from PyQt5.QtCore import pyqtSlot

from product.controllers.productcontroller import ProductController

from workshop.views.selectoddmentdraw import Ui_Dialog

import user

class SelectoddmentdrawModule(QDialog, Ui_Dialog):

    def __init__(self, ppid, parent=None):
        super(SelectoddmentdrawModule, self).__init__(parent)
        self.setupUi(self)
        self.ppid = ppid
        self.prodid = ''
        self.id_list = tuple()
        self.PC = ProductController()
        # 获取ppid对应的prodid
        self.get_prodid()
        # 查找prodid相同的记录，并保存这些记录的id_list
        self.get_id_list()
        # 零头记录中找到没有过有效期且没有被领取（dppid=0）
        self.get_oddments_list()

    def get_prodid(self):
        values_list = ('prodid',)
        key_dict = {
            'autoid': self.ppid
        }
        res = self.PC.get_producingplan(True, *values_list, **key_dict)
        if len(res) == 1:
            self.prodid = res[0]

    def get_id_list(self):
        if not self.prodid:
            return
        values_list = ('autoid', )
        key_dict = {
            'prodid': self.prodid
        }
        res = self.PC.get_producingplan(True, *values_list, **key_dict)
        # 去除本批ppid
        self.id_list = list(res)
        if self.ppid in self.id_list:
            self.id_list .remove(self.ppid)

    def get_oddments_list(self):
        self.treeWidget_oddmentdrawlist.clear()
        self.treeWidget_oddmentdrawlist.hideColumn(0)

        values_list = (
            'autoid', 'batchno', 'amount', 'unit', 'registerid', 'registername',
            'regdate', 'invaliddate'
        )
        # 没有发放，没有寄库，没有过期
        key_dict = {
            'ppid__in': self.id_list,
            'dppid': 0,
            'status':0,
            'invaliddate__gte': user.now_date,
        }
        res = self.PC.get_oddmentdrawnotes(False, *values_list, **key_dict)
        if not len(res):
            return
        for item in res:
            qtreeitem = QTreeWidgetItem(self.treeWidget_oddmentdrawlist)
            qtreeitem.setText(0, str(item['autoid']))
            qtreeitem.setText(1, item['batchno'])
            qtreeitem.setText(2, str(item['amount']))
            qtreeitem.setText(3, item['unit'])
            qtreeitem.setText(
                4, item['registerid'] + ' ' + item['registername']
            )
            qtreeitem.setText(5, str(item['regdate']))
            qtreeitem.setText(6, str(item['invaliddate']))
            qtreeitem.setCheckState(1, 0)
        for i in range(1, 7):
            self.treeWidget_oddmentdrawlist.resizeColumnToContents(i)

    @pyqtSlot(QTreeWidgetItem, int)
    def on_treeWidget_oddmentdrawlist_itemDoubleClicked(self, qtreeitem, p_int):
        state = qtreeitem.checkState(1)
        if state == 0:
            qtreeitem.setCheckState(1, 2)
        else:
            qtreeitem.setCheckState(1, 0)
    @pyqtSlot()
    def on_pushButton_accept_clicked(self):
        it = QTreeWidgetItemIterator(self.treeWidget_oddmentdrawlist)
        select_id_list = []
        while it.value():
            qtreeitem = it.value()
            if qtreeitem.checkState(1) == 2:
                select_id_list.append(int(qtreeitem.text(0)))
            it += 1
        if not len(select_id_list):
            return
        detail = {
            'dppid': self.ppid,
            'drawerid': user.user_id,
            'drawername': user.user_name,
            'drawdate': user.now_date
        }
        res = self.PC.update_oddmentdrawnotes(select_id_list, **detail)
        if res:
            self.accept()

    @pyqtSlot()
    def on_pushButton_cancel_clicked(self):
        self.close()


import json

from django.db.models import Q
from django.http import HttpResponse
from django.views import View
from mymodels.models import Posts, CustomUser
from mypackage.MixinClasses import GetUserMixin, SlicerMixin, ExcludeDelPostsMixin
import datetime
from django.utils import timezone

class PostsList(View, SlicerMixin, GetUserMixin, ExcludeDelPostsMixin):
    '''Базовый клас для отображения списка постов в тех или иных местах сайта, включая главную. Не используется ListView т.к. там запрещен POST'''
    posts_in_page = 5
    model = Posts
    template_name = 'index/post_list.html'
    context_object_name = 'posts'

    def get_tags_filter(self):
        '''Смотрит куки и возвращет спиок запрещенных постов'''
        cookies = self.request.COOKIES
        out_list = ['nsfw', 'политика', 'жесть']
        if cookies.get('is_18_age') and cookies.get('nsfw'):
            out_list.remove('nsfw')
        if cookies.get('politic'):
            out_list.remove('политика')
        if cookies.get('gore'):
            out_list.remove('жесть')

        return out_list

    def get_tags_user(self, value=True):
        '''Возвращает список подписаных или забаненых тегов пользователя в зависимоти от value'''
        user = self.get_user_session()
        if user==False:
            return []

        list_tags = user.subtags_set.filter(status = value).values('tag_id__tag')
        list_tags = [key['tag_id__tag'] for key in list_tags]

        return list_tags

    def get_queryset(self):
        '''Возварщает соответсвующий запрос'''
        start, end = self.get_slice(self.request, self.posts_in_page)

        list_views_posts = self.get_list_views_posts()

        block_tags = self.get_tags_filter() + self.get_tags_user(value=False)

        slice_time = timezone.now() - datetime.timedelta(days=90)  #расчитываем срез в -90 день от сегодня, напомниаю что у settings стоит USE_TZ = True

        queryset = Posts.objects.select_related('author').prefetch_related('evaluationposts_set', 'tags_set', 'comments_set')
        queryset = queryset.filter(date_create__gt=slice_time)
        queryset = queryset.exclude(Q(id__in=list_views_posts) | Q(tags__tag__in=block_tags))
        queryset = self.exclude_del_posts(queryset)
        queryset = queryset.order_by('-pk')

        return queryset[start:end]

    def get_list_views_posts(self):
        '''Возвращает массив просмотренных постов из POST если надо, или возвращает пустой массив'''
        if (self.request.COOKIES.get('hide_view')):
            return self.request.POST.getlist('views_posts[]')
        else:
            return []

    def get_posts(self):
        user = self.get_user_session()

        posts_list = self.get_queryset()

        out_data = []

        for one_post in posts_list:

            post_data = {}

            if not user:
                post_data['eval'] = 0
            # если get(), то будет кидать исключение при отсутствии объекта
            post_eval = one_post.evaluationposts_set.filter(
                user_id=user).first()
            if (post_eval == None):
                post_data['eval'] = 0
            elif post_eval.evaluation == True:
                post_data['eval'] = 1
            elif post_eval.evaluation == False:
                post_data['eval'] = -1

            tags_list = []
            for tag in one_post.tags_set.all():
                tags_list.append(tag.tag)
            post_data['tags'] = tags_list

            # В базу записывается словарь в формате JSON. Но python по умочанию, видит ее как строку.
            # для того, что бы python видел ее как объект нужная другая БД с поддержкой JSONField
            post_data['content'] = json.loads(one_post.content)

            # Функция превращает ее в словарь. Просто послать сроку нелья, т.к. js неадеквато ее пробразует
            post_data['name'] = one_post.name
            post_data['author'] = one_post.author.username
            post_data['avatar'] = one_post.author.profile.avatar.url
            post_data['rating'] = one_post.post_rating
            post_data['views'] = one_post.views
            post_data['comments'] = len(one_post.comments_set.all())
            post_data['post_id'] = one_post.pk
            if hasattr(one_post, 'deletedposts'): #проверка, удален ли пост, на 16.03.21 удаленные посты get_queryset() возврщает только на странице пользователя и то его собственные
                post_data['is_deleted'] = True

            if user:
                if one_post.author.username == user.username or user.group > 100:
                    post_data['may_change'] = True

            out_data.append(post_data)

        return json.dumps(out_data)

    def get(self, request, *args, **kwargs):
        return HttpResponse('Вообще то я жду POST')

    def post(self, request, *args, **kwargs):

        data = self.get_posts()

        return HttpResponse(data)


import csv
from elasticsearch import Elasticsearch
from elasticsearch import helpers

es = Elasticsearch([{'host': 'localhost', 'port': 9200}])
#es.indices.delete(index='movies', ignore=[400, 404])
print(es.ping())

def convert(filename,indexname,type):
    with open(filename, encoding="utf8") as file:
        r = csv.DictReader(file)
        helpers.bulk(es, r, index=indexname, doc_type=type)
es.indices.create(index = 'movies',ignore=400,body={
    "settings":{
        "index":{
            "similarity":{
                "default":{
                    "type":"BM25"
                }
            }
        }
    }})
es.indices.create(index = 'ratings',ignore=400,body={
    "settings":{
        "index":{
            "similarity":{
                "default":{
                    "type":"BM25"
                }
            }
        }
    }})
convert('movies.csv','movies','movies')
convert('ratings.csv','ratings','ratings')
print(es.indices.exists('movies'))
print(es.indices.exists('ratings'))


# PopGen 1.1 is A Synthetic Population Generator for Advanced
# Microsimulation Models of Travel Demand
# Copyright (C) 2009, Arizona State University
# See PopGen/License

from PyQt4.QtCore import *
from PyQt4.QtGui import *
from PyQt4.QtSql import *

from misc.widgets import *

import os, shutil

class SummaryPage(QWizardPage):
    def __init__(self, parent=None):
        super(SummaryPage, self).__init__(parent)

        self.projectLocationDummy = False
        self.projectDatabaseDummy = False

        self.setTitle("Step 6: Project Summary")
        vlayoutCol1 = QVBoxLayout()
        vlayoutCol1.addWidget(QLabel("Project name"))
        vlayoutCol1.addWidget(QLabel("Project location"))
        vlayoutCol1.addWidget(QLabel("Project description"))
        vlayoutCol1.addWidget(QLabel("Selected counties"))
        vlayoutCol1.addWidget(Separator())
        vlayoutCol1.addWidget(QLabel("Geographic resolution of population synthesis"))
        vlayoutCol1.addWidget(QLabel("Geographic correspondence data provided by the user"))
        vlayoutCol1.addWidget(QLabel("Location of the geographic correspondence file"))
        vlayoutCol1.addWidget(Separator())
        vlayoutCol1.addWidget(QLabel("Location data provided by the user"))
        vlayoutCol1.addWidget(QLabel("Location of the household sample file"))
        vlayoutCol1.addWidget(QLabel("Location of the groupquarter sample file"))
        vlayoutCol1.addWidget(QLabel("Location of the person sample file"))
        vlayoutCol1.addWidget(Separator())
        vlayoutCol1.addWidget(QLabel("Marginals data provided by the user"))
        vlayoutCol1.addWidget(QLabel("Location of the household marginals data file"))
        vlayoutCol1.addWidget(QLabel("Location of the groupquarter marginals data file"))
        vlayoutCol1.addWidget(QLabel("Location of the person marginals data file"))

        vlayoutCol2 = QVBoxLayout()

        self.projectNameLineEdit = DisplayLineEdit()
        vlayoutCol2.addWidget(self.projectNameLineEdit)

        self.projectLocationLineEdit = DisplayLineEdit()
        vlayoutCol2.addWidget(self.projectLocationLineEdit)

        self.projectDescLineEdit = DisplayLineEdit()
        vlayoutCol2.addWidget(self.projectDescLineEdit)

        self.projectRegionLineEdit = DisplayLineEdit()
        vlayoutCol2.addWidget(self.projectRegionLineEdit)

        vlayoutCol2.addWidget(Separator())

        #self.projectResolutionLineEdit = DisplayLineEdit()
        self.projectResolutionComboBox = ComboBoxFile()
        self.projectResolutionComboBox.setEnabled(False)
        self.projectResolutionComboBox.addItems(['County', 'Census Tract', 'Census Blockgroup'])
        vlayoutCol2.addWidget(self.projectResolutionComboBox)

        self.geocorrUserProvLineEdit = DisplayLineEdit()
        vlayoutCol2.addWidget(self.geocorrUserProvLineEdit)

        self.geocorrUserProvLocationLineEdit = DisplayLineEdit()
        vlayoutCol2.addWidget(self.geocorrUserProvLocationLineEdit)

        vlayoutCol2.addWidget(Separator())

        self.sampleUserProvLineEdit = DisplayLineEdit()
        vlayoutCol2.addWidget(self.sampleUserProvLineEdit)

        self.sampleHHLocationLineEdit = DisplayLineEdit()
        vlayoutCol2.addWidget(self.sampleHHLocationLineEdit)

        self.sampleGQLocationLineEdit = DisplayLineEdit()
        vlayoutCol2.addWidget(self.sampleGQLocationLineEdit)

        self.samplePersonLocationLineEdit = DisplayLineEdit()
        vlayoutCol2.addWidget(self.samplePersonLocationLineEdit)

        vlayoutCol2.addWidget(Separator())

        self.controlUserProvLineEdit = DisplayLineEdit()
        vlayoutCol2.addWidget(self.controlUserProvLineEdit)

        self.controlHHLocationLineEdit = DisplayLineEdit()
        vlayoutCol2.addWidget(self.controlHHLocationLineEdit)

        self.controlGQLocationLineEdit = DisplayLineEdit()
        vlayoutCol2.addWidget(self.controlGQLocationLineEdit)

        self.controlPersonLocationLineEdit = DisplayLineEdit()
        vlayoutCol2.addWidget(self.controlPersonLocationLineEdit)

        hlayout = QHBoxLayout()
        hlayout.addLayout(vlayoutCol1)
        hlayout.addLayout(vlayoutCol2)
        self.setLayout(hlayout)

    def fillPage(self, project):
        self.project = project
        self.projectNameLineEdit.setText(self.project.name)
        self.projectLocationLineEdit.setText(self.project.location)
        self.projectDescLineEdit.setText(self.project.description)
        dummy = ""
        if self.project.region is not None:
            for i in self.project.region.keys():
                dummy = dummy + i + ", "+ self.project.region[i]+ "; "
        self.projectRegionLineEdit.setText("%s"%dummy[:-2])
        #self.projectResolutionLineEdit.setText(self.project.resolution)

        resolutionText = self.project.resolution

        if resolutionText == "Tract":
            resolution = 'Census Tract'
        elif resolutionText == "Blockgroup":
            resolution = 'Census Blockgroup'
        elif resolutionText == 'TAZ':
            resolution = 'Traffic Analysis Zone (TAZ)'
            self.projectResolutionComboBox.clear()
            self.projectResolutionComboBox.addItems(['Traffic Analysis Zone (TAZ)',])

        else:
            resolution = 'County'

        self.projectResolutionComboBox.findAndSet(resolution)

        geocorrUserProv = self.convertBoolToString(self.project.geocorrUserProv.userProv)
        self.geocorrUserProvLineEdit.setText("%s" %geocorrUserProv)
        self.geocorrUserProvLocationLineEdit.setText(self.project.geocorrUserProv.location)

        sampleUserProv = self.convertBoolToString(self.project.sampleUserProv.userProv)
        self.sampleUserProvLineEdit.setText("%s" %sampleUserProv)
        self.sampleHHLocationLineEdit.setText(self.project.sampleUserProv.hhLocation)
        self.sampleGQLocationLineEdit.setText(self.project.sampleUserProv.gqLocation)
        self.samplePersonLocationLineEdit.setText(self.project.sampleUserProv.personLocation)

        controlUserProv = self.convertBoolToString(self.project.controlUserProv.userProv)
        self.controlUserProvLineEdit.setText("%s" %controlUserProv)
        self.controlHHLocationLineEdit.setText(self.project.controlUserProv.hhLocation)
        self.controlGQLocationLineEdit.setText(self.project.controlUserProv.gqLocation)
        self.controlPersonLocationLineEdit.setText(self.project.controlUserProv.personLocation)

    def convertBoolToString(self, value):
        if value:
            text = 'Yes'
        else:
            text = 'No, default data will be used'
        return text

    def enableEditableWidgets(self):
        self.projectDescLineEdit.setEnabled(True)
        self.projectResolutionComboBox.setEnabled(True)
        #self.geocorrUserProvLineEdit.setEnabled(True)
        #self.geocorrUserProvLocationLineEdit.setEnabled(True)
        #self.sampleUserProvLineEdit.setEnabled(True)
        #self.sampleHHLocationLineEdit.setEnabled(True)
        #self.sampleGQLocationLineEdit.setEnabled(True)
        #self.samplePersonLocationLineEdit.setEnabled(True)
        #self.controlUserProvLineEdit.setEnabled(True)
        #self.controlHHLocationLineEdit.setEnabled(True)
        #self.controlGQLocationLineEdit.setEnabled(True)
        #self.controlPersonLocationLineEdit.setEnabled(True)
        pass

    def updateProject(self):
        self.project.description = self.projectDescLineEdit.text()
        resolutionText = self.projectResolutionComboBox.currentText()
        if resolutionText == "Census Tract":
            resolution = 'Tract'
        elif resolutionText == "Census Blockgroup":
            resolution = 'Blockgroup'
        elif resolutionText == 'Traffic Analysis Zone (TAZ)':
            resolution = 'TAZ'
        else:
            resolution = 'County'

        self.project.resolution = resolution


from django.conf.urls import include, url
from teacherMan import views as tech

app_name = 'teacher'
urlpatterns = [
    url(r'main', tech.main, name='teacher'),
    url(r'data', tech.getData),

    url(r'edit', tech.editTech),

    url(r'delTechInfo', tech.delTech),
    url(r'addTechInfo', tech.addTech),
    #url(r'fillSelect', std.fillSelect)

]


def factorialRecursion(n):
    if n <= 0:
        return
    if n == 1:
        return n

    return n * factorialRecursion(n - 1)

def factorialIterative(n):
    fact = 1

    for num in range(2, n - 1):
        fact += num

    return fact

print(factorialRecursion(5))


import logging
import json
from invokust.aws_lambda import LambdaLoadTest, results_aggregator

logging.basicConfig(level=logging.INFO)

###
# SETTINGS
###

# How long should the test run for in minutes?
# Note that Lambda invokations that are started cannot be stopped.
# Test times will actually be run in intervals of 3 minutes.
test_time = 6
# How many concurrent users to test for?
# threads x 25 = num of concurrent users
threads = 20
# What test file are we using?
test_file_name = 'locust_test_mail_list.py'

def print_stat(type, name, req_count, median, avg, min, max, rps):
    return "%-7s %-50s %10s %9s %9s %9s %9s %10s" % (
        type,
        name,
        req_count,
        median,
        avg,
        min,
        max,
        rps,
    )

def print_stats_exit(load_test_state):
    summ_stats = load_test_state.get_summary_stats()
    agg_results = results_aggregator(load_test_state.get_locust_results())
    agg_results["request_fail_ratio"] = summ_stats["request_fail_ratio"]
    agg_results["invocation_error_ratio"] = summ_stats["invocation_error_ratio"]
    agg_results["locust_settings"] = load_test_state.lambda_payload
    agg_results["lambda_function_name"] = load_test_state.lambda_function_name
    agg_results["threads"] = load_test_state.threads
    agg_results["ramp_time"] = load_test_state.ramp_time
    agg_results["time_limit"] = load_test_state.time_limit
    logging.info("Aggregated results: {0}".format(json.dumps(agg_results)))

    logging.info(
        "\n============================================================"
        f"\nRamp up time: {agg_results['ramp_time']}s"
        f"\nStarted ramp down after {agg_results['time_limit']}s (time_limit)"
        f"\nThread count: {agg_results['threads']}"
        f"\nLambda invocation count: {agg_results['lambda_invocations']}"
        f"\nLambda invocation error ratio: {agg_results['invocation_error_ratio']}"
        f"\nCumulative lambda execution time: {agg_results['total_lambda_execution_time']}ms"
        f"\nTotal requests sent: {agg_results['num_requests']}"
        f"\nTotal requests failed: {agg_results['num_requests_fail']}"
        f"\nTotal request failure ratio: {agg_results['request_fail_ratio']}\n"
    )
    logging.info(
        "==========================================================================================================================="
    )
    logging.info(
        print_stat(
            "TYPE", "NAME", "#REQUESTS", "MEDIAN", "AVERAGE", "MIN", "MAX", "#REQS/SEC"
        )
    )
    logging.info(
        "==========================================================================================================================="
    )

    reqs = agg_results["requests"]
    for k in reqs.keys():
        k_arr = k.split("_")
        type = k_arr[0]
        del k_arr[0]
        name = "_".join(k_arr)
        logging.info(
            print_stat(
                type,
                name,
                reqs[k]["num_requests"],
                round(reqs[k]["median_response_time"], 2),
                round(reqs[k]["avg_response_time"], 2),
                round(reqs[k]["min_response_time"], 2),
                round(reqs[k]["max_response_time"], 2),
                round(reqs[k]["total_rps"], 2) * agg_results["threads"],
            )
        )
    total_rps = sum(
        [
            reqs[path]["total_rps"]
            for path in reqs
        ]
    ) * agg_results["threads"]
    logging.info(
        print_stat(
            "Total",
            f"Concurrent Users: {agg_results['threads']*25}",
            agg_results["num_requests"],
            "",
            "",
            "",
            "",
            round(total_rps, 2),
        )
    )

    logging.info("Exiting...")

if __name__ == "__main__":

    lambda_runtime = f"{test_time}m" if test_time < 3 else "3m"

    lambda_payload = {
        'locustfile': f"./tests/{test_file_name}",
        'host': 'https://forms-staging.cdssandbox.xyz',
        'num_users': 25,
        'spawn_rate': 5,
        'run_time': lambda_runtime
    }

    load_test = LambdaLoadTest(
    lambda_function_name='LoadTesting',
    threads=threads,
    ramp_time=0,
    time_limit=test_time*60,
    lambda_payload=lambda_payload
    )

    load_test.run()
    print_stats_exit(load_test)

    output_file = open("threads_output.json", "w")
    thread_output = {"threads": load_test.get_locust_results() }
    json.dump(thread_output, output_file)
    output_file.close()


import datetime
import json

from flask_restful import Resource
from flask import request
from init import app, db
from Models.player import Player
from Models.playerRequest import PlayerRequest
from decorators import json_required
import random

class RPlayerPost(Resource):

    def post(self, **kwargs):
        """ Создание нового игрока """
        unicue = random.randrange(9999999999)
        new_player_model = Player(unique=unicue)
        db.session.add(new_player_model)
        db.session.commit()
        player_model = Player.query.filter_by(unique=unicue).first()
        if player_model is not None:
            return {"message": "new player created", "player_id": player_model.id, "unique": player_model.unique}
        return {"message": "new player error created"}

class RPlayer(Resource):

    def get(self, player_id, **kwargs):
        """ Получение положения персонажа"""
        player_model = Player.query.filter_by(id=player_id).first()
        if player_model is None:
            return {"message": "player not found"}

        enemies_model = Player.query.filter(Player.id != player_id).all()
        enemies = dict()
        if enemies_model is not None:
            for enemy in enemies_model:
                enemy_data = {
                    "id": enemy.id,
                    "name": enemy.name,
                    "icon": enemy.icon,
                    "type": enemy.type,
                    "world_position": [enemy.y_world, enemy.x_world],
                }
                enemies[enemy.id] = enemy_data

        #print(F"\n\nenemies - {enemies}\n\n")
        answer = {
            "icon": player_model.icon,
            "name": player_model.name,
            "type": player_model.type,
            "dynamic_position": [player_model.y_dynamic, player_model.x_dynamic],
            "world_position": [player_model.y_world, player_model.x_world],
            "chunks_use_map": player_model.chunk,
            "level": player_model.level,
            "vertices": player_model.vertices,
            "direction": player_model.direction,
            "enemies": enemies,
            "assemblage_point": [player_model.assemblage_point_y, player_model.assemblage_point_x]
        }
        return answer

    def post(self, player_id, **kwargs):
        """ Отправка запроса """
        player_model = Player.query.filter_by(id=player_id).first()
        if player_model is None:
            return {"message": "player not found"}
        data = request.get_json(force=True)
        player_request = PlayerRequest.query.filter_by(player_id=player_id).first()
        if player_request is None:
            player_request = PlayerRequest(player_id=player_id, time=datetime.datetime.today(), type=data['type'],
                                           description=data['description'])
            db.session.add(player_request)
        else:
            player_request.date = datetime.datetime.today()
            player_request.type = data['type']
            player_request.description = data['description']

        db.session.commit()
        return {"message": "true"}


import django_filters as filters
from django_filters.rest_framework import FilterSet
from recipes.models import Ingredient, Recipe

class RecipeFilter(filters.FilterSet):
    "Filters recipes againts tags and author"
    tags = filters.AllValuesMultipleFilter(
        field_name='tags__slug'
    )

    class Meta:
        model = Recipe
        fields = (
            'author',
        )

    def filter_queryset(self, queryset):
        queryset = super().filter_queryset(queryset)
        if self.request.query_params.get('is_favorited') in [
            '1', 'true', 'True'
        ]:
            queryset = queryset.filter(watching__user=self.request.user)
        if self.request.query_params.get('is_in_shopping_cart') in [
            '1', 'true', 'True'
        ]:
            queryset = queryset.filter(goods__user=self.request.user)
        return queryset  # noqa

class IngredientNameFilter(filters.FilterSet):
    "Filters ingredients againts name and measurement_unit"
    name = filters.CharFilter(field_name='name', lookup_expr='istartswith')

    class Meta:
        model = Ingredient
        fields = ('name',)


import loadData as data
import numpy as np
import computeCost as cost

def gradientDescent(X, y, theta, alpha, num_iters):
    # Gradient descent is used to minimize cost function J
    # Vectorisation implementation. Applicable for any no for features.
    # Theta should be a nx1 vector where n = No of feature + 1

    #GRADIENTDESCENT Performs gradient descent to learn theta
    #   theta = GRADIENTDESCENT(X, y, theta, alpha, num_iters) updates theta by
    #   taking num_iters gradient steps with learning rate alpha

    # Algorithm
    # theta_i= theta_i - alpha * (1/m) Sum [i=1 to m] (h(x)-y)*x_i
    #    where i= training data

    m = y.size  # number of training examples
    J_history = np.zeros((num_iters, 1))

    for iter in range(num_iters):

        hx = X * theta  # computing h(x)
        sqrErrors = hx - y
        newX = sqrErrors.T * X
        theta = theta - ((alpha/m) * newX.T)

        # Save the cost J in every iteration
        J_history[iter,0] = cost.computeCost(X, y, theta)

    return theta, J_history
#  End of function

def main():
    # Calculate for theta [0,0]
    theta = np.zeros((2,1))
    alpha = 0.01
    iterations = 1500

    theta, J_history = gradientDescent(data.X, data.y, theta, alpha, iterations)

    print(theta)
    print(J_history)

# If this script is executed, then main() will be executed
if __name__ == '__main__':
    main()

Supervised Learning/Linear Regression/Linear Regression with One Variable/Python/gradientDescent.py


# coding: utf-8

# # 보스톤 집 값 예측하기
#
# * 보스턴 주택 데이터는 여러 개의 측정지표들을 포함한, 보스톤 인근의 주택가 중앙값
#
# * Variable in order:
#   - CRIM : 마을별 1인당 범죄율
#   - ZN : 25,000 평방미터를 초과하는 거주지역의 비율
#   - INDUS : 비소매 상업지역이 점유하고 있는 토지의 비율
#   - CHAS : 찰스 강에 대한 더미변수(강의 경계에 위치한 경우는 1, 아니면 0
#   - NOX : 10ppm 당 농축 일산화질소
#   - RM  : 주책 1가구당 평균 방의 개수
#   - AGE : 1940년 이전에 건축된 소유 주택의 비율
#   - DIS : 5개의 보스톤 직업센터까지의 접근성 지수
#   - RAD : 방사형 돌까지의 접근성 지수
#   - TAX : 10,000달러 당 재산세율
#   - PTRATIO : 시별 학생/교사 비율
#   - B   : 1000(Bk - 0.63) ^2, 여기서 Bk는 시별 흑인의 비율
#   - LSTAT : 모집단의 하위계층의 비율(%)
#   - MEDV : 본인 소유의 주택가격(중앙값) (단위:$1,000)

# In[24]:

from keras.models import Sequential
from keras.layers import Dense
from sklearn.model_selection import train_test_split

import numpy as np
import pandas as pd
import tensorflow as tf

# In[25]:

## 랜덤 시드값 설정
seed = 0
np.random.seed(seed)
tf.set_random_seed(seed)

# In[26]:

df =pd.read_csv('./data/housing.csv', delim_whitespace = True, header =None)

dataset = df.values
X = dataset[:, 0:13]
Y = dataset[:, 13]  ## 집값이 지 ㅇ낳을까?

# In[33]:

X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size = 0.3, random_state=seed)

# In[30]:

## 어떤 알고리즘을 적용할 것인지 선택
model = Sequential()
model.add(Dense(30, input_dim = 13, activation='relu'))
model.add(Dense(6, activation='relu'))
model.add(Dense(1))

# In[31]:

model.compile(loss = 'mean_squared_error', optimizer = 'adam')
model.fit(X_train, Y_train, epochs = 200, batch_size =10)

# In[34]:

# 예측값과 실제 값의 비교

Y_prediction = model.predict(X_test).flatten()
for i in range(10):
    label = Y_test[i]
    prediction = Y_prediction[i]
    print('실제가격 : {:.3f}, 예상가격 : {:.3f}'.format(label, prediction))

# In[ ]:

# In[ ]:

# In[ ]:

# In[ ]: