import three;

import three;
import cpkcolors;

// A sample Protein Data Bank file for this example is available from
// http://ndbserver.rutgers.edu/files/ftp/NDB/coordinates/na-biol/100d.pdb1

currentlight=White;
//currentlight=nolight;

//defaultrender.merge=true; // Fast low-quality PRC rendering
//defaultrender.merge=false; // Slow high-quality PRC rendering

bool pixel=false; // Set to true to draw dots as pixels.
real width=6;

size(200);
currentprojection=perspective(30,30,15);

pen chainpen=green;
pen hetpen=purple;

string filename="100d.pdb1";
//string filename=getstring("filename");

string prefix=stripextension(filename);
file data=input(filename);

pen color(string e)
{
e=replace(e," ","");
int n=length(e);
if(n < 1) return currentpen;
if(n > 1) e=substr(e,0,1)+downcase(substr(e,1,n-1));
int index=find(Element == e);
if(index < 0) return currentpen;
return rgb(Hexcolor[index]);
}

// ATOM
string[] name,altLoc,resName,chainID,iCode,element,charge;
int[] serial,resSeq;
real[][] occupancy,tempFactor;

bool newchain=true;

struct bond
{
int i,j;
void operator init(int i, int j) {
this.i=i;
this.j=j;
}
}

bond[] bonds;

struct atom
{
string name;
triple v;
void operator init(string name, triple v) {
this.name=name;
this.v=v;
}
}

struct chain
{
int[] serial;
atom[] a;
}

int[] serials;
chain[] chains;
atom[] atoms;

while(true) {
string line=data;
if(eof(data)) break;
string record=replace(substr(line,0,6)," ","");
if(record == "TER") {newchain=true; continue;}
bool ATOM=record == "ATOM";
bool HETATOM=record == "HETATM";
int serial;

atom a;
if(ATOM || HETATOM) {
serial=(int) substr(line,6,5);
a.name=substr(line,76,2);
a.v=((real) substr(line,30,8),
(real) substr(line,38,8),
(real) substr(line,46,8));
}
if(ATOM) {
if(newchain) {
chains.push(new chain);
newchain=false;
}
chain c=chains[chains.length-1];
c.serial.push(serial);
c.a.push(a);
continue;
}
if(HETATOM) {
serials.push(serial);
atoms.push(a);
}
if(record == "CONECT") {
int k=0;
int i=(int) substr(line,6,5);
while(true) {
string s=replace(substr(line,11+k,5)," ","");
if(s == "") break;
k += 5;
int j=(int) s;
if(j <= i) continue;
bonds.push(bond(i,j));
}
}
}

write("Number of atomic chains: ",chains.length);

int natoms;
begingroup3("chained");
for(chain c : chains) {
for(int i=0; i < c.a.length-1; ++i)
draw(c.a[i].v--c.a[i+1].v,chainpen,currentlight);
for(atom a : c.a)
if(pixel)
pixel(a.v,color(a.name),width);
else
dot(a.v,color(a.name),currentlight);
natoms += c.a.length;
}
endgroup3();

write("Number of chained atoms: ",natoms);
write("Number of hetero atoms: ",atoms.length);

begingroup3("hetero");
for(atom h : atoms)
if(pixel)
pixel(h.v,color(h.name),width);
else
dot(h.v,color(h.name),currentlight);
endgroup3();

write("Number of hetero bonds: ",bonds.length);

begingroup3("bonds");
for(bond b : bonds) {
triple v(int i) {return atoms[find(serials == i)].v;}
draw(v(b.i)--v(b.j),hetpen,currentlight);
}
endgroup3();

string options;
string viewfilename=prefix+".views";

if(!error(input(viewfilename,check=false)))
options="3Dviews="+locatefile(viewfilename);

shipout(options=options);