/*--------------------------------------------MECHANO---------------------------------------------*/

/*--Builds loops from both N and C anchors, to a total number specified by the '-decoys' option*/
/*--Uses H3-specific dihedral data, generated using Jin's code*/
/*--Includes 'switching' - starting from the N direction, if a loop is discarded then the next loop
    is built in the opposite direction*/

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <math.h>

#define pi 3.14159265359

double ALA_CCD_dist[5184];
double ARG_CCD_dist[5184];
double ASN_CCD_dist[5184];
double ASP_CCD_dist[5184];
double CYS_CCD_dist[5184];
double GLN_CCD_dist[5184];
double GLU_CCD_dist[5184];
double GLY_CCD_dist[5184];
double HIS_CCD_dist[5184];
double ILE_CCD_dist[5184];
double LEU_CCD_dist[5184];
double LYS_CCD_dist[5184];
double MET_CCD_dist[5184];
double PHE_CCD_dist[5184];
double PRO_CCD_dist[5184];
double SER_CCD_dist[5184];
double THR_CCD_dist[5184];
double TRP_CCD_dist[5184];
double TYR_CCD_dist[5184];
double VAL_CCD_dist[5184];

int *ptr_array[20];
int dihedral_array_lengths[20];

int **ALA_dihedrals;
int **ARG_dihedrals;
int **ASN_dihedrals;
int **ASP_dihedrals;
int **CYS_dihedrals;
int **GLN_dihedrals;
int **GLU_dihedrals;
int **GLY_dihedrals;
int **HIS_dihedrals;
int **ILE_dihedrals;
int **LEU_dihedrals;
int **LYS_dihedrals;
int **MET_dihedrals;
int **PHE_dihedrals;
int **PRO_dihedrals;
int **SER_dihedrals;
int **THR_dihedrals;
int **TRP_dihedrals;
int **TYR_dihedrals;
int **VAL_dihedrals;

/* Files containing MECHANO functions and data */
#include "include_data.h"
#include "other_things.h"
#include "other_functions.h"
#include "MECHANO_functions.h"
#include "MECHANO_functions_C.h"

int main(char pdb_file[], char chain[], char loop_sequence[], int loop_length, int start, int end, int decoys,
         double FREADbond_lengthsN[], double FREADbond_anglesN[], double FREADdihedralsN[],
         double FREADbond_lengthsC[], double FREADbond_anglesC[], double FREADdihedralsC[],
         char bitstring[], char log_file[], char decoy_file[], double C_native_coords[][3], double N_anchor[][3],
         double C_anchor[][3], double nitrogens[][3], double a_carbons[][3], double carbons[][3], 
         double oxygens[][3], int atom_frequencies[], int info_array[], char N_seq[], char C_seq[], char decoy_code[],         char sphinx_version[])

{
    char results_file[200];
    char results_N[200];
    char results_C[200];
    char seq_inc_anchors[200];
    int ins_res;     /*Residue number on which to place insertion codes*/
    int opt1 = 0;
    int opt2 = 0;
    int i, j, k;
    int no_extra_residues;
    int last_atom_no;
    int native_present;
    int generated = 0;
    int clashed = 0;
    int closure_outcome_N;
    int closure_outcome_C;
    int clash_outcome;
    int self_clash_outcome;
    int clash_count = 0;
    int failed_to_close = 0;
    int generated_N = 0;
    int generated_C = 0;
    int direction = 0;  /*0 means from N anchor, 1 from C anchor*/
    double first_O_dihedral;
    double last_psi;
    double last_O_dihedral;
    double last_N_anchor[4][3];
    double first_C_anchor[4][3];
    double **native_loop;
    double **loop_model_N;
    double **loop_model_C;
    double **loop_model_C_reordered;
    double **final_model;
    double **final_model_Cb;
    double RMSD;
    double ex_time;
    double **open_loop;
    double **half_closed;
    int limit;
    
    FILE *output_log;
    FILE *output;
    clock_t begin_time, end_time;
    time_t mytime;
    
    mytime = time(NULL);
    begin_time = clock();
    srand(time(NULL));

    /* Fill the *_CCD_dist and *_dihedral data arrays with the correct data for the requested Sphinx version */
    if (strcmp(sphinx_version, "h3") == 0) {
        copy_arrays_h3();
    }
    else if (strcmp(sphinx_version, "membrane") == 0) {
        copy_arrays_membrane();
    }
    else {
        copy_arrays_general();
    }
        
    if (opt2 == 1) {
        if (ins_res < start || ins_res > end) {
            printf("Error: The residue number on which to put insertion codes must be between the "
                   "start and end values.\n");
            return 0;
        }
    }

    /*printf("%s\n", decoy_file);*/

    int res_numbers[loop_length];
    allocate_array(&native_loop, loop_length*4);
    allocate_array(&loop_model_N, (loop_length+2)*4);
    allocate_array(&loop_model_C, (loop_length+2)*4);
    allocate_array(&loop_model_C_reordered, (loop_length+2)*4);
    allocate_array(&final_model, loop_length*4);
    allocate_array(&open_loop, loop_length*4);
    allocate_array(&half_closed, loop_length*4);
    allocate_array(&final_model_Cb, loop_length*5);
     
    /*output = fopen(decoy_file, "w");
    if (output == NULL) {
        printf("Error: Could not open output file.\n");
        return 0;
    }*/
        
    output_log = fopen(log_file, "a");
    if (output_log == NULL) {
        printf("Error: Could not open log file.\n");
        return 0;
    }

    /*Sort out whether insertion codes/gaps are required in the residue numbering*/
    if ((end - start) + 1 == loop_length) {
        for (i = 0; i < loop_length; i++) {
            res_numbers[i] = start + i;
        }
    }
    else if ((end - start) + 1 < loop_length) {
        /*need to use insertion codes*/
        if (opt2 == 0) {  /*Optional argument to specified which residue number to use not given*/ 
            ins_res = ((end - start) + 1) / 2 + start;
        }
        no_extra_residues = loop_length - ((end - start) + 1);
        for (i = 0; i < ins_res - start; i++) {
            res_numbers[i] = start + i;
        }
        for (j = 0; j < no_extra_residues + 1; j++) {
            res_numbers[i+j] = ins_res;
        }
        for (k = 0; k < loop_length-(i+j); k++) {
            res_numbers[i+j+k] = ins_res + k + 1;
        }
    }
    else if ((end - start) + 1 > loop_length) {
        if (loop_length%2 == 0) {
            for (i = 0; i < loop_length/2; i++) {
                res_numbers[i] = start + i;
            }
        }
        else {
            for (i = 0; i < loop_length/2 + 1; i++) {
                res_numbers[i] = start + i;
            }
        }
        for (j = 0; j < loop_length/2; j++) {
            res_numbers[i+j] = end - (loop_length/2 - j - 1);
        }
    }

    
    /*anchor_extractor(pdb_file, chain, start, end, N_anchor, N_seq, C_anchor, C_seq, native_loop,
                     loop_length, info_array);*/
                     
    /*count_atoms(pdb_file, chain, start, end, atom_frequencies);*/
                     
    /*allocate_array(&nitrogens, atom_frequencies[0]);
    allocate_array(&a_carbons, atom_frequencies[1]);
    allocate_array(&carbons, atom_frequencies[2]);
    allocate_array(&oxygens, atom_frequencies[3]);*/
    
                       
    /*structure_extractor(pdb_file, chain, start, end, nitrogens, a_carbons, carbons, oxygens, 
                        atom_frequencies);*/
                                      
    native_present = info_array[0];
    last_atom_no = info_array[1];
                                   
    
    /*Generate sequence string which includes all loop residues and the residue on each side*/
    seq_inc_anchors[0] = N_seq[2];
    for (i=1;i<loop_length+1;i++){
        seq_inc_anchors[i] = loop_sequence[i-1];
    }
    seq_inc_anchors[loop_length+1] = C_seq[0];
    seq_inc_anchors[loop_length+2] = '\0';

    /*Get last residue of N anchor and first of C anchor (i.e. 1 residue each side of loop)*/
    for (i=0;i<4;i++){
        for (j=0;j<3;j++){
            last_N_anchor[i][j] = N_anchor[i+8][j];
            first_C_anchor[i][j] = C_anchor[i][j];
        }
    }

    /*Calculate angles required for loop building 
      (do this here because it only needs to be done once, these angles don't change)*/
    first_O_dihedral = calculate_dihedral(N_anchor[8], N_anchor[9], N_anchor[10], N_anchor[11]);
    last_psi = calculate_dihedral(C_anchor[0], C_anchor[1], C_anchor[2], C_anchor[4]);
    last_O_dihedral = calculate_dihedral(C_anchor[0], C_anchor[1], C_anchor[2], C_anchor[3]);

    
    /*----------------------------------------------------------------------------------------------*/
    /*LOOP BUILDING---------------------------------------------------------------------------------*/
    
    if (decoys == 1) {
        limit = 50;
    }
    else {
        limit = decoys*30;
    }
    
    direction = 0;

    while (generated < decoys) {

        if (clash_count + failed_to_close > limit) {
            printf("    %d loop structures generated                \n\n", generated);
            fprintf(output_log, " %d closure failures, %d clashes\n", failed_to_close, clash_count);
            fprintf(output_log, " %d loop structures generated\n", generated);
            if (generated > 0) {
                fclose(output);
            }
            fclose(output_log);
            return 0;
        }
        
        if (direction == 0) {
            /*Build loop from N-anchor*/
            build_loop_N(seq_inc_anchors, last_N_anchor, first_C_anchor, loop_model_N, 
                         first_O_dihedral, last_psi, last_O_dihedral, FREADbond_lengthsN, 
                         FREADbond_anglesN, FREADdihedralsN, bitstring);    

            for (i=0; i<loop_length*4; i++) {
                for (j=0; j<3; j++) {
                    open_loop[i][j] = loop_model_N[i+4][j];
                }
            }
                       
            //Initial closure, not changing dihedrals taken from FREAD structure
            //Until C anchor RMSD < 1.0A
            //closure_outcome_N = close_loop_N_first(loop_model_N, C_anchor, loop_length, seq_inc_anchors, bitstring);
            
            /*for (i=0; i<loop_length*4; i++) {
                for (j=0; j<3; j++) {
                    half_closed[i][j] = loop_model_N[i+4][j];
                }
            }*/
            closure_outcome_N = 1;
                        
            if (closure_outcome_N != 0) {
                
                //Complete closure
                closure_outcome_N = close_loop_N(loop_model_N, C_anchor, loop_length, seq_inc_anchors);
                //closure_outcome_N = close_loop_N_limited(loop_model_N, C_anchor, loop_length, seq_inc_anchors, bitstring, FREADdihedralsN);               
                
                if (closure_outcome_N != 0) {
                    self_clash_outcome = self_clash_check(loop_model_N, loop_length);
                    
                    if (self_clash_outcome != 0) {
                        clash_outcome = clash_check(loop_model_N, loop_length, nitrogens, a_carbons, 
                                                    carbons, oxygens, atom_frequencies);
                                                    
                        if (clash_outcome != 0) {
                    
                            for (i=0; i<loop_length*4; i++) {
                                for (j=0; j<3; j++) {
                                    final_model[i][j] = loop_model_N[i+4][j];
                                }
                            }
                            generated_N++;
                            generated++;
                            
                            if (generated == 1) {
                                output = fopen(decoy_file, "a");
                                if (output == NULL) {
                                    printf("Error: Could not open output file.\n");
                                    return 0;
                                }
                            }
                            
                            if (native_present == 1) {
                                RMSD = calculate_RMSD(final_model, C_native_coords, 0, loop_length*4);
                            }
                            else {
                                RMSD = 0;
                            }
                                                   
                            build_Cb_atoms(final_model, final_model_Cb, loop_length, loop_sequence);
                                                   
                            write_structure_to_file_Cb(final_model_Cb, generated, output, loop_sequence, chain,
                                                       loop_length, RMSD, res_numbers, last_atom_no, 0, decoy_code);
                                                             
                            if (generated%10 == 0){
                                printf("    %d loop structures generated...       \r",generated);
                                fflush(stdout);
                            }
                        }
                        else {
                            clash_count++;
                            direction = 1;
                        }
                    }
                    else {
                        clash_count++;
                        direction = 1;
                    }
                }
                else {
                    failed_to_close++;
                    direction = 1;
                }
            }
            else {
                failed_to_close++;
                direction = 1;
            }
        }
        else if (direction == 1) {
            /*Build loops from C-anchor*/          
            build_loop_C(seq_inc_anchors, last_N_anchor, first_C_anchor, loop_model_C, first_O_dihedral,
                       last_psi, last_O_dihedral, FREADbond_lengthsC, FREADbond_anglesC, FREADdihedralsC,
                       bitstring);
                       
            for (j=0; j<(loop_length+2)*4; j++) {
                if (j%4 == 0) {
                    loop_model_C_reordered[j+1][0] = loop_model_C[j][0];
                    loop_model_C_reordered[j+1][1] = loop_model_C[j][1];
                    loop_model_C_reordered[j+1][2] = loop_model_C[j][2];
                }
                else if ((j-1)%4 == 0) {
                    loop_model_C_reordered[j-1][0] = loop_model_C[j][0];
                    loop_model_C_reordered[j-1][1] = loop_model_C[j][1];
                    loop_model_C_reordered[j-1][2] = loop_model_C[j][2];
                }
                else {
                    loop_model_C_reordered[j][0] = loop_model_C[j][0];
                    loop_model_C_reordered[j][1] = loop_model_C[j][1];
                    loop_model_C_reordered[j][2] = loop_model_C[j][2];
                }
            }
    
            //closure_outcome_C = close_loop_C_first(loop_model_C_reordered, last_N_anchor, loop_length, seq_inc_anchors, bitstring);
            closure_outcome_C = 1;
            
            if (closure_outcome_C != 0) {
                
                closure_outcome_C = close_loop_C(loop_model_C_reordered, last_N_anchor, loop_length, seq_inc_anchors);
                //closure_outcome_C = close_loop_C_limited(loop_model_C_reordered, last_N_anchor, loop_length, seq_inc_anchors, bitstring, FREADdihedralsN);
                
                
                if (closure_outcome_C != 0) {
                    self_clash_outcome = self_clash_check(loop_model_C_reordered, loop_length);
                    
                    if (self_clash_outcome != 0) {
                        clash_outcome = clash_check(loop_model_C_reordered, loop_length, nitrogens, 
                                                    a_carbons, carbons, oxygens, atom_frequencies);
                                                    
                        if (clash_outcome != 0) {
                            
                            for (j=0; j<(loop_length+2)*4; j++) {
                                loop_model_C[j][0] = loop_model_C_reordered[(loop_length+2)*4-1-j][0];
                                loop_model_C[j][1] = loop_model_C_reordered[(loop_length+2)*4-1-j][1];
                                loop_model_C[j][2] = loop_model_C_reordered[(loop_length+2)*4-1-j][2];
                            }
                            
                            for (i=0; i<loop_length*4; i++) {
                                for (j=0; j<3; j++) {
                                    final_model[i][j] = loop_model_C[i+4][j];
                                }
                            }
                            generated_C++;
                            generated++;
                            
                            if (generated == 1) {
                                output = fopen(decoy_file, "a");
                                if (output == NULL) {
                                    printf("Error: Could not open output file.\n");
                                    return 0;
                                }
                            }
                            
                            if (native_present == 1) {
                                RMSD = calculate_RMSD(final_model, C_native_coords, 0, loop_length*4);
                            }
                            else {
                                RMSD = 0;
                            }
                            
                            build_Cb_atoms(final_model, final_model_Cb, loop_length, loop_sequence);
                            
                            write_structure_to_file_Cb(final_model_Cb, generated, output, loop_sequence, 
                                                   chain, loop_length, RMSD, res_numbers, last_atom_no, 1, decoy_code);
                                                    
                            if (generated%10 == 0){
                                printf("    %d loop structures generated...       \r",generated);
                                fflush(stdout);
                            }
                        }
                        else {
                            clash_count++;
                            direction = 0;
                        }
                    }
                    else {
                        clash_count++;
                        direction = 0;
                    }
                }
                else {
                    failed_to_close++;
                    direction = 0;
                }
            }
            else {
                failed_to_close++;
                direction = 0;
            }
        }
    }
    
    printf("    %d loop structures generated!!                   \n\n",generated);
    
    fclose(output);
    
    end_time = clock();
    ex_time = (double)(end_time - begin_time) / CLOCKS_PER_SEC;
    
    /*fprintf(output_log, " %d closure failures, %d clashes\n", failed_to_close, clash_count);*/
    fprintf(output_log, " %d loop structures generated\n", generated);
    
    fclose(output_log);
    
    free(native_loop);
    free(loop_model_N);
    free(loop_model_C);
    free(loop_model_C_reordered);
    free(final_model);
    free(final_model_Cb);
    /*free(nitrogens);
    free(a_carbons);
    free(carbons);
    free(oxygens);*/

    free(ALA_dihedrals);
    free(ARG_dihedrals);
    free(ASN_dihedrals);
    free(ASP_dihedrals);
    free(CYS_dihedrals);
    free(GLN_dihedrals);
    free(GLU_dihedrals);
    free(GLY_dihedrals);
    free(HIS_dihedrals);
    free(ILE_dihedrals);
    free(LEU_dihedrals);
    free(LYS_dihedrals);
    free(MET_dihedrals);
    free(PHE_dihedrals);
    free(PRO_dihedrals);
    free(SER_dihedrals);
    free(THR_dihedrals);
    free(TRP_dihedrals);
    free(TYR_dihedrals);
    free(VAL_dihedrals);

}
