//**********************************************
// triads.cpp, Thomas R. Ioerger, copyright 1998
//   for information on this program, see...
//     http://www.cs.tamu.edu/faculty/ioerger/triads/triads.html
//**********************************************

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include "utils.h"
#include "aas.h"

#define MAXATOMS 100000

typedef struct {
  int num,res;
  int aa;
  char chain;
  char name[4];
  float x,y,z;
  char alt,ins;
} ATOM;

int natoms=0;
ATOM* atoms;

// constraints...

#define MAX_CA1_CA2  8 // 6.5
#define MAX_CA1_CA3  6 // 4.5
#define MAX_CA2_CA3 12 // 9.5
#define MAX_CH2_SS   7 // 4.5
#define MAX_S_S      3 // 2.0

//----------------------------------------------------------------------

#define PI 3.1415927

#define DIST(a,b,c,d,e,f) sqrt(((a)-(d))*((a)-(d))+((b)-(e))*((b)-(e))+((c)-(f))*((c)-(f)))

float angle(float a1,float a2,float a3,float b1,float b2,float b3,float c1,float c2,float c3)
{
  float d1,d2,d3,e1,e2,e3,d;

  d1 = a1-b1; d2 = a2-b2; d3 = a3-b3;
  e1 = c1-b1; e2 = c2-b2; e3 = c3-b3;
  
  d = (180.0/PI)*acos((d1*e1+d2*e2+d3*e3)/
      (sqrt(d1*d1+d2*d2+d3*d3)*sqrt(e1*e1+e2*e2+e3*e3)));

  return d;
}

// multiple 3x3 matrix A by 3-vector B and put answer in C (another 3 vector)

void mult_mat_vec(float A[][3],float B[],float C[])
{
  int j,k;

  for (j=0 ; j<3 ; j++) 
  {
    C[j] = 0.0;
    for (k=0 ; k<3 ; k++) C[j] += A[j][k]*B[k];
  }
}

//------------------------------------------------------------------

// i0....i1-----i2....i3
// translate i1 to origin
// rotate i2 to Z-axis
// project i0 and i3 onto XY-plane
// compute difference angle

float dihedral_angle(float x[4],float y[4],float z[4])
{
  int i,j,k;
  float rot[3][3],VA[3],VB[3];
  float a,b;

  x[0] -= x[1]; x[2] -= x[1]; x[3] -= x[1]; x[1] = 0;
  y[0] -= y[1]; y[2] -= y[1]; y[3] -= y[1]; y[1] = 0;
  z[0] -= z[1]; z[2] -= z[1]; z[3] -= z[1]; z[1] = 0;

  // rotate i2 around Z-axis to XZ-plane (make angle in XY-plane 0)

  a = asin(y[2]/sqrt(x[2]*x[2]+y[2]*y[2]));
  if (x[2]<0) a = PI-a;

  rot[0][0] = cos(-a);  rot[0][1] = -sin(-a); rot[0][2] = 0;
  rot[1][0] = sin(-a);  rot[1][1] = cos(-a);  rot[1][2] = 0;
  rot[2][0] = 0;        rot[2][1] = 0 ;       rot[2][2] = 1;

  VA[0] = x[0]; VA[1] = y[0]; VA[2] = z[0];
  mult_mat_vec(rot,VA,VB);
  x[0] = VB[0]; y[0] = VB[1]; z[0] = VB[2]; 

  VA[0] = x[2]; VA[1] = y[2]; VA[2] = z[2];
  mult_mat_vec(rot,VA,VB);
  x[2] = VB[0]; y[2] = VB[1]; z[2] = VB[2]; 

  VA[0] = x[3]; VA[1] = y[3]; VA[2] = z[3];
  mult_mat_vec(rot,VA,VB);
  x[3] = VB[0]; y[3] = VB[1]; z[3] = VB[2]; 

  // rotate i2 around Y-axis to Z-axis (make angle in XZ-plane 0)

  a = asin(x[2]/sqrt(x[2]*x[2]+z[2]*z[2]));
  if (z[2]<0) a = PI-a;

  rot[0][0] = cos(-a);  rot[0][1] = 0; rot[0][2] = sin(-a);
  rot[1][0] = 0;        rot[1][1] = 1; rot[1][2] = 0;
  rot[2][0] = -sin(-a); rot[2][1] = 0; rot[2][2] = cos(-a);

  VA[0] = x[0]; VA[1] = y[0]; VA[2] = z[0];
  mult_mat_vec(rot,VA,VB);
  x[0] = VB[0]; y[0] = VB[1]; z[0] = VB[2]; 

  VA[0] = x[2]; VA[1] = y[2]; VA[2] = z[2];
  mult_mat_vec(rot,VA,VB);
  x[2] = VB[0]; y[2] = VB[1]; z[2] = VB[2]; 

  VA[0] = x[3]; VA[1] = y[3]; VA[2] = z[3];
  mult_mat_vec(rot,VA,VB);
  x[3] = VB[0]; y[3] = VB[1]; z[3] = VB[2]; 

  // SGj is on Z-axis, project CBi and CBj onto XY-plane to compute angles

  a = asin(y[0]/sqrt(x[0]*x[0]+y[0]*y[0]));
  if (x[0]<0) a = PI-a;
  b = asin(y[3]/sqrt(x[3]*x[3]+y[3]*y[3]));
  if (x[3]<0) b = PI-b;
 
  a = 180.0*a/PI;
  b = 180.0*b/PI;
  a = b-a;
  if (a<-180) a = a+360;
  if (a>180) a = a-360;
  return a;
}

//----------------------------------------------------------------------

// atom name technically starts one column earlier (12)

void read_record(char* buf)
{
  int num,res;
  float x,y,z;
  char name[4];

  sscanf(buf+6,"%d",&num);
  sscanf(buf+22,"%d",&res);
  sscanf(buf+30,"%f%f%f",&x,&y,&z);

  atoms[natoms].num = num;
  atoms[natoms].res = res;
  atoms[natoms].chain = buf[21];
  atoms[natoms].aa = aastr2int(buf+17);
  atoms[natoms].alt = buf[16];
  atoms[natoms].ins = buf[26];
  
  atoms[natoms].x = x;
  atoms[natoms].y = y;
  atoms[natoms].z = z;

  atoms[natoms].name[0] = buf[13];
  atoms[natoms].name[1] = buf[14];
  atoms[natoms].name[2] = buf[15];
  atoms[natoms].name[3] = '\0';

  natoms++;
}

//----------------------------------------------------------------------

// pass in atom index and name of other atom in same residue/chain to find
// return -1 if not found
// insertion codes must match; what about alternative atoms?

int find_atom(int i,char* name)
{
  int m;

  m = 0;
  while ((m<natoms)&&((atoms[m].res!=atoms[i].res) 
          ||(atoms[m].chain!=atoms[i].chain)||(atoms[m].ins!=atoms[i].ins)
          ||(strcmp(atoms[m].name,name)!=0))) m++;
  if (m==natoms) return -1;
  else return m;
}

//----------------------------------------------------------------------

void print_constraints()
{
  printf("triad constraints...\n");
  printf("  atom 1     atom 2      typical traid dist     max dist cutoff\n");
  printf("  Cys1-CA    Cys2-CA           6.5 A                  %d A\n",MAX_CA1_CA2);
  printf("  Cys1-SG    Cys2-SG           2.0 A                  %d A\n",MAX_S_S);
  printf("  Cys1-CA    Trp-CA            4.5 A                  %d A\n",MAX_CA1_CA3);
  printf("  Cys2-CA    Trp-CA            9.5 A                  %d A\n",MAX_CA2_CA3);
  printf("  Trp-Cn2    S-S midp.         4.5 A                  %d A\n",MAX_CH2_SS);
}

//----------------------------------------------------------------------

void print_record(int i)
{
  printf("%s-%d%c-%s (%d): <%g,%g,%g>\n",
    aaint2str(atoms[i].aa),atoms[i].res,atoms[i].chain,atoms[i].name,
    atoms[i].num,atoms[i].x,atoms[i].y,atoms[i].z);
}

#define BUFSIZE 1000

// note: quits reading PDB file after ENDMDL

main(int argc,char** argv)
{
  FILE* fil;
  char buf[BUFSIZE];
  int i,j,k,l,m;
  int cys,trp;
  float a,b,d1,d2,d3,d4;
  int cb1,cb2,sg1,sg2,ch2;
  float x[4],y[4],z[4];

  printf("***************************************************************\n");
  printf("* TRIADS                                                      *\n");
  printf("*   Thomas R. Ioerger, copyright 1998                         *\n");
  printf("*   http://www.cs.tamu.edu/faculty/ioerger/triads/triads.html *\n");
  printf("***************************************************************\n");
  printf("\n");

  atoms = (ATOM*)malloc(MAXATOMS*sizeof(ATOM));
  if (atoms==NULL) error("unable to allocate sufficient memory");

  if (argc!=2) error("please give a PDB filename on the command line");
  if (strcmp(argv[1],"-c")==0) { print_constraints(); exit(0); }
  if ((fil=fopen(argv[1],"r"))==NULL) error("failed to open file"); 

  while (fgets(buf,BUFSIZE,fil)!=NULL)
  {
    if (strncmp(buf,"ENDMDL",6)==0) break;
    if (strncmp(buf,"ATOM",4)==0) read_record(buf);
  }  

  fclose(fil);

  cys = aastr2int("CYS");
  trp = aastr2int("TRP");

  for (i=0 ; i<natoms ; i++)
  {
    if (atoms[i].aa!=cys) continue;
    if (strcmp(atoms[i].name,"CA ")!=0) continue;
    for (j=0 ; j<natoms ; j++)
    {
      if (i==j) continue;
      if (atoms[j].aa!=cys) continue;
      if (strcmp(atoms[j].name,"CA ")!=0) continue;
      d1 = DIST(atoms[i].x,atoms[i].y,atoms[i].z,
                atoms[j].x,atoms[j].y,atoms[j].z);
      if (d1>MAX_CA1_CA2) continue; // ***
      for (k=0 ; k<natoms ; k++)
      {
        if (atoms[k].aa!=trp) continue;
        if (strcmp(atoms[k].name,"CA ")!=0) continue;
        d2 = DIST(atoms[i].x,atoms[i].y,atoms[i].z,
                  atoms[k].x,atoms[k].y,atoms[k].z);
        d3 = DIST(atoms[k].x,atoms[k].y,atoms[k].z,
                  atoms[j].x,atoms[j].y,atoms[j].z);
        if (d2>MAX_CA1_CA2) continue; // ***
        if (d3>MAX_CA2_CA3) continue; // ***

        // get C-beta's, S-gamma's, and Ch2
     
        if ((cb1=find_atom(i,"CB "))==-1) {
          printf("warning: C-beta atom for residue %s-%d%c not found\n",
            aaint2str(atoms[i].aa),atoms[i].res,atoms[i].chain); 
          continue; }

        if ((cb2=find_atom(j,"CB "))==-1) {
          printf("warning: C-beta atom for residue %s-%d%c not found\n",
            aaint2str(atoms[j].aa),atoms[j].res,atoms[j].chain);
          continue; }

        if ((sg1=find_atom(i,"SG "))==-1) {
          printf("warning: S-gamma atom for residue %s-%d%c not found\n",
            aaint2str(atoms[i].aa),atoms[i].res,atoms[i].chain);
          continue; }

        if ((sg2=find_atom(j,"SG "))==-1) {
          printf("warning: S-gamma atom for residue %s-%d%c not found\n",
            aaint2str(atoms[j].aa),atoms[j].res,atoms[j].chain);
          continue; }

        if ((ch2=find_atom(k,"CH2"))==-1) {
          printf("warning: C-eta-2 atom for residue %s-%d%c not found\n",
            aaint2str(atoms[k].aa),atoms[k].res,atoms[k].chain);
          continue; }

        a = angle(atoms[cb1].x-atoms[i].x,atoms[cb1].y-atoms[i].y,atoms[cb1].z-atoms[i].z,   0,0,0,   atoms[cb2].x-atoms[j].x,atoms[cb2].y-atoms[j].y,atoms[cb2].z-atoms[j].z);

        d4 = DIST((atoms[sg1].x+atoms[sg2].x)/2.0,(atoms[sg1].y+atoms[sg2].y)/2.0,(atoms[sg1].z+atoms[sg2].z)/2.0,atoms[ch2].x,atoms[ch2].y,atoms[ch2].z);

        if (DIST(atoms[sg1].x,atoms[sg1].y,atoms[sg1].z,atoms[sg2].x,atoms[sg2].y,atoms[sg2].z)>MAX_S_S) continue; // ***

        if (d4>MAX_CH2_SS) continue; // ***

        x[0] = atoms[cb1].x; y[0] = atoms[cb1].y; z[0] = atoms[cb1].z;
        x[1] = atoms[sg1].x; y[1] = atoms[sg1].y; z[1] = atoms[sg1].z;
        x[2] = atoms[sg2].x; y[2] = atoms[sg2].y; z[2] = atoms[sg2].z;
        x[3] = atoms[cb2].x; y[3] = atoms[cb2].y; z[3] = atoms[cb2].z;
        b = dihedral_angle(x,y,z);

        printf("triad:\n");

        printf("  1: ");
        print_record(i); 
        printf("  2: ");
        print_record(j); 
        printf("  3: ");
        print_record(k); 
     
        printf("  dist(CA1,CA2)=%g\n",d1);
        printf("  dist(CA1,CA3)=%g\n",d2);
        printf("  dist(CA2,CA3)=%g\n",d3);
        printf("  dist(Trp-Ch2,SS-midpoint)=%g\n",d4);

        printf("  angle between Cys C-alpha/C-beta vectors = %g degrees\n",a);
        printf("  dihedral angle across S-S bond = %g degrees\n",b);

        printf("\n");
      }
    }
  }
}