#!/usr/bin/env python
# -*- coding: utf-8 -*-

#PCAによって得られた極構造を0->1にコーンを描く
#コマンドライン引数にはPCAの出力を１つ(複数モデル)を入れる


from Bio.PDB.PDBParser import PDBParser
import sys
import math
import colorsys

class Point3D:
    def __init__(self, x1, y1, z1):
        self.x = x1
        self.y = y1
        self.z = z1

    def distance(p1, p2):
        dx = p1.x - p2.x
        dy = p1.y - p2.y
        dz = p1.z - p2.z
        return math.sqrt(dx * dx + dy * dy + dz * dz)

    def vector(p1,p2):
        vect_x = p1.x - p2.x
        vect_y = p1.y - p2.y
        vect_z = p1.z - p2.z       
        return (Point3D(vect_x,vect_y,vect_z))

    def vec_to_dis(vec):
        return math.sqrt(vec.x * vec.x + vec.y * vec.y + vec.z * vec.z)

def get_vec(center_model,top_model):
    vec_list=[]
    for i in range(len(center_model)):
        #p1->p2
        p1=Point3D(center_model[i][0],center_model[i][1],center_model[i][2])
        p2=Point3D(top_model[i][0],top_model[i][1],top_model[i][2])
        vec=p2.vector(p1)
        vec_list.append(vec)
    return vec_list

def cal_dist(vec_list): #vec_listの中の型はPoint3D
    dis_list=[]
    for vec_i in vec_list:
        #p1->p2
        item=vec_i.vec_to_dis()
        dis_list.append(item)
    min_dis=min(dis_list)
    for i in range(len(dis_list)):
       dis_list[i]=dis_list[i]-min_dis
       i+=1
    return dis_list 


def get_coord_of_model(model):
    models_coord=[]
    for chain in model.get_list():
        for residue in chain.get_list():
            if residue.has_id("CA"):
                ca_atom=residue["CA"]
                ca_coord=ca_atom.get_coord()
                models_coord.append(ca_coord)
    return (models_coord)

def get_color(max,dis_v):
    weight=1-(float(dis_v)/float(max))
    hue=(240.0/360.0)*weight
    RGB=colorsys.hsv_to_rgb(hue,1,1)

    if(dis_v >= max*0.8):
        RGB=['255','0','0']
    else:
        RGB=['128','0','128']

    return (RGB[0]+'\t'+RGB[1]+'\t'+RGB[2])

def write_cone(ofile,center,top,dis_coord):
    r='0.37' #size of cone
    drline = 2
    max_v=max(dis_coord)
    Length=len(center)
    for i in range(Length)[:-1]:
        if dis_coord[i] > drline:
            out_line='\t'.join(map(str,center[i]))+'\t'+'\t'.join(map(str,top[i]))+'\t'+r
            color=get_color(max_v,dis_coord[i])
            ofile.write('.color\t'+color+'\n')
            ofile.write('.cone\t'+out_line+'\n')
            ofile.flush()

#ここまでimportや関数定義
#ここからmainプログラム

pdb_models_file=sys.argv[1]


parser=PDBParser(PERMISSIVE=1)
models_str=parser.get_structure(pdb_models_file[-3], pdb_models_file)
start_coord=get_coord_of_model(models_str.get_list()[0])
end_coord=get_coord_of_model(models_str.get_list()[-1])

vec_list=get_vec(start_coord,end_coord)
dis_list=cal_dist(vec_list)
print 'max distance=%.2f' %(max(dis_list))
out_file=open('PCA_vec.bild','w')
print 'create "PCA_vec.bild" file'
write_cone(out_file,start_coord,end_coord,dis_list)
out_file.close()