#! /usr/bin/env python2
# coding=utf-8

# prolog wymagany przez PyMOL-a
import __main__
__main__.pymol_argv = ['pymol', '-qc']
import pymol
pymol.finish_launching()

# własne importy
import scipy.cluster.hierarchy

# PDB-id struktur, które nas interesują
IDS = ('1ehz', '1evv', '4tna', '4tra', '6tna')

if __name__ == '__main__':
    # ściągnij każdą strukturę
    for i in IDS:
        pymol.cmd.fetch(i)

    # przygotuj macierz porównań parami
    count = len(IDS)
    matrix = [[-1 for i in range(count)] for i in range(count)]

    # porównaj parami każde dwie struktury (i wypisz od razu macierz)
    print 'Matrix:'
    for i in range(count):
        print '[', 
        for j in range(i):
            print '%0.2f' % matrix[i][j],
        for j in range(i, count):
            matrix[j][i] = matrix[i][j] = pymol.cmd.fit(IDS[i], IDS[j])
            print '%0.2f' % matrix[i][j],
        print ']'

    # wykonaj klastrowanie wyników
    # uznaj, że klastry oddzielone są o 3A
    Z = scipy.cluster.hierarchy.average(matrix)
    print scipy.cluster.hierarchy.fcluster(Z, 3, 'distance')